root _text preservation/inference on transforms; atoms with the same root keep tok_pos/text_span

Author: telmomenezes

src/hyperbase/hyperedge.py

@@ -233,9 +233,9 @@ def replace_atom(
         unique -- match only the exact same instance of the atom, i.e.
         UniqueAtom(self) == UniqueAtom(old) (default: False)
         """
-        from hyperbase.transforms import replace_atom
+        from hyperbase.transforms import _propagate_root_text, replace_atom
 
-        return replace_atom(self, old, new, unique=unique)
+        return _propagate_root_text(self, replace_atom(self, old, new, unique=unique))
 
     def simplify(self, subtypes: bool = False, namespaces: bool = False) -> Hyperedge:
         """Returns a version of the edge with simplified atoms.
@@ -244,9 +244,11 @@ def simplify(self, subtypes: bool = False, namespaces: bool = False) -> Hyperedg
         subtypes -- include subtypes (default: True).
         namespaces -- include namespaces (default: True).
         """
-        from hyperbase.transforms import simplify
+        from hyperbase.transforms import _propagate_root_text, simplify
 
-        return simplify(self, subtypes=subtypes, namespaces=namespaces)
+        return _propagate_root_text(
+            self, simplify(self, subtypes=subtypes, namespaces=namespaces)
+        )
 
     def transform(
         self,
@@ -255,9 +257,12 @@ def transform(
         recursive: bool = True,
     ) -> Hyperedge:
         """Pattern-driven rewrite. See ``hyperbase.transforms.transform``."""
-        from hyperbase.transforms import transform
+        from hyperbase.transforms import _propagate_root_text, transform
 
-        return transform(self, origin_pattern, target_pattern, recursive=recursive)
+        return _propagate_root_text(
+            self,
+            transform(self, origin_pattern, target_pattern, recursive=recursive),
+        )
 
     def type(self) -> str:
         """Returns the type of this edge as a string.
@@ -376,27 +381,27 @@ def replace_argroles(self, argroles: str | None) -> Hyperedge:
         """Returns an edge with the argroles of the connector atom replaced
         with the provided string.
         Returns same edge if the atom does not contain a role part."""
-        from hyperbase.transforms import replace_argroles
+        from hyperbase.transforms import _propagate_root_text, replace_argroles
 
-        return replace_argroles(self, argroles)
+        return _propagate_root_text(self, replace_argroles(self, argroles))
 
     def _insert_argrole(self, argrole: str, pos: int) -> Hyperedge:
         """Returns an edge with the given argrole inserted at the specified
         position in the argroles of the connector atom.
         Same restrictions as in replace_argroles() apply."""
-        from hyperbase.transforms import insert_argrole
+        from hyperbase.transforms import _propagate_root_text, insert_argrole
 
-        return insert_argrole(self, argrole, pos)
+        return _propagate_root_text(self, insert_argrole(self, argrole, pos))
 
     def add_argument(
         self, edge: Hyperedge, argrole: str, pos: int | None = None
     ) -> Hyperedge:
         """Returns a new edge with the provided edge and its argroles inserted
         at the specified position. If pos is not provided, the argument is
         appended at the end."""
-        from hyperbase.transforms import add_argument
+        from hyperbase.transforms import _propagate_root_text, add_argument
 
-        return add_argument(self, edge, argrole, pos)
+        return _propagate_root_text(self, add_argument(self, edge, argrole, pos))
 
     def arguments_with_role(self, argrole: str) -> list[Hyperedge]:
         """Returns the list of edges with the given argument role."""
@@ -418,9 +423,9 @@ def check_correctness(self) -> dict[Hyperedge, list[tuple[str, str]]]:
         return check_correctness(self)
 
     def normalise(self) -> Hyperedge:
-        from hyperbase.transforms import normalise
+        from hyperbase.transforms import _propagate_root_text, normalise
 
-        return normalise(self)
+        return _propagate_root_text(self, normalise(self))
 
     ############
     # patterns #

src/hyperbase/transforms.py

@@ -75,19 +75,44 @@ def replace_atom(
 ) -> Hyperedge:
     """Return a copy with every instance of *old* replaced by *new*."""
     if edge.atom:
+        edge_atom = cast(Atom, edge)
         if unique:
-            if UniqueAtom(edge) == UniqueAtom(old):  # type: ignore[arg-type]
-                return new
+            matched = UniqueAtom(edge_atom) == UniqueAtom(old)  # type: ignore[arg-type]
         else:
-            if edge == old:
-                return new
+            matched = edge_atom == old
+        if matched:
+            return _maybe_inherit_atom_metadata(new, edge_atom)
         return edge
     else:
         return Hyperedge(
             tuple(replace_atom(item, old, new, unique=unique) for item in edge)
         )
 
 
+def _maybe_inherit_atom_metadata(new: Hyperedge, source: Atom) -> Hyperedge:
+    """If *new* is an atom with the same root as *source* and carries no
+    source-position metadata of its own, return a copy that inherits
+    ``tok_pos``/``text_span``/``text`` from *source*. Otherwise *new* is
+    returned unchanged.
+    """
+    if not new.atom:
+        return new
+    new_atom = cast(Atom, new)
+    if new_atom.tok_pos is not None or new_atom.text_span is not None:
+        return new
+    if new_atom.root() != source.root():
+        return new
+    if source.tok_pos is None and source.text_span is None and source._text is None:
+        return new
+    return Atom(
+        new_atom.atom_str,
+        new_atom.parens,
+        text=source._text,
+        tok_pos=source.tok_pos,
+        text_span=source.text_span,
+    )
+
+
 def replace_argroles(edge: Hyperedge, argroles: str | None) -> Hyperedge:
     """Return a copy with the connector's argument roles replaced."""
     if edge.atom:
@@ -328,18 +353,21 @@ def _instantiate(
 
 
 def _attach_metadata_from_original(target_atom: Atom, original: Hyperedge) -> Atom:
-    """If ``target_atom`` matches an atom in ``original``, return a copy with
-    that atom's tok_pos / text_span. Otherwise return ``target_atom`` as-is.
+    """If an atom with the same root as ``target_atom`` exists in ``original``,
+    return a copy that inherits its ``tok_pos``/``text_span``/``text``.
+    Otherwise return ``target_atom`` as-is.
     """
     if target_atom.tok_pos is not None or target_atom.text_span is not None:
         return target_atom
+    target_root = target_atom.root()
     for atom in _walk_origin_atoms(original):
-        if atom == target_atom and (
+        if atom.root() == target_root and (
             atom.tok_pos is not None or atom.text_span is not None
         ):
             return Atom(
                 target_atom.atom_str,
                 target_atom.parens,
+                text=atom._text,
                 tok_pos=atom.tok_pos,
                 text_span=atom.text_span,
             )
@@ -374,6 +402,7 @@ def _substitute_var_atom(
         new_parts = [binding_atom.root(), *target_parts[1:]]
         return Atom(
             "/".join(p for p in new_parts if p),
+            text=binding_atom._text,
             tok_pos=binding_atom.tok_pos,
             text_span=binding_atom.text_span,
         )
@@ -402,6 +431,7 @@ def _substitute_var_atom(
     if inner.type() == binding.type():
         new_inner = Atom(
             "/".join(p for p in [inner.root(), *target_parts[1:]] if p),
+            text=inner._text,
             tok_pos=inner.tok_pos,
             text_span=inner.text_span,
         )
@@ -463,3 +493,47 @@ def _consumed_arg_indices(
                 consumed.add(i)
                 break
     return consumed
+
+
+def _propagate_root_text(original: Hyperedge, result: Hyperedge) -> Hyperedge:
+    """Carry the original root's source-text fingerprint onto a transform's
+    result.
+
+    Mutates ``result``'s ``_text``/``tokens`` in place (safe: the result is
+    freshly constructed by the transform). When the original carries no
+    source text or when ``result is original``, the result is returned
+    unchanged. When the set of atom roots is unchanged, ``_text`` is
+    preserved verbatim. Otherwise it is derived as the contiguous slice of
+    the original text that spans the result atoms' ``text_span`` extents.
+    """
+    if result is original:
+        return result
+    if original.atom or result.atom:
+        return result
+    src_text = original.text
+    if src_text is None:
+        return result
+
+    orig_roots = sorted(a.root() for a in original.all_atoms())
+    new_roots = sorted(a.root() for a in result.all_atoms())
+
+    if orig_roots == new_roots:
+        new_text: str | None = src_text
+    else:
+        new_text = _derive_text_from_spans(result, src_text)
+        if new_text is None:
+            return result
+
+    object.__setattr__(result, "_text", new_text)
+    if original.tokens is not None:
+        object.__setattr__(result, "tokens", original.tokens)
+    return result
+
+
+def _derive_text_from_spans(edge: Hyperedge, source_text: str) -> str | None:
+    spans = [atom.text_span for atom in edge.all_atoms() if atom.text_span is not None]
+    if not spans:
+        return None
+    start = min(s[0] for s in spans)
+    end = max(s[1] for s in spans)
+    return source_text[start:end]

tests/test_hyperedge.py

@@ -1337,6 +1337,135 @@ def test_transform_partial_match_unchanged(self):
         )
         assert result == edge
 
+    def test_transform_preserves_root_text_same_roots(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(eats/Pd.so john/C apples/C)"),
+            text="John eats apples",
+            tokens=["John", "eats", "apples"],
+            tok_pos=hedge("(1 0 2)"),
+        )
+        loaded = hedge(pr)
+        # Decoration-only changes leave the atom roots untouched: text and
+        # source tokens flow through to the result root verbatim.
+        result = loaded.transform(
+            hedge("(eats/Pd.so X Y)"),
+            hedge("(eats/Pd.so X/Cp Y/Cp)"),
+        )
+        assert result.text == "John eats apples"
+        assert result.tokens == ("John", "eats", "apples")
+
+    def test_transform_derives_root_text_when_roots_change(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(eats/Pd.so john/C apples/C)"),
+            text="John eats apples",
+            tokens=["John", "eats", "apples"],
+            tok_pos=hedge("(1 0 2)"),
+        )
+        loaded = hedge(pr)
+        # Drop the predicate via the {} preserve form: the resulting root has
+        # a different atom-root set, so text is derived from atom spans.
+        result = loaded.transform(
+            hedge("(eats/Pd.so X Y)"),
+            hedge("Y"),
+        )
+        # Y -> apples/C with text_span (10, 16) in "John eats apples".
+        assert str(result) == "apples/C"
+        assert result.text == "apples"
+
+    def test_replace_argroles_preserves_root_text(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(is/P.sc john/C tired/C)"),
+            text="John is tired",
+            tokens=["John", "is", "tired"],
+            tok_pos=hedge("(1 0 2)"),
+        )
+        loaded = hedge(pr)
+        result = loaded.replace_argroles("os")
+        assert result.text == "John is tired"
+        assert result.tokens == ("John", "is", "tired")
+
+    def test_add_argument_preserves_root_text_when_arg_in_source(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(eats/Pd.so john/C)"),
+            text="John eats apples",
+            tokens=["John", "eats", "apples"],
+            tok_pos=hedge("(1 0)"),
+        )
+        loaded = hedge(pr)
+        # Build the new arg from the parse so it carries source spans.
+        apples_pr = ParseResult(
+            edge=hedge("apples/C"),
+            text="John eats apples",
+            tokens=["John", "eats", "apples"],
+            tok_pos=hedge("2"),
+        )
+        apples = hedge(apples_pr)
+        result = loaded.add_argument(apples, "o")
+        # Roots changed (added apples), so text is derived. Slice spans
+        # from John to apples — the full source text.
+        assert result.text == "John eats apples"
+
+    def test_transform_constant_inherits_metadata_by_root(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(yesterday/Mt died/Mn fido/Cc)"),
+            text="Yesterday died Fido",
+            tokens=["Yesterday", "died", "Fido"],
+            tok_pos=hedge("(0 1 2)"),
+        )
+        loaded = hedge(pr)
+        # died/Md (constant target) shares root with died/Mn in the original;
+        # tok_pos and text_span must be inherited even though the decoration
+        # changed.
+        result = loaded.transform(hedge("died/Mn"), hedge("died/Md"))
+        assert str(result) == "(yesterday/Mt died/Md fido/Cc)"
+        assert result[1].tok_pos == 1
+        assert result[1].text_span == (10, 14)
+
+    def test_replace_atom_inherits_metadata_when_same_root(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(eats/Pd.so john/C apples/C)"),
+            text="John eats apples",
+            tokens=["John", "eats", "apples"],
+            tok_pos=hedge("(1 0 2)"),
+        )
+        loaded = hedge(pr)
+        old = loaded[2]  # apples/C with tok_pos=2, text_span=(10, 16)
+        new = hedge("apples/Cp")  # bare atom, no metadata
+        result = loaded.replace_atom(old, new)
+        # Metadata flows through because the root matches.
+        assert str(result[2]) == "apples/Cp"
+        assert result[2].tok_pos == 2
+        assert result[2].text_span == (10, 16)
+
+    def test_replace_atom_no_inherit_when_root_differs(self):
+        from hyperbase.parsers.parse_result import ParseResult
+
+        pr = ParseResult(
+            edge=hedge("(eats/Pd.so john/C apples/C)"),
+            text="John eats apples",
+            tokens=["John", "eats", "apples"],
+            tok_pos=hedge("(1 0 2)"),
+        )
+        loaded = hedge(pr)
+        old = loaded[2]
+        new = hedge("oranges/C")
+        result = loaded.replace_atom(old, new)
+        assert str(result[2]) == "oranges/C"
+        assert result[2].tok_pos is None
+        assert result[2].text_span is None
+
 
 if __name__ == "__main__":
     unittest.main()