python (3.12.0)
1 import ast
2
3 class ESC[4;38;5;81mASTTestMixin:
4 """Test mixing to have basic assertions for AST nodes."""
5
6 def assertASTEqual(self, ast1, ast2):
7 # Ensure the comparisons start at an AST node
8 self.assertIsInstance(ast1, ast.AST)
9 self.assertIsInstance(ast2, ast.AST)
10
11 # An AST comparison routine modeled after ast.dump(), but
12 # instead of string building, it traverses the two trees
13 # in lock-step.
14 def traverse_compare(a, b, missing=object()):
15 if type(a) is not type(b):
16 self.fail(f"{type(a)!r} is not {type(b)!r}")
17 if isinstance(a, ast.AST):
18 for field in a._fields:
19 value1 = getattr(a, field, missing)
20 value2 = getattr(b, field, missing)
21 # Singletons are equal by definition, so further
22 # testing can be skipped.
23 if value1 is not value2:
24 traverse_compare(value1, value2)
25 elif isinstance(a, list):
26 try:
27 for node1, node2 in zip(a, b, strict=True):
28 traverse_compare(node1, node2)
29 except ValueError:
30 # Attempt a "pretty" error ala assertSequenceEqual()
31 len1 = len(a)
32 len2 = len(b)
33 if len1 > len2:
34 what = "First"
35 diff = len1 - len2
36 else:
37 what = "Second"
38 diff = len2 - len1
39 msg = f"{what} list contains {diff} additional elements."
40 raise self.failureException(msg) from None
41 elif a != b:
42 self.fail(f"{a!r} != {b!r}")
43 traverse_compare(ast1, ast2)