Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions aymurai/transforms/anonymization_postprocess/core.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import re
from copy import deepcopy
from string import punctuation

from aymurai.meta.pipeline_interfaces import Transform
from aymurai.meta.types import DataItem
from aymurai.utils.misc import get_element
from aymurai.meta.pipeline_interfaces import Transform


class AnonymizationEntityCleaner(Transform):
Expand Down Expand Up @@ -45,6 +44,9 @@ def process(self, ent: dict) -> dict:
# Clean the text
cleaned_text = pattern.sub("", original_text)

if not cleaned_text:
return None
Comment on lines +47 to +48

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: Align the return type annotation and docstring with the new None return path.

The method can now return None when cleaned_text is empty, but the type hint and docstring still indicate it always returns a dict. Please update the signature (e.g. -> dict | None) and docstring to match, or avoid returning None by returning ent and handling the filtering elsewhere.

Suggested implementation:

class AnonymizationEntityCleaner(Transform):
    def _clean_entity(
        self,
        ent: dict,
        pattern,
        original_text: str,
        start_char: int,
        leading_chars_removed: int,
    ) -> dict | None:
        """Clean a single anonymization entity.

        The entity's alternative text and index attributes are updated in place.

        Args:
            ent: The entity to clean.
            pattern: Compiled regex used to remove unwanted text.
            original_text: The original text for the entity.
            start_char: The original start character index of the entity.
            leading_chars_removed: Number of leading characters removed from the original text.

        Returns:
            dict | None: The updated entity, or ``None`` if no cleaned text remains.
        """
        # Clean the text
        cleaned_text = pattern.sub("", original_text)

        if not cleaned_text:
            return None

        # Update the entity's alt text and indices
        ent["attrs"]["aymurai_alt_text"] = cleaned_text
        ent["attrs"]["aymurai_alt_start_char"] = start_char + leading_chars_removed
        return ent

    def __call__(self, item: DataItem) -> DataItem:
        """
        Post-process anonymization entities in the given item.

        Returns:
            DataItem: processed item
        """
        item = deepcopy(item)
        ents = get_element(item, [self.field, "entities"]) or []

I had to reconstruct some surrounding context because only a fragment of the class was provided:

  1. If the helper method currently has a different name or signature than _clean_entity(...): -> dict, adjust the def _clean_entity(...) line in the replacement block to match your existing parameters, and only change the return annotation to -> dict | None.
  2. If there is already a __call__ (or other transform entrypoint) method in AnonymizationEntityCleaner, remove or merge the __call__ implementation I added so you do not duplicate the method. Keep your existing logic, only updating any calls to the helper method to handle the new None return (e.g. by filtering out None values).
  3. Ensure your project’s minimum Python version supports dict | None. If you are constrained to Python < 3.10, change it to -> Optional[dict] and import Optional from typing.


Comment on lines +47 to +49
# Update the entity's alt text and indices
ent["attrs"]["aymurai_alt_text"] = cleaned_text
ent["attrs"]["aymurai_alt_start_char"] = start_char + leading_chars_removed
Expand All @@ -61,11 +63,11 @@ def __call__(self, item: DataItem) -> DataItem:
DataItem: processed item
"""
item = deepcopy(item)

ents = get_element(item, [self.field, "entities"]) or []

# Filter out predictions that are punctuation marks only
ents = [ent for ent in ents if ent["text"] not in punctuation]
ents = [self.process(ent) for ent in ents]
# Filter out predictions with empty alt text and update the rest
item[self.field]["entities"] = [
Comment on lines 66 to +69

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

issue (bug_risk): Guard against missing self.field / entities when assigning back to item.

get_element can return None when self.field or its "entities" key is missing, but the new code always does item[self.field]["entities"] = ..., which will raise in that case. Since the previous version only read from ents, this wasn’t exposed before. If these keys might be absent, initialize them first (e.g. item.setdefault(self.field, {}).setdefault("entities", [])) before assigning.

Comment on lines +68 to +69
out for ent in ents if (out := self.process(ent)) is not None
]
Comment on lines +68 to +71

return item
Loading