From 137554f5226ff9b338c9c602b4e40833ea3b84d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=83=91=E4=BD=A9=E8=8C=B9?= <zhengpeiru21@protonmail.com>
Date: Tue, 3 May 2022 15:59:33 -0600
Subject: [PATCH] Add getDeckStats action

---
 README.md           | 40 ++++++++++++++++++++++++++++++++++++++++
 plugin/__init__.py  | 27 +++++++++++++++++++++++++++
 tests/test_decks.py |  5 +++++
 3 files changed, 72 insertions(+)

diff --git a/README.md b/README.md
index a24789e..15be80f 100644
--- a/README.md
+++ b/README.md
@@ -960,6 +960,46 @@ corresponding to when the API was available for use.
     }
     ```
 
+*   **getDeckStats**
+
+    Gets statistics such as total cards and cards due for the given decks.
+
+    *Sample request*:
+    ```json
+    {
+        "action": "getDeckStats",
+        "version": 6,
+        "params": {
+            "decks": ["Japanese::JLPT N5", "Easy Spanish"]
+        }
+    }
+    ```
+
+    *Sample result*:
+    ```json
+    {
+        "result": {
+            "1651445861967": {
+                "deck_id": 1651445861967,
+                "name": "Japanese::JLPT N5",
+                "new_count": 20,
+                "learn_count": 0,
+                "review_count": 0,
+                "total_in_deck": 1506
+            },
+            "1651445861960": {
+                "deck_id": 1651445861960,
+                "name": "Easy Spanish",
+                "new_count": 26,
+                "learn_count": 10,
+                "review_count": 5,
+                "total_in_deck": 852
+            }
+        },
+        "error": null
+    }
+    ```
+
 #### Graphical Actions
 
 *   **guiBrowse**
diff --git a/plugin/__init__.py b/plugin/__init__.py
index 898f81b..eaad2c4 100644
--- a/plugin/__init__.py
+++ b/plugin/__init__.py
@@ -346,6 +346,21 @@ class AnkiConnect:
         except NotFoundError:
             raise NotFoundError('Note was not found: {}'.format(note_id))
 
+    def deckStatsToJson(self, due_tree):
+        return {'deck_id': due_tree.deck_id,
+                'name': due_tree.name,
+                'new_count': due_tree.new_count,
+                'learn_count': due_tree.learn_count,
+                'review_count': due_tree.review_count,
+                'total_in_deck': due_tree.total_in_deck}
+
+    def collectDeckTreeChildren(self, parent_node):
+        allNodes = {parent_node.deck_id: parent_node}
+        for child in parent_node.children:
+            for deckId, childNode in self.collectDeckTreeChildren(child).items():
+                allNodes[deckId] = childNode
+        return allNodes
+
     #
     # Miscellaneous
     #
@@ -617,6 +632,18 @@ class AnkiConnect:
         collection.decks.remConf(configId)
         return True
 
+    @util.api()
+    def getDeckStats(self, decks):
+        collection = self.collection()
+        scheduler = self.scheduler()
+        responseDict = {}
+        deckIds = list(map(lambda d: collection.decks.id(d), decks))
+
+        allDeckNodes = self.collectDeckTreeChildren(scheduler.deck_due_tree())
+        for deckId, deckNode in allDeckNodes.items():
+            if deckId in deckIds:
+                responseDict[deckId] = self.deckStatsToJson(deckNode)
+        return responseDict
 
     @util.api()
     def storeMediaFile(self, filename, data=None, path=None, url=None, skipHash=None, deleteExisting=True):
diff --git a/tests/test_decks.py b/tests/test_decks.py
index c388b03..822ba1c 100755
--- a/tests/test_decks.py
+++ b/tests/test_decks.py
@@ -67,3 +67,8 @@ def test_removedDeckConfigId_fails_with_invalid_id(session_with_profile_loaded):
     new_config_id = ac.cloneDeckConfigId(cloneFrom=1, name="test")
     assert ac.removeDeckConfigId(configId=new_config_id) is True
     assert ac.removeDeckConfigId(configId=new_config_id) is False
+
+
+def test_getDeckStats(session_with_profile_loaded):
+    result = ac.getDeckStats(decks=["Default"])
+    assert result["name"] == "Default"