import ast class ASTTestMixin: """Test mixing to have basic assertions for AST nodes.""" def assertASTEqual(self, ast1, 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)