Repository: jtauber/termdoc

Issue:
allow loading at arbitrary point in tree
`load` should take an option `prefix` which gets prepended to all addresses loaded

Comments:

Patch:
diff --git a/README.md b/README.md
index c3e2ffe..2e85876 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@ pip install termdoc
 
 ## HTDM
 
-The core data structure is a term-document matrix supporting hierarchical documents. Documents are labelled with a tuple such as `(1, 7, 5)` or `("Plato", "Republic", "Book 5")` (the type of each item in the tuple does not matter). This hierarchy could represent divisions of a work, grouping of multiple works, or some combination. Counts are aggregated at each level of the hierarchy (including at the top level to get totals across all documents).
+The core data structure is a term-document matrix supporting hierarchical documents. Documents are labelled with a delimited string such as "1.7.5" or "Plato.Republic.5". This hierarchy could represent divisions of a work, grouping of multiple works, or some combination. Counts are aggregated at each level of the hierarchy (including at the top level to get totals across all documents).
 
 HTDMs can be loaded with `load`:
 
diff --git a/termdoc/htdm.py b/termdoc/htdm.py
index d68eb0b..e44c6f9 100644
--- a/termdoc/htdm.py
+++ b/termdoc/htdm.py
@@ -42,7 +42,7 @@ class HTDM:
             address = self.address_sep.join(address.split(self.address_sep)[:-1])
             first = False
 
-    def load(self, filename, field_sep="\t", address_sep=None):
+    def load(self, filename, field_sep="\t", address_sep=None, prefix=None):
         address_sep = address_sep or self.address_sep
         with open(filename) as f:
             for line in f:
@@ -55,6 +55,8 @@ class HTDM:
                     count = 1
                 else:
                     raise ValueError(f"{fields} should have 2 or 3 fields")
+                if prefix:
+                    address = prefix + address_sep + address
                 self.increment_count(address, term, count)
 
     def get_counts(self, prefix=""):


Test patch:
diff --git a/tests.py b/tests.py
index 8859571..4a531fe 100755
--- a/tests.py
+++ b/tests.py
@@ -99,6 +99,25 @@ class Test1(unittest.TestCase):
         c = termdoc.HTDM()
         self.assertRaises(ValueError, c.load, "test_data/test4e.tsv")
 
+    def test_load5(self):
+        import termdoc
+
+        c = termdoc.HTDM()
+        c.load("test_data/test2.tsv", prefix="xxx")
+
+        self.assertEqual(c.get_counts()["foo"], 10)
+        self.assertEqual(c.get_counts()["bar"], 5)
+        self.assertEqual(c.get_counts()["baz"], 5)
+        self.assertEqual(c.get_counts("xxx")["foo"], 10)
+        self.assertEqual(c.get_counts("xxx")["bar"], 5)
+        self.assertEqual(c.get_counts("xxx")["baz"], 5)
+        self.assertEqual(c.get_counts("xxx.1")["foo"], 9)
+        self.assertEqual(c.get_counts("xxx.1")["bar"], 5)
+        self.assertEqual(c.get_counts("xxx.1")["baz"], 0)
+        self.assertEqual(c.get_counts("xxx.2")["foo"], 1)
+        self.assertEqual(c.get_counts("xxx.2")["bar"], 0)
+        self.assertEqual(c.get_counts("xxx.2")["baz"], 5)
+
     def test_prune(self):
         import termdoc