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)