Dedupe refactor

compass/pipeline/collection/base.py

@@ -160,7 +160,7 @@ async def execute(self, *, eager_extract=False):
                 "Collected the following documents for %s:\n\n%s",
                 self.workflow.jurisdiction.full_name,
                 "\n\n".join(
-                    [f"{info['doc']!r}" for info in self.de_duplicator.values]
+                    [f"{info.doc!r}" for info in self.de_duplicator.values()]
                 ),
             )
         else:

compass/pipeline/collection/dedupe.py

@@ -1,16 +1,35 @@
 """Document deduplication for collected artifacts"""
 
 import logging
+from collections import UserDict
+from dataclasses import dataclass
+
+from elm.web.document import BaseDocument
 
 
 logger = logging.getLogger(__name__)
 
 
-class DocumentDeDuplicator:
-    """Domain Service for deduplicating collected documents"""
+@dataclass
+class _DocInfo:
+    """Information about a collected document"""
+
+    doc: BaseDocument
+    from_steps: list[str]
+
+    def add_step(self, step_name: str | None):
+        """Add a collection step to the provenance of this document"""
+        if step_name and step_name not in self.from_steps:
+            self.from_steps.append(step_name)
 
-    def __init__(self):
-        self._docs = {}
+    @classmethod
+    def from_doc(cls, doc: BaseDocument):
+        """Create a new _DocInfo from a document"""
+        return cls(doc=doc, from_steps=list(doc.attrs.get("from_steps", [])))
+
+
+class DocumentDeDuplicator(UserDict):
+    """Domain Service for deduplicating collected documents"""
 
     def add_docs(self, docs, *, step_name=None):
         """Add documents to the collection mapping
@@ -33,23 +52,8 @@ def add_docs(self, docs, *, step_name=None):
         logger.debug("Adding %d doc(s) to collection", len(docs))
         for doc in docs:
             key = _collection_doc_key(doc.attrs)
-            entry = self._docs.setdefault(
-                key,
-                {
-                    "doc": doc,
-                    "from_steps": list(doc.attrs.get("from_steps", [])),
-                },
-            )
-            if step_name and step_name not in entry["from_steps"]:
-                entry["from_steps"].append(step_name)
-
-    @property
-    def values(self):
-        """Deduplicated collected docs"""
-        return self._docs.values()
-
-    def __bool__(self):
-        return bool(self._docs)
+            doc_info = self.data.setdefault(key, _DocInfo.from_doc(doc))
+            doc_info.add_step(step_name)
 
 
 def _collection_doc_key(doc_info):

compass/pipeline/collection/persistence.py

@@ -300,8 +300,8 @@ async def persist_documents(
         Jurisdiction whose deduplicated documents will be persisted and
         serialized into collection metadata.
     collected_docs : compass.pipeline.collection.dedupe.DocumentDeDuplicator
-        Deduplicated document collection containing ``{"doc",
-        "from_steps"}`` entries for each persisted document.
+        Deduplicated document collection containing document info
+        entries for each persisted document.
     completed_steps : iterable of str
         Collection step names that were completed for this jurisdiction,
         used to record the ``"completed_step_document_counts"`` in the
@@ -399,10 +399,10 @@ async def _store_docs_as_needed(collected_docs, jurisdiction, relative_to):
     """Store collected documents and their parsed text when needed"""
     document_metadata = []
     left_to_store = []
-    for info in collected_docs.values:
-        doc = info["doc"]
+    for info in collected_docs.values():
+        doc = info.doc
         if "parsed_fp" in doc.attrs and "source_fp" in doc.attrs:
-            doc.attrs["from_steps"] = list(info["from_steps"])
+            doc.attrs["from_steps"] = list(info.from_steps)
             document_metadata.append(doc.attrs)
         else:
             left_to_store.append(info)
@@ -413,9 +413,9 @@ async def _store_docs_as_needed(collected_docs, jurisdiction, relative_to):
     ):
         task = asyncio.create_task(
             _persist_doc(
-                info["doc"],
+                info.doc,
                 out_stem=f"{jurisdiction.full_name}_{index}",
-                from_steps=info["from_steps"],
+                from_steps=info.from_steps,
                 relative_to=relative_to,
             ),
             name=jurisdiction.full_name,

docs/source/conf.py

@@ -295,7 +295,7 @@ def _skip_builtin_methods(name, obj):
     if name in {"items", "keys", "values"} and "Mapping" in str(obj):
         return True
 
-    return name in {"copy", "get"} and "UserDict" in str(obj)
+    return name in {"copy", "get", "fromkeys"} and "UserDict" in str(obj)
 
 
 def _skip_internal_api(name, obj):

tests/python/unit/pipeline/test_pipeline_collection.py

@@ -49,9 +49,9 @@ async def _load_existing_collection_shard():  # ruff:ignore[unused-async]
     async def _write_collection_shard_no_fail(deduplicator, completed_steps):
         await asyncio.sleep(0)
         documents = []
-        for entry in deduplicator.values:
-            document = dict(entry["doc"].attrs)
-            document["from_steps"] = list(entry["from_steps"])
+        for entry in deduplicator.values():
+            document = dict(entry.doc.attrs)
+            document["from_steps"] = list(entry.from_steps)
             documents.append(document)
         collection_info = {
             "documents": documents,

tests/python/unit/pipeline/test_pipeline_collection_dedupe.py

@@ -22,10 +22,10 @@ def test_add_docs_keeps_from_steps_unique_for_same_doc_and_step():
         step_name="Look for document on jurisdiction website",
     )
 
-    values = list(deduplicator.values)
+    values = list(deduplicator.values())
 
     assert len(values) == 1
-    assert values[0]["from_steps"] == [
+    assert values[0].from_steps == [
         "Look for document on jurisdiction website"
     ]
 
@@ -55,14 +55,11 @@ def test_add_docs_preserves_restored_artifacts_and_merges_provenance():
         step_name="search_engine",
     )
 
-    values = list(deduplicator.values)
+    values = list(deduplicator.values())
 
     assert len(values) == 1
-    assert values[0]["doc"] is saved_doc
-    assert values[0]["from_steps"] == [
-        "known_local_docs",
-        "search_engine",
-    ]
+    assert values[0].doc is saved_doc
+    assert values[0].from_steps == ["known_local_docs", "search_engine"]
 
 
 if __name__ == "__main__":