bpo-47131: Speedup AST comparisons in test_unparse by using node traversal (GH-32132)

This commit is contained in:
Jeremy Kloth 2022-04-01 19:54:04 -06:00 committed by GitHub
parent b183f48649
commit 0f68c208fa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -130,7 +130,43 @@ docstring_prefixes = (
class ASTTestCase(unittest.TestCase):
def assertASTEqual(self, ast1, ast2):
self.assertEqual(ast.dump(ast1), ast.dump(ast2))
# Ensure the comparisons start at an AST node
self.assertIsInstance(ast1, ast.AST)
self.assertIsInstance(ast2, ast.AST)
# An AST comparison routine modeled after ast.dump(), but
# instead of string building, it traverses the two trees
# in lock-step.
def traverse_compare(a, b, missing=object()):
if type(a) is not type(b):
self.fail(f"{type(a)!r} is not {type(b)!r}")
if isinstance(a, ast.AST):
for field in a._fields:
value1 = getattr(a, field, missing)
value2 = getattr(b, field, missing)
# Singletons are equal by definition, so further
# testing can be skipped.
if value1 is not value2:
traverse_compare(value1, value2)
elif isinstance(a, list):
try:
for node1, node2 in zip(a, b, strict=True):
traverse_compare(node1, node2)
except ValueError:
# Attempt a "pretty" error ala assertSequenceEqual()
len1 = len(a)
len2 = len(b)
if len1 > len2:
what = "First"
diff = len1 - len2
else:
what = "Second"
diff = len2 - len1
msg = f"{what} list contains {diff} additional elements."
raise self.failureException(msg) from None
elif a != b:
self.fail(f"{a!r} != {b!r}")
traverse_compare(ast1, ast2)
def check_ast_roundtrip(self, code1, **kwargs):
with self.subTest(code1=code1, ast_parse_kwargs=kwargs):