diff --git a/compass/pipeline/collection/base.py b/compass/pipeline/collection/base.py index 3eb075480..e271147b0 100644 --- a/compass/pipeline/collection/base.py +++ b/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: diff --git a/compass/pipeline/collection/dedupe.py b/compass/pipeline/collection/dedupe.py index 416be6775..c4ab41b63 100644 --- a/compass/pipeline/collection/dedupe.py +++ b/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): diff --git a/compass/pipeline/collection/persistence.py b/compass/pipeline/collection/persistence.py index 431d970c4..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 @@ -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, 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): diff --git a/tests/python/unit/pipeline/test_pipeline_collection.py b/tests/python/unit/pipeline/test_pipeline_collection.py index a1ff92453..c698ea8bb 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection.py +++ b/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, diff --git a/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py b/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py index 0d633f32e..9f6c22d0d 100644 --- a/tests/python/unit/pipeline/test_pipeline_collection_dedupe.py +++ b/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__":