Repository: jtauber/termdoc

Issue:
add a strict mode to load that prevents duplicate counts

Discussion:

Patch:
diff --git a/README.md b/README.md
index a720141..9ad2bae 100644
--- a/README.md
+++ b/README.md
@@ -97,3 +97,45 @@ You can **prune** a HTDM to just `n` levels with the method `prune(n)`.
 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.
+
+### 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.
+
+```python
+>>> c = termdoc.HTDM()
+>>> c.increment_count((), "foo", 3)
+>>> c.increment_count((), "foo", 2)
+>>> c.get_counts()["foo"]
+5
+```
+
+is the same as
+
+```python
+>>> c = termdoc.HTDM(duplicates=termdoc.Duplicates.ALLOW)
+>>> c.increment_count((), "foo", 3)
+>>> c.increment_count((), "foo", 2)
+>>> c.get_counts()["foo"]
+5
+```
+
+But you can tell an HTDM to ignore attempts to update an existing count:
+
+```python
+>>> c = termdoc.HTDM(duplicates=termdoc.Duplicates.IGNORE)
+>>> c.increment_count((), "foo", 3)
+>>> c.increment_count((), "foo", 2)
+>>> c.get_counts()["foo"]
+3
+```
+
+or to raise an exception:
+
+```python
+>>> c = termdoc.HTDM(duplicates=termdoc.Duplicates.ERROR)
+>>> c.increment_count((), "foo", 3)
+>>> c.increment_count((), "foo", 2)  # this will raise a ValueError
+```
+
+Note that duplicates are only checked in the leaves of the document tree.
diff --git a/termdoc/__init__.py b/termdoc/__init__.py
index 2c09770..69142a9 100644
--- a/termdoc/__init__.py
+++ b/termdoc/__init__.py
@@ -1,1 +1,1 @@
-from .htdm import HTDM  # noqa: F401
+from .htdm import HTDM, Duplicates  # noqa: F401
diff --git a/termdoc/htdm.py b/termdoc/htdm.py
index 6493735..15c6e3e 100644
--- a/termdoc/htdm.py
+++ b/termdoc/htdm.py
@@ -1,9 +1,17 @@
 import collections
+from enum import Enum, auto
+
+
+class Duplicates(Enum):
+    ALLOW = auto()
+    IGNORE = auto()
+    ERROR = auto()
 
 
 class HTDM:
-    def __init__(self):
+    def __init__(self, duplicates=Duplicates.ALLOW):
         self.counters = []
+        self.duplicates = duplicates
 
     def get_or_create_counter(self, depth):
         while depth > len(self.counters) - 1:
@@ -11,12 +19,20 @@ class HTDM:
         return self.counters[depth]
 
     def increment_count(self, address, term, count):
+        first = True
         while True:
             depth = len(address)
-            self.get_or_create_counter(depth)[address][term] += count
+            counter = self.get_or_create_counter(depth)[address]
+            if first and term in counter:
+                if self.duplicates == Duplicates.IGNORE:
+                    return
+                elif self.duplicates == Duplicates.ERROR:
+                    raise ValueError(f"{term} already in {address}")
+            counter[term] += count
             if depth == 0:
                 break
             address = address[:-1]
+            first = False
 
     def load(self, filename, field_sep="\t", address_sep="."):
         with open(filename) as f:


Test patch:
diff --git a/tests.py b/tests.py
index 54805b9..ac12be5 100755
--- a/tests.py
+++ b/tests.py
@@ -118,5 +118,56 @@ class Test1(unittest.TestCase):
         )
 
 
+class TestDuplicates(unittest.TestCase):
+    def test_allow(self):
+        import termdoc
+
+        c = termdoc.HTDM()
+        c.increment_count((), "foo", 3)
+        c.increment_count((), "foo", 2)
+        self.assertEqual(c.get_counts()["foo"], 5)
+
+    def test_explicit_allow(self):
+        import termdoc
+
+        c = termdoc.HTDM(duplicates=termdoc.Duplicates.ALLOW)
+        c.increment_count((), "foo", 3)
+        c.increment_count((), "foo", 2)
+        self.assertEqual(c.get_counts()["foo"], 5)
+
+    def test_ignore(self):
+        import termdoc
+
+        c = termdoc.HTDM(duplicates=termdoc.Duplicates.IGNORE)
+        c.increment_count((), "foo", 3)
+        c.increment_count((), "foo", 2)
+        self.assertEqual(c.get_counts()["foo"], 3)
+
+    def test_error(self):
+        import termdoc
+
+        c = termdoc.HTDM(duplicates=termdoc.Duplicates.ERROR)
+        c.increment_count((), "foo", 3)
+        self.assertRaises(ValueError, c.increment_count, (), "foo", 2)
+        self.assertEqual(c.get_counts()["foo"], 3)
+
+    def test_multi_level_error(self):
+        import termdoc
+
+        # this test makes sure the `if first` is working in `increment_count`
+        c = termdoc.HTDM(duplicates=termdoc.Duplicates.ERROR)
+        c.increment_count((1, 1), "foo", 3)
+        c.increment_count((1, 2), "foo", 3)
+        c.increment_count((1, 1), "bar", 2)
+        c.increment_count((1, 2), "bar", 2)
+        c.increment_count((2, 1), "foo", 2)
+        c.increment_count((2, 2), "foo", 2)
+        c.increment_count((2, 1), "bar", 1)
+        c.increment_count((2, 2), "bar", 1)
+        self.assertEqual(c.get_counts()["foo"], 10)
+        self.assertEqual(c.get_counts()["bar"], 6)
+        self.assertEqual(c.get_counts((2,))["foo"], 4)
+
+
 if __name__ == "__main__":
     unittest.main()