add support for printing the diff of AST trees when running tests (#3902)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

docs/contributing/the_basics.md

@@ -37,6 +37,44 @@ the root of the black repo:
 (.venv)$ tox -e run_self
 ```
 
+### Development
+
+Further examples of invoking the tests
+
+```console
+# Run all of the above mentioned, in parallel
+(.venv)$ tox --parallel=auto
+
+# Run tests on a specific python version
+(.venv)$ tox -e py39
+
+# pass arguments to pytest
+(.venv)$ tox -e py -- --no-cov
+
+# print full tree diff, see documentation below
+(.venv)$ tox -e py -- --print-full-tree
+
+# disable diff printing, see documentation below
+(.venv)$ tox -e py -- --print-tree-diff=False
+```
+
+`Black` has two pytest command-line options affecting test files in `tests/data/` that
+are split into an input part, and an output part, separated by a line with`# output`.
+These can be passed to `pytest` through `tox`, or directly into pytest if not using
+`tox`.
+
+#### `--print-full-tree`
+
+Upon a failing test, print the full concrete syntax tree (CST) as it is after processing
+the input ("actual"), and the tree that's yielded after parsing the output ("expected").
+Note that a test can fail with different output with the same CST. This used to be the
+default, but now defaults to `False`.
+
+#### `--print-tree-diff`
+
+Upon a failing test, print the diff of the trees as described above. This is the
+default. To turn it off pass `--print-tree-diff=False`.
+
 ### News / Changelog Requirement
 
 `Black` has CI that will check for an entry corresponding to your PR in `CHANGES.md`. If

src/black/debug.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass
-from typing import Iterator, TypeVar, Union
+from dataclasses import dataclass, field
+from typing import Any, Iterator, List, TypeVar, Union
 
 from black.nodes import Visitor
 from black.output import out
@@ -14,26 +14,33 @@
 @dataclass
 class DebugVisitor(Visitor[T]):
     tree_depth: int = 0
+    list_output: List[str] = field(default_factory=list)
+    print_output: bool = True
+
+    def out(self, message: str, *args: Any, **kwargs: Any) -> None:
+        self.list_output.append(message)
+        if self.print_output:
+            out(message, *args, **kwargs)
 
     def visit_default(self, node: LN) -> Iterator[T]:
         indent = " " * (2 * self.tree_depth)
         if isinstance(node, Node):
             _type = type_repr(node.type)
-            out(f"{indent}{_type}", fg="yellow")
+            self.out(f"{indent}{_type}", fg="yellow")
             self.tree_depth += 1
             for child in node.children:
                 yield from self.visit(child)
 
             self.tree_depth -= 1
-            out(f"{indent}/{_type}", fg="yellow", bold=False)
+            self.out(f"{indent}/{_type}", fg="yellow", bold=False)
         else:
             _type = token.tok_name.get(node.type, str(node.type))
-            out(f"{indent}{_type}", fg="blue", nl=False)
+            self.out(f"{indent}{_type}", fg="blue", nl=False)
             if node.prefix:
                 # We don't have to handle prefixes for `Node` objects since
                 # that delegates to the first child anyway.
-                out(f" {node.prefix!r}", fg="green", bold=False, nl=False)
-            out(f" {node.value!r}", fg="blue", bold=False)
+                self.out(f" {node.prefix!r}", fg="green", bold=False, nl=False)
+            self.out(f" {node.value!r}", fg="blue", bold=False)
 
     @classmethod
     def show(cls, code: Union[str, Leaf, Node]) -> None:

tests/conftest.py

@@ -1 +1,28 @@
+import pytest
+
 pytest_plugins = ["tests.optional"]
+
+PRINT_FULL_TREE: bool = False
+PRINT_TREE_DIFF: bool = True
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--print-full-tree",
+        action="store_true",
+        default=False,
+        help="print full syntax trees on failed tests",
+    )
+    parser.addoption(
+        "--print-tree-diff",
+        action="store_true",
+        default=True,
+        help="print diff of syntax trees on failed tests",
+    )
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    global PRINT_FULL_TREE
+    global PRINT_TREE_DIFF
+    PRINT_FULL_TREE = config.getoption("--print-full-tree")
+    PRINT_TREE_DIFF = config.getoption("--print-tree-diff")

tests/test_black.py

@@ -9,7 +9,6 @@
 import re
 import sys
 import types
-import unittest
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager, redirect_stderr
 from dataclasses import replace
@@ -1047,9 +1046,10 @@ def test_endmarker(self) -> None:
         self.assertEqual(len(n.children), 1)
         self.assertEqual(n.children[0].type, black.token.ENDMARKER)
 
+    @patch("tests.conftest.PRINT_FULL_TREE", True)
+    @patch("tests.conftest.PRINT_TREE_DIFF", False)
     @pytest.mark.incompatible_with_mypyc
-    @unittest.skipIf(os.environ.get("SKIP_AST_PRINT"), "user set SKIP_AST_PRINT")
-    def test_assertFormatEqual(self) -> None:
+    def test_assertFormatEqual_print_full_tree(self) -> None:
         out_lines = []
         err_lines = []
 
@@ -1068,6 +1068,29 @@ def err(msg: str, **kwargs: Any) -> None:
         self.assertIn("Actual tree:", out_str)
         self.assertEqual("".join(err_lines), "")
 
+    @patch("tests.conftest.PRINT_FULL_TREE", False)
+    @patch("tests.conftest.PRINT_TREE_DIFF", True)
+    @pytest.mark.incompatible_with_mypyc
+    def test_assertFormatEqual_print_tree_diff(self) -> None:
+        out_lines = []
+        err_lines = []
+
+        def out(msg: str, **kwargs: Any) -> None:
+            out_lines.append(msg)
+
+        def err(msg: str, **kwargs: Any) -> None:
+            err_lines.append(msg)
+
+        with patch("black.output._out", out), patch("black.output._err", err):
+            with self.assertRaises(AssertionError):
+                self.assertFormatEqual("j = [1, 2, 3]\n", "j = [1, 2, 3,]\n")
+
+        out_str = "".join(out_lines)
+        self.assertIn("Tree Diff:", out_str)
+        self.assertIn("+          COMMA", out_str)
+        self.assertIn("+ ','", out_str)
+        self.assertEqual("".join(err_lines), "")
+
     @event_loop()
     @patch("concurrent.futures.ProcessPoolExecutor", MagicMock(side_effect=OSError))
     def test_works_in_mono_process_only_environment(self) -> None:

tests/util.py

@@ -12,6 +12,8 @@
 from black.mode import TargetVersion
 from black.output import diff, err, out
 
+from . import conftest
+
 PYTHON_SUFFIX = ".py"
 ALLOWED_SUFFIXES = (PYTHON_SUFFIX, ".pyi", ".out", ".diff", ".ipynb")
 
@@ -34,22 +36,34 @@
 
 
 def _assert_format_equal(expected: str, actual: str) -> None:
-    if actual != expected and not os.environ.get("SKIP_AST_PRINT"):
+    if actual != expected and (conftest.PRINT_FULL_TREE or conftest.PRINT_TREE_DIFF):
         bdv: DebugVisitor[Any]
-        out("Expected tree:", fg="green")
+        actual_out: str = ""
+        expected_out: str = ""
+        if conftest.PRINT_FULL_TREE:
+            out("Expected tree:", fg="green")
         try:
             exp_node = black.lib2to3_parse(expected)
-            bdv = DebugVisitor()
+            bdv = DebugVisitor(print_output=conftest.PRINT_FULL_TREE)
             list(bdv.visit(exp_node))
+            expected_out = "\n".join(bdv.list_output)
         except Exception as ve:
             err(str(ve))
-        out("Actual tree:", fg="red")
+        if conftest.PRINT_FULL_TREE:
+            out("Actual tree:", fg="red")
         try:
             exp_node = black.lib2to3_parse(actual)
-            bdv = DebugVisitor()
+            bdv = DebugVisitor(print_output=conftest.PRINT_FULL_TREE)
             list(bdv.visit(exp_node))
+            actual_out = "\n".join(bdv.list_output)
         except Exception as ve:
             err(str(ve))
+        if conftest.PRINT_TREE_DIFF:
+            out("Tree Diff:")
+            out(
+                diff(expected_out, actual_out, "expected tree", "actual tree")
+                or "Trees do not differ"
+            )
 
     if actual != expected:
         out(diff(expected, actual, "expected", "actual"))

tox.ini

@@ -1,6 +1,6 @@
 [tox]
 isolated_build = true
-envlist = {,ci-}py{37,38,39,310,311,py3},fuzz,run_self
+envlist = {,ci-}py{38,39,310,311,py3},fuzz,run_self
 
 [testenv]
 setenv =