diff --git a/graphtage/dataclasses.py b/graphtage/dataclasses.py index dc68544..1fbf5a0 100644 --- a/graphtage/dataclasses.py +++ b/graphtage/dataclasses.py @@ -3,7 +3,7 @@ from . import AbstractCompoundEdit, Edit, Range, Replace from .printer import Fore, Printer -from .tree import ContainerNode, TreeNode +from .tree import ContainerNode, GraphtageFormatter, TreeNode class DataClassEdit(AbstractCompoundEdit): @@ -30,6 +30,20 @@ def edits(self) -> Iterator[Edit]: def tighten_bounds(self) -> bool: return any(edit.tighten_bounds() for edit in self.slot_edits) + def print(self, formatter: GraphtageFormatter, printer: Printer): + """Prints this edit by delegating to the formatter for the node being edited. + + The default :meth:`graphtage.AbstractCompoundEdit.print` implementation prints the slot edits back to back, + which drops whatever syntax the node's formatter writes between the slots. Delegating to the node formatter + keeps that syntax, and the formatter reaches the slot edits as it prints each child. + + This is equivalent to:: + + formatter.get_formatter(self.from_node)(printer, self.from_node) + + """ + formatter.get_formatter(self.from_node)(printer, self.from_node) + class DataClassNode(ContainerNode): """A container node that can be initialized similar to a Python :func:`dataclasses.dataclass`""" diff --git a/test/test_pydiff.py b/test/test_pydiff.py index ed70fd8..bc43c10 100644 --- a/test/test_pydiff.py +++ b/test/test_pydiff.py @@ -1,13 +1,26 @@ import ast import dataclasses +from io import StringIO from unittest import TestCase import graphtage +from graphtage.printer import Printer from graphtage.pydiff import PyDiffFormatter, ast_to_tree, build_tree, print_diff from .timing import run_with_time_limit +def render_diff(from_source: str, to_source: str) -> str: + """Diffs two Python sources and returns the rendering produced by :class:`PyDiffFormatter`.""" + from_tree = ast_to_tree(ast.parse(from_source)) + to_tree = ast_to_tree(ast.parse(to_source)) + stream = StringIO() + printer = Printer(out_stream=stream, ansi_color=False) + with printer: + PyDiffFormatter.DEFAULT_INSTANCE.print(printer, from_tree.diff(to_tree)) + return stream.getvalue().strip() + + class TestPyDiff(TestCase): def test_build_tree(self): self.assertIsInstance(build_tree([1, 2, 3, 4]), graphtage.ListNode) @@ -47,6 +60,21 @@ def test_nested_tuple_diff(self): self.assertIsInstance(kvp.key, graphtage.StringNode) self.assertIsInstance(kvp.value, graphtage.ListNode) + def test_dataclass_edit_preserves_syntax(self): + """Reproduces https://github.com/trailofbits/graphtage/issues/150 + + Without `DataClassEdit.print`, an edited `DataClassNode` falls back to `AbstractCompoundEdit.print`, which + prints the slot edits back to back and drops the syntax the node formatter writes between them. The first + case below rendered as `[x]foo[1,2,++3++]`. + """ + self.assertEqual("x = foo(1, 2, ++3++)", render_diff("x = foo(1, 2)", "x = foo(1, 2, 3)")) + self.assertEqual("x -> y = foo(1, 2)", render_diff("x = foo(1, 2)", "y = foo(1, 2)")) + self.assertEqual("x = ~~foo~~++bar++(1, 2)", render_diff("x = foo(1, 2)", "x = bar(1, 2)")) + + def test_attribute_edit_preserves_separators(self): + """An edited `PyObjAttribute` keeps the dots between its slots; it used to render as `[x]abc -> d`.""" + self.assertEqual("x = a.b.c -> d", render_diff("x = a.b.c", "x = a.b.d")) + def test_infinite_loop(self): """Reproduces https://github.com/trailofbits/graphtage/issues/82"""