Improve processor and processor group abstraction

vmenger · vmenger · Dec 20, 2023 · Nov 30, 2023 · Nov 30, 2023 · Dec 1, 2023
commit fbdbd55bc94d1d94147ad8073e0333c804598c2c
diff --git a/docdeid/process/doc.py b/docdeid/process/doc.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from collections import OrderedDict
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Union
 
 from docdeid.document import Document
 
@@ -19,15 +19,17 @@ def process(self, doc: Document, **kwargs) -> None:
         """
 
 
-class DocProcessorGroup(DocProcessor):
+class DocProcessorGroup:
     """
     A group of :class:`.DocProcessor`, that executes the containing processors in order.
 
     A :class:`.DocProcessorGroup` can itself be part of a :class:`.DocProcessorGroup`.
     """
 
     def __init__(self) -> None:
-        self._processors: OrderedDict[str, DocProcessor] = OrderedDict()
+        self._processors: OrderedDict[
+            str, Union[DocProcessor | DocProcessorGroup]
+        ] = OrderedDict()
 
     def get_names(self, recursive: bool = True) -> list[str]:
         """
@@ -52,7 +54,10 @@ def get_names(self, recursive: bool = True) -> list[str]:
         return names
 
     def add_processor(
-        self, name: str, processor: DocProcessor, position: Optional[int] = None
+        self,
+        name: str,
+        processor: Union[DocProcessor, "DocProcessorGroup"],
+        position: Optional[int] = None,
     ) -> None:
         """
         Add a document processor to the group.
@@ -90,7 +95,7 @@ def remove_processor(self, name: str) -> None:
         """
         del self._processors[name]
 
-    def __getitem__(self, name: str) -> DocProcessor:
+    def __getitem__(self, name: str) -> Union[DocProcessor, "DocProcessorGroup"]:
         """
         Get a document processor by name.
 
@@ -103,23 +108,25 @@ def __getitem__(self, name: str) -> DocProcessor:
 
         return self._processors[name]
 
-    def process(self, doc: Document, **kwargs) -> None:
+    def process(
+        self,
+        doc: Document,
+        enabled: Optional[set[str]] = None,
+        disabled: Optional[set[str]] = None,
+    ) -> None:
         """
-        Process a document, by passing it to this groups processors.
+        Process a document, by passing it to this group's processors.
 
         Args:
             doc: The document to be processed.
             enabled: A set of strings, indicating which document processors to run for
-            this document. By default all document processors are used. In case of
+            this document. By default, all document processors are used. In case of
             nested, it's necessary to supply both the name of the processor group,
-            as well as all of its containing processors (or a subset thereof).
+            and all of its containing processors (or a subset thereof).
             disabled: A set of strings, indicating which document processors not to
             run for this document. Cannot be used together with `enabled`.
         """
 
-        enabled = kwargs.get("enabled", None)
-        disabled = kwargs.get("disabled", None)
-
         if (enabled is not None) and (disabled is not None):
             raise RuntimeError("Cannot use enabled and disabled simultaneously")
 
@@ -131,7 +138,10 @@ def process(self, doc: Document, **kwargs) -> None:
             if (disabled is not None) and (name in disabled):
                 continue
 
-            proc.process(doc, enabled=enabled, disabled=disabled)
+            if isinstance(proc, DocProcessor):
+                proc.process(doc)
+            elif isinstance(proc, DocProcessorGroup):
+                proc.process(doc, enabled=enabled, disabled=disabled)
 
     def __iter__(self) -> Iterator:
         """