EVOLUTION-MANAGER

Edit File: annotate_test.py

# coding=utf-8
"""Tests for annotate."""
# Copyright 2017 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import ast
import difflib
import itertools
import os.path
from six import with_metaclass
import sys
import textwrap
import unittest

import pasta
from pasta.base import annotate
from pasta.base import ast_utils
from pasta.base import codegen
from pasta.base import formatting as fmt
from pasta.base import test_utils

TESTDATA_DIR = os.path.realpath(
    os.path.join(os.path.dirname(pasta.__file__), '../testdata'))

class PrefixSuffixTest(test_utils.TestCase):

def test_block_suffix(self):
    src_tpl = textwrap.dedent('''\
        {open_block}
          pass #a
          #b
            #c

#d
        #e
        a
        ''')
    test_cases = (
        # first: attribute of the node with the last block
        # second: code snippet to open a block
        ('body', 'def x():'),
        ('body', 'class X:'),
        ('body', 'if x:'),
        ('orelse', 'if x:\n  y\nelse:'),
        ('body', 'if x:\n  y\nelif y:'),
        ('body', 'while x:'),
        ('orelse', 'while x:\n  y\nelse:'),
        ('finalbody', 'try:\n  x\nfinally:'),
        ('body', 'try:\n  x\nexcept:'),
        ('orelse', 'try:\n  x\nexcept:\n  y\nelse:'),
        ('body', 'with x:'),
        ('body', 'with x, y:'),
        ('body', 'with x:\n with y:'),
        ('body', 'for x in y:'),
    )
    def is_node_for_suffix(node, children_attr):
      # Return True if this node contains the 'pass' statement
      val = getattr(node, children_attr, None)
      return isinstance(val, list) and type(val[0]) == ast.Pass

for children_attr, open_block in test_cases:
      src = src_tpl.format(open_block=open_block)
      t = pasta.parse(src)
      node_finder = ast_utils.FindNodeVisitor(
          lambda node: is_node_for_suffix(node, children_attr))
      node_finder.visit(t)
      node = node_finder.results[0]
      expected = '  #b\n    #c\n\n  #d\n'
      actual = str(fmt.get(node, 'block_suffix_%s' % children_attr))
      self.assertMultiLineEqual(
          expected, actual,
          'Incorrect suffix for code:\n%s\nNode: %s (line %d)\nDiff:\n%s' % (
              src, node, node.lineno, '\n'.join(_get_diff(actual, expected))))
      self.assertMultiLineEqual(src, pasta.dump(t))

def test_module_suffix(self):
    src = 'foo\n#bar\n\n#baz\n'
    t = pasta.parse(src)
    self.assertEqual(src[src.index('#bar'):], fmt.get(t, 'suffix'))

def test_no_block_suffix_for_single_line_statement(self):
    src = 'if x:  return y\n  #a\n#b\n'
    t = pasta.parse(src)
    self.assertIsNone(fmt.get(t.body[0], 'block_suffix_body'))

def test_expression_prefix_suffix(self):
    src = 'a\n\nfoo\n\n\nb\n'
    t = pasta.parse(src)
    self.assertEqual('\n', fmt.get(t.body[1], 'prefix'))
    self.assertEqual('\n', fmt.get(t.body[1], 'suffix'))

def test_statement_prefix_suffix(self):
    src = 'a\n\ndef foo():\n  return bar\n\n\nb\n'
    t = pasta.parse(src)
    self.assertEqual('\n', fmt.get(t.body[1], 'prefix'))
    self.assertEqual('', fmt.get(t.body[1], 'suffix'))

class IndentationTest(test_utils.TestCase):

def test_indent_levels(self):
    src = textwrap.dedent('''\
        foo('begin')
        if a:
          foo('a1')
          if b:
            foo('b1')
            if c:
              foo('c1')
            foo('b2')
          foo('a2')
        foo('end')
        ''')
    t = pasta.parse(src)
    call_nodes = ast_utils.find_nodes_by_type(t, (ast.Call,))
    call_nodes.sort(key=lambda node: node.lineno)
    begin, a1, b1, c1, b2, a2, end = call_nodes

self.assertEqual('', fmt.get(begin, 'indent'))
    self.assertEqual('  ', fmt.get(a1, 'indent'))
    self.assertEqual('    ', fmt.get(b1, 'indent'))
    self.assertEqual('      ', fmt.get(c1, 'indent'))
    self.assertEqual('    ', fmt.get(b2, 'indent'))
    self.assertEqual('  ', fmt.get(a2, 'indent'))
    self.assertEqual('', fmt.get(end, 'indent'))

def test_indent_levels_same_line(self):
    src = 'if a: b; c\n'
    t = pasta.parse(src)
    if_node = t.body[0]
    b, c = if_node.body
    self.assertIsNone(fmt.get(b, 'indent_diff'))
    self.assertIsNone(fmt.get(c, 'indent_diff'))

def test_indent_depths(self):
    template = 'if a:\n{first}if b:\n{first}{second}foo()\n'
    indents = (' ', ' ' * 2, ' ' * 4, ' ' * 8, '\t', '\t' * 2)

for first, second in itertools.product(indents, indents):
      src = template.format(first=first, second=second)
      t = pasta.parse(src)
      outer_if_node = t.body[0]
      inner_if_node = outer_if_node.body[0]
      call_node = inner_if_node.body[0]

self.assertEqual('', fmt.get(outer_if_node, 'indent'))
      self.assertEqual('', fmt.get(outer_if_node, 'indent_diff'))
      self.assertEqual(first, fmt.get(inner_if_node, 'indent'))
      self.assertEqual(first, fmt.get(inner_if_node, 'indent_diff'))
      self.assertEqual(first + second, fmt.get(call_node, 'indent'))
      self.assertEqual(second, fmt.get(call_node, 'indent_diff'))

def test_indent_multiline_string(self):
    src = textwrap.dedent('''\
        class A:
          """Doc
             string."""
          pass
        ''')
    t = pasta.parse(src)
    docstring, pass_stmt = t.body[0].body
    self.assertEqual('  ', fmt.get(docstring, 'indent'))
    self.assertEqual('  ', fmt.get(pass_stmt, 'indent'))

def test_indent_multiline_string_with_newline(self):
    src = textwrap.dedent('''\
        class A:
          """Doc\n
             string."""
          pass
        ''')
    t = pasta.parse(src)
    docstring, pass_stmt = t.body[0].body
    self.assertEqual('  ', fmt.get(docstring, 'indent'))
    self.assertEqual('  ', fmt.get(pass_stmt, 'indent'))

def test_scope_trailing_comma(self):
    template = 'def foo(a, b{trailing_comma}): pass'
    for trailing_comma in ('', ',', ' , '):
      tree = pasta.parse(template.format(trailing_comma=trailing_comma))
      self.assertEqual(trailing_comma.lstrip(' ') + ')',
                       fmt.get(tree.body[0], 'args_suffix'))

template = 'class Foo(a, b{trailing_comma}): pass'
    for trailing_comma in ('', ',', ' , '):
      tree = pasta.parse(template.format(trailing_comma=trailing_comma))
      self.assertEqual(trailing_comma.lstrip(' ') + ')',
                       fmt.get(tree.body[0], 'bases_suffix'))

template = 'from mod import (a, b{trailing_comma})'
    for trailing_comma in ('', ',', ' , '):
      tree = pasta.parse(template.format(trailing_comma=trailing_comma))
      self.assertEqual(trailing_comma + ')',
                       fmt.get(tree.body[0], 'names_suffix'))

def test_indent_extra_newlines(self):
    src = textwrap.dedent('''\
        if a:

b
        ''')
    t = pasta.parse(src)
    if_node = t.body[0]
    b = if_node.body[0]
    self.assertEqual('  ', fmt.get(b, 'indent_diff'))

def test_indent_extra_newlines_with_comment(self):
    src = textwrap.dedent('''\
        if a:
            #not here

b
        ''')
    t = pasta.parse(src)
    if_node = t.body[0]
    b = if_node.body[0]
    self.assertEqual('  ', fmt.get(b, 'indent_diff'))

def test_autoindent(self):
    src = textwrap.dedent('''\
        def a():
            b
            c
        ''')
    expected = textwrap.dedent('''\
        def a():
            b
            new_node
        ''')
    t = pasta.parse(src)
    # Repace the second node and make sure the indent level is corrected
    t.body[0].body[1] = ast.Expr(ast.Name(id='new_node'))
    self.assertMultiLineEqual(expected, codegen.to_str(t))

@test_utils.requires_features('mixed_tabs_spaces')
  def test_mixed_tabs_spaces_indentation(self):
    pasta.parse(textwrap.dedent('''\
        if a:
                b
        {ONETAB}c
        ''').format(ONETAB='\t'))

@test_utils.requires_features('mixed_tabs_spaces')
  def test_tab_below_spaces(self):
    for num_spaces in range(1, 8):
      t = pasta.parse(textwrap.dedent('''\
          if a:
          {WS}if b:
          {ONETAB}c
          ''').format(ONETAB='\t', WS=' ' * num_spaces))
      node_c = t.body[0].body[0].body[0]
      self.assertEqual(fmt.get(node_c, 'indent_diff'), ' ' * (8 - num_spaces))

@test_utils.requires_features('mixed_tabs_spaces')
  def test_tabs_below_spaces_and_tab(self):
    for num_spaces in range(1, 8):
      t = pasta.parse(textwrap.dedent('''\
          if a:
          {WS}{ONETAB}if b:
          {ONETAB}{ONETAB}c
          ''').format(ONETAB='\t', WS=' ' * num_spaces))
      node_c = t.body[0].body[0].body[0]
      self.assertEqual(fmt.get(node_c, 'indent_diff'), '\t')

def _is_syntax_valid(filepath):
  with open(filepath, 'r') as f:
    try:
      ast.parse(f.read())
    except SyntaxError:
      return False
  return True

class SymmetricTestMeta(type):

def __new__(mcs, name, bases, inst_dict):
    # Helper function to generate a test method
    def symmetric_test_generator(filepath):
      def test(self):
        with open(filepath, 'r') as handle:
          src = handle.read()
        t = ast_utils.parse(src)
        annotator = annotate.AstAnnotator(src)
        annotator.visit(t)
        self.assertMultiLineEqual(codegen.to_str(t), src)
        self.assertEqual([], annotator.tokens._parens, 'Unmatched parens')
      return test

# Add a test method for each input file
    test_method_prefix = 'test_symmetric_'
    data_dir = os.path.join(TESTDATA_DIR, 'ast')
    for dirpath, dirs, files in os.walk(data_dir):
      for filename in files:
        if filename.endswith('.in'):
          full_path = os.path.join(dirpath, filename)
          inst_dict[test_method_prefix + filename[:-3]] = unittest.skipIf(
                    not _is_syntax_valid(full_path),
                    'Test contains syntax not supported by this version.',
                  )(symmetric_test_generator(full_path))
    return type.__new__(mcs, name, bases, inst_dict)

class SymmetricTest(with_metaclass(SymmetricTestMeta, test_utils.TestCase)):
  """Validates the symmetry property.

After parsing + annotating a module, regenerating the source code for it
  should yield the same result.
  """

def _get_node_identifier(node):
  for attr in ('id', 'name', 'attr', 'arg', 'module'):
    if isinstance(getattr(node, attr, None), str):
      return getattr(node, attr, '')
  return ''

class PrefixSuffixGoldenTestMeta(type):

def __new__(mcs, name, bases, inst_dict):
    # Helper function to generate a test method
    def golden_test_generator(input_file, golden_file):
      def test(self):
        with open(input_file, 'r') as handle:
          src = handle.read()
        t = ast_utils.parse(src)
        annotator = annotate.AstAnnotator(src)
        annotator.visit(t)

def escape(s):
          return '' if s is None else s.replace('\n', '\\n')

result = '\n'.join(
            "{0:12} {1:20} \tprefix=|{2}|\tsuffix=|{3}|\tindent=|{4}|".format(
                str((getattr(n, 'lineno', -1), getattr(n, 'col_offset', -1))),
                type(n).__name__ + ' ' + _get_node_identifier(n),
                escape(fmt.get(n, 'prefix')),
                escape(fmt.get(n, 'suffix')),
                escape(fmt.get(n, 'indent')))
            for n in ast.walk(t)) + '\n'

# If specified, write the golden data instead of checking it
        if getattr(self, 'generate_goldens', False):
          if not os.path.isdir(os.path.dirname(golden_file)):
            os.makedirs(os.path.dirname(golden_file))
          with open(golden_file, 'w') as f:
            f.write(result)
          print('Wrote: ' + golden_file)
          return

try:
          with open(golden_file, 'r') as f:
            golden = f.read()
        except IOError:
          self.fail('Missing golden data.')

self.assertMultiLineEqual(golden, result)
      return test

# Add a test method for each input file
    test_method_prefix = 'test_golden_prefix_suffix_'
    data_dir = os.path.join(TESTDATA_DIR, 'ast')
    python_version = '%d.%d' % sys.version_info[:2]
    for dirpath, dirs, files in os.walk(data_dir):
      for filename in files:
        if filename.endswith('.in'):
          full_path = os.path.join(dirpath, filename)
          golden_path = os.path.join(dirpath, 'golden', python_version,
                                     filename[:-3] + '.out')
          inst_dict[test_method_prefix + filename[:-3]] = unittest.skipIf(
                    not _is_syntax_valid(full_path),
                    'Test contains syntax not supported by this version.',
                  )(golden_test_generator(full_path, golden_path))
    return type.__new__(mcs, name, bases, inst_dict)

class PrefixSuffixGoldenTest(with_metaclass(PrefixSuffixGoldenTestMeta,
                                            test_utils.TestCase)):
  """Checks the prefix and suffix on each node in the AST.

This uses golden files in testdata/ast/golden. To regenerate these files, run
  python setup.py test -s pasta.base.annotate_test.generate_goldens
  """

maxDiff = None

class ManualEditsTest(test_utils.TestCase):
  """Tests that we can handle ASTs that have been modified.

Such ASTs may lack position information (lineno/col_offset) on some nodes.
  """

def test_call_no_pos(self):
    """Tests that Call node traversal works without position information."""
    src = 'f(a)'
    t = pasta.parse(src)
    node = ast_utils.find_nodes_by_type(t, (ast.Call,))[0]
    node.keywords.append(ast.keyword(arg='b', value=ast.Num(n=0)))
    self.assertEqual('f(a, b=0)', pasta.dump(t))

def test_call_illegal_pos(self):
    """Tests that Call node traversal works even with illegal positions."""
    src = 'f(a)'
    t = pasta.parse(src)
    node = ast_utils.find_nodes_by_type(t, (ast.Call,))[0]
    node.keywords.append(ast.keyword(arg='b', value=ast.Num(n=0)))

# This position would put b=0 before a, so it should be ignored.
    node.keywords[-1].value.lineno = 0
    node.keywords[-1].value.col_offset = 0

self.assertEqual('f(a, b=0)', pasta.dump(t))

class FstringTest(test_utils.TestCase):
  """Tests fstring support more in-depth."""

@test_utils.requires_features('fstring')
  def test_fstring(self):
    src = 'f"a {b} c d {e}"'
    t = pasta.parse(src)
    node = t.body[0].value
    self.assertEqual(
        fmt.get(node, 'content'),
        'f"a {__pasta_fstring_val_0__} c d {__pasta_fstring_val_1__}"')

@test_utils.requires_features('fstring')
  def test_fstring_escaping(self):
    src = 'f"a {{{b} {{c}}"'
    t = pasta.parse(src)
    node = t.body[0].value
    self.assertEqual(
        fmt.get(node, 'content'),
        'f"a {{{__pasta_fstring_val_0__} {{c}}"')

def _get_diff(before, after):
  return difflib.ndiff(after.splitlines(), before.splitlines())

def suite():
  result = unittest.TestSuite()
  result.addTests(unittest.makeSuite(ManualEditsTest))
  result.addTests(unittest.makeSuite(SymmetricTest))
  result.addTests(unittest.makeSuite(PrefixSuffixTest))
  result.addTests(unittest.makeSuite(PrefixSuffixGoldenTest))
  result.addTests(unittest.makeSuite(FstringTest))
  return result

def generate_goldens():
  result = unittest.TestSuite()
  result.addTests(unittest.makeSuite(PrefixSuffixGoldenTest))
  setattr(PrefixSuffixGoldenTest, 'generate_goldens', True)
  return result

if __name__ == '__main__':
  unittest.main()