Preserve node syntax when printing data class edits

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`"""

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"""