Repository: jtauber/termdoc

Issue:
implement copying of one hierarchical counter as a subtree of another

Discussion:

Patch:
diff --git a/README.md b/README.md
index 9d1d41d..ffed1b7 100644
--- a/README.md
+++ b/README.md
@@ -77,4 +77,10 @@ with a period-separated hierarchical document ID, term, and count all separated
 
 Both the period and tab are just defaults and can be override by padding `address_sep` and/or `field_sep` to `load`.
 
-You can prune a HTDM to just `n` levels with the method `prune(n)`.
+You can **prune** a HTDM to just `n` levels with the method `prune(n)`.
+
+## Unreleased Features
+
+You can iterate over the document-term counts at the leaves of the HDTM with the method `leaf_entries()` (this returns a generator yielding `(document_address, term, count)` tuples). This is effectively a traditional TDM (the document IDs will still reflect the hierarchy but the aggregate counts aren't present).
+
+You can **graft** one HTDM under another by using the `graft(prefix, subtree)` method, specifying as `prefix` the document address you want to add the subtree under. This is useful if you have an HTDM for, say, a single work by an author, with chapters as documents and you want to incorporate that into a higher-level HTDM of multiple works by the author, or a collection of works by different authors.
diff --git a/termdoc/htdm.py b/termdoc/htdm.py
index ce4aa6a..6493735 100644
--- a/termdoc/htdm.py
+++ b/termdoc/htdm.py
@@ -32,3 +32,15 @@ class HTDM:
 
     def prune(self, level):
         self.counters = self.counters[:level]
+
+    def leaves(self):
+        return self.counters[-1]
+
+    def leaf_entries(self):
+        for document, counter in self.leaves().items():
+            for term, count in counter.items():
+                yield document, term, count
+
+    def graft(self, prefix, subtree):
+        for address, term, count in subtree.leaf_entries():
+            self.increment_count(prefix + address, term, count)


Test patch:
diff --git a/tests.py b/tests.py
index 49ccfe8..fe7ab9b 100755
--- a/tests.py
+++ b/tests.py
@@ -75,6 +75,48 @@ class Test1(unittest.TestCase):
         self.assertEqual(c.get_counts()["bar"], 3)
         self.assertEqual(len(c.counters), 1)
 
+    def test5(self):
+        import termdoc
+
+        c = termdoc.HTDM()
+        c.increment_count((1,), "foo", 3)
+        c.increment_count((1,), "bar", 2)
+        c.increment_count((2,), "foo", 2)
+        c.increment_count((2,), "bar", 1)
+        self.assertEqual(
+            list(c.leaf_entries()),
+            [
+                ((1,), "foo", 3),
+                ((1,), "bar", 2),
+                ((2,), "foo", 2),
+                ((2,), "bar", 1),
+            ],
+        )
+
+    def test6(self):
+        import termdoc
+
+        c1 = termdoc.HTDM()
+        c1.increment_count((1,), "foo", 3)
+        c1.increment_count((1,), "bar", 2)
+        c1.increment_count((2,), "foo", 2)
+        c1.increment_count((2,), "bar", 1)
+        c2 = termdoc.HTDM()
+        c2.increment_count((4, 1), "foo", 5)
+        c2.increment_count((5, 3), "foo", 6)
+        c2.graft((5,), c1)
+        self.assertEqual(
+            list(c2.leaf_entries()),
+            [
+                ((4, 1), "foo", 5),
+                ((5, 3), "foo", 6),
+                ((5, 1), "foo", 3),
+                ((5, 1), "bar", 2),
+                ((5, 2), "foo", 2),
+                ((5, 2), "bar", 1),
+            ],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()