From 7f392dbc7983e44bbecfdfbd921be7101d78feeb Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 10:51:58 -0600 Subject: [PATCH 1/4] `DocumentDeDuplicator` is not a dict --- compass/pipeline/collection/base.py | 5 ++++- compass/pipeline/collection/dedupe.py | 16 +++------------- compass/pipeline/collection/persistence.py | 2 +- .../unit/pipeline/test_pipeline_collection.py | 2 +- .../pipeline/test_pipeline_collection_dedupe.py | 4 ++-- 5 files changed, 11 insertions(+), 18 deletions(-) diff --git a/compass/pipeline/collection/base.py b/compass/pipeline/collection/base.py index 3eb075480..bef1661d1 100644 --- a/compass/pipeline/collection/base.py +++ b/compass/pipeline/collection/base.py @@ -160,7 +160,10 @@ 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: diff --git a/compass/pipeline/collection/dedupe.py b/compass/pipeline/collection/dedupe.py index 416be6775..ba6877882 100644 --- a/compass/pipeline/collection/dedupe.py +++ b/compass/pipeline/collection/dedupe.py @@ -1,17 +1,15 @@ """Document deduplication for collected artifacts""" import logging +from collections import UserDict logger = logging.getLogger(__name__) -class DocumentDeDuplicator: +class DocumentDeDuplicator(UserDict): """Domain Service for deduplicating collected documents""" - def __init__(self): - self._docs = {} - def add_docs(self, docs, *, step_name=None): """Add documents to the collection mapping @@ -33,7 +31,7 @@ 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( + entry = self.data.setdefault( key, { "doc": doc, @@ -43,14 +41,6 @@ def add_docs(self, docs, *, step_name=None): 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) - def _collection_doc_key(doc_info): """Build the deduplication key for a collected document""" diff --git a/compass/pipeline/collection/persistence.py b/compass/pipeline/collection/persistence.py index 431d970c4..90a0e29d9 100644 --- a/compass/pipeline/collection/persistence.py +++ b/compass/pipeline/collection/persistence.py @@ -399,7 +399,7 @@ 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: + 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"]) diff --git a/tests/python/unit/pipeline/test_pipeline_collection.py b/tests/python/unit/pipeline/test_pipeline_collection.py index a1ff92453..9d3cb18c3 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection.py +++ b/tests/python/unit/pipeline/test_pipeline_collection.py @@ -49,7 +49,7 @@ 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: + for entry in deduplicator.values(): document = dict(entry["doc"].attrs) document["from_steps"] = list(entry["from_steps"]) documents.append(document) diff --git a/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py b/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py index 0d633f32e..b37ca862b 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py +++ b/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py @@ -22,7 +22,7 @@ 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"] == [ @@ -55,7 +55,7 @@ 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 From ea392c5dcbba2f2c3a6dcfc027fe9ba1c5b12610 Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 11:04:33 -0600 Subject: [PATCH 2/4] Add `_DocInfo` dataclass --- compass/pipeline/collection/base.py | 5 +-- compass/pipeline/collection/dedupe.py | 32 +++++++++++++------ compass/pipeline/collection/persistence.py | 12 +++---- .../unit/pipeline/test_pipeline_collection.py | 4 +-- .../test_pipeline_collection_dedupe.py | 9 ++---- 5 files changed, 35 insertions(+), 27 deletions(-) diff --git a/compass/pipeline/collection/base.py b/compass/pipeline/collection/base.py index bef1661d1..e271147b0 100644 --- a/compass/pipeline/collection/base.py +++ b/compass/pipeline/collection/base.py @@ -160,10 +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: diff --git a/compass/pipeline/collection/dedupe.py b/compass/pipeline/collection/dedupe.py index ba6877882..5bcdc10d1 100644 --- a/compass/pipeline/collection/dedupe.py +++ b/compass/pipeline/collection/dedupe.py @@ -2,11 +2,32 @@ import logging from collections import UserDict +from dataclasses import dataclass + +from elm.web.document import BaseDocument logger = logging.getLogger(__name__) +@dataclass(frozen=True) +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) + + @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""" @@ -31,15 +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.data.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) + doc_info = self.data.setdefault(key, _DocInfo.from_doc(doc)) + doc_info.add_step(step_name) def _collection_doc_key(doc_info): diff --git a/compass/pipeline/collection/persistence.py b/compass/pipeline/collection/persistence.py index 90a0e29d9..420428429 100644 --- a/compass/pipeline/collection/persistence.py +++ b/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 @@ -400,9 +400,9 @@ async def _store_docs_as_needed(collected_docs, jurisdiction, relative_to): document_metadata = [] left_to_store = [] for info in collected_docs.values(): - doc = info["doc"] + 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, diff --git a/tests/python/unit/pipeline/test_pipeline_collection.py b/tests/python/unit/pipeline/test_pipeline_collection.py index 9d3cb18c3..c698ea8bb 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection.py +++ b/tests/python/unit/pipeline/test_pipeline_collection.py @@ -50,8 +50,8 @@ 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"]) + document = dict(entry.doc.attrs) + document["from_steps"] = list(entry.from_steps) documents.append(document) collection_info = { "documents": documents, diff --git a/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py b/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py index b37ca862b..9f6c22d0d 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py +++ b/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py @@ -25,7 +25,7 @@ def test_add_docs_keeps_from_steps_unique_for_same_doc_and_step(): values = list(deduplicator.values()) assert len(values) == 1 - assert values[0]["from_steps"] == [ + assert values[0].from_steps == [ "Look for document on jurisdiction website" ] @@ -58,11 +58,8 @@ def test_add_docs_preserves_restored_artifacts_and_merges_provenance(): 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__": From b1545609b805b24fff98fd7d8583a2beb07d5be2 Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 12:06:14 -0600 Subject: [PATCH 3/4] fix docs --- docs/source/conf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index e81090d2e..d5e48580f 100644 --- a/docs/source/conf.py +++ b/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): From a9890fa965def4173485a7a5dadd6d3d027c92f8 Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 12:06:19 -0600 Subject: [PATCH 4/4] PR review --- compass/pipeline/collection/dedupe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compass/pipeline/collection/dedupe.py b/compass/pipeline/collection/dedupe.py index 5bcdc10d1..c4ab41b63 100644 --- a/compass/pipeline/collection/dedupe.py +++ b/compass/pipeline/collection/dedupe.py @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) -@dataclass(frozen=True) +@dataclass class _DocInfo: """Information about a collected document"""