added support for the count field to be omitted if 1

Author: jtauber

README.md

@@ -98,9 +98,11 @@ You can iterate over the document-term counts at the leaves of the HDTM with the
 
 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.
 
+The third (count) field in a loaded file can be omitted if the count is 1 and a document ID + term may be repeated with the counts accumulating.
+
 ### Duplicates Policy
 
-You can optionall pass in a `duplicates` setting to the constructorr indicating the policy you want to follow if a term-document count is updated more than once.
+You can optionally pass in a `duplicates` setting to the constructorr indicating the policy you want to follow if a term-document count is updated more than once.
 
 ```python
 >>> c = termdoc.HTDM()

termdoc/htdm.py

@@ -64,9 +64,16 @@ def increment_count(self, address, term, count):
     def load(self, filename, field_sep="\t", address_sep="."):
         with open(filename) as f:
             for line in f:
-                address_string, term, count_string = line.strip().split(field_sep)
+                fields = line.strip().split(field_sep)
+                if len(fields) == 3:
+                    address_string, term, count_string = fields
+                    count = int(count_string)
+                elif len(fields) == 2:
+                    address_string, term = fields
+                    count = 1
+                else:
+                    raise ValueError(f"{fields} should have 2 or 3 fields")
                 address = tuple(address_string.split(address_sep))
-                count = int(count_string)
                 self.increment_count(address, term, count)
 
     def get_counts(self, prefix=()):

test_data/test2.tsv

@@ -0,0 +1,20 @@
+1.1	foo
+1.1	foo
+1.1	foo
+1.1	foo
+1.1	foo
+1.1	foo
+1.1	foo
+1.1	bar
+1.1	bar
+1.1	bar
+1.1	bar
+1.2	foo
+1.2	foo
+1.3	bar
+2.1	baz
+2.1	baz
+2.1	baz
+2.1	baz
+2.1	baz
+2.1	foo

test_data/test3e.tsv

@@ -0,0 +1 @@
+1

test_data/test4e.tsv

@@ -0,0 +1 @@
+1	2	3	4

tests.py

@@ -31,7 +31,7 @@ def test_multi_level(self):
         self.assertEqual(c.get_counts()["bar"], 6)
         self.assertEqual(c.get_counts((2,))["foo"], 4)
 
-    def test_load(self):
+    def test_load1(self):
         import termdoc
 
         c = termdoc.HTDM()
@@ -59,6 +59,46 @@ def test_load(self):
         self.assertEqual(c.get_counts(("2", "1"))["bar"], 0)
         self.assertEqual(c.get_counts(("2", "1"))["baz"], 5)
 
+    def test_load2(self):
+        import termdoc
+
+        c = termdoc.HTDM()
+        c.load("test_data/test2.tsv")
+
+        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(("1",))["foo"], 9)
+        self.assertEqual(c.get_counts(("1",))["bar"], 5)
+        self.assertEqual(c.get_counts(("1",))["baz"], 0)
+        self.assertEqual(c.get_counts(("2",))["foo"], 1)
+        self.assertEqual(c.get_counts(("2",))["bar"], 0)
+        self.assertEqual(c.get_counts(("2",))["baz"], 5)
+        self.assertEqual(c.get_counts(("1", "1"))["foo"], 7)
+        self.assertEqual(c.get_counts(("1", "1"))["bar"], 4)
+        self.assertEqual(c.get_counts(("1", "1"))["baz"], 0)
+        self.assertEqual(c.get_counts(("1", "2"))["foo"], 2)
+        self.assertEqual(c.get_counts(("1", "2"))["bar"], 0)
+        self.assertEqual(c.get_counts(("1", "2"))["baz"], 0)
+        self.assertEqual(c.get_counts(("1", "3"))["foo"], 0)
+        self.assertEqual(c.get_counts(("1", "3"))["bar"], 1)
+        self.assertEqual(c.get_counts(("1", "3"))["baz"], 0)
+        self.assertEqual(c.get_counts(("2", "1"))["foo"], 1)
+        self.assertEqual(c.get_counts(("2", "1"))["bar"], 0)
+        self.assertEqual(c.get_counts(("2", "1"))["baz"], 5)
+
+    def test_load3(self):
+        import termdoc
+
+        c = termdoc.HTDM()
+        self.assertRaises(ValueError, c.load, "test_data/test3e.tsv")
+
+    def test_load4(self):
+        import termdoc
+
+        c = termdoc.HTDM()
+        self.assertRaises(ValueError, c.load, "test_data/test4e.tsv")
+
     def test_prune(self):
         import termdoc