diff --git a/graphtage/ast.py b/graphtage/ast.py index aeabb13..6c87f58 100644 --- a/graphtage/ast.py +++ b/graphtage/ast.py @@ -87,7 +87,7 @@ def print(self, printer: Printer): self.value.print(printer) with printer.color(Fore.LIGHTBLUE_EX): printer.write("[") - self.slice.write(printer) + self.slice.print(printer) with printer.color(Fore.LIGHTBLUE_EX): printer.write("]") diff --git a/graphtage/pydiff.py b/graphtage/pydiff.py index ac4c358..b92d8f5 100644 --- a/graphtage/pydiff.py +++ b/graphtage/pydiff.py @@ -433,7 +433,7 @@ def print_Subscript(self, printer: Printer, node: Subscript): printer.write("[") self.print(printer, node.slice) with printer.color(Fore.BLUE): - printer.write("[") + printer.write("]") def diff(from_py_obj, to_py_obj, options: BuildOptions | None = None): diff --git a/test/test_pydiff.py b/test/test_pydiff.py index d64f261..effb501 100644 --- a/test/test_pydiff.py +++ b/test/test_pydiff.py @@ -4,6 +4,7 @@ from unittest import TestCase import graphtage +from graphtage.ast import Subscript from graphtage.printer import Printer from graphtage.pydiff import PyDiffFormatter, PyObjAttribute, ast_to_tree, build_tree, print_diff @@ -71,6 +72,28 @@ def test_nested_tuple_diff(self): self.assertIsInstance(kvp.key, graphtage.StringNode) self.assertIsInstance(kvp.value, graphtage.ListNode) + def _only_subscript(self, source: str) -> Subscript: + subscripts = [node for node in ast_to_tree(ast.parse(source)).dfs() if isinstance(node, Subscript)] + self.assertEqual(1, len(subscripts)) + return subscripts[0] + + def test_subscript_node_print(self): + """Reproduces the ``TreeNode.write`` half of https://github.com/trailofbits/graphtage/issues/154 + + ``Subscript.print`` is the fallback that :meth:`graphtage.tree.GraphtageFormatter.print` uses when no + formatter resolves the node type, so this calls it directly rather than through a formatter. + + """ + stream = StringIO() + self._only_subscript("a[1]").print(Printer(out_stream=stream, ansi_color=False)) + self.assertEqual("a[1]", stream.getvalue()) + + def test_subscript_formatter_print(self): + """Reproduces the unbalanced bracket half of https://github.com/trailofbits/graphtage/issues/154""" + stream = StringIO() + node = self._only_subscript("a[1]") + PyDiffFormatter.DEFAULT_INSTANCE.print(Printer(out_stream=stream, ansi_color=False), node) + self.assertEqual("a[1]", stream.getvalue()) def test_dataclass_edit_preserves_syntax(self): """Reproduces https://github.com/trailofbits/graphtage/issues/150