示例#1
0
 def test_parallel_walk_inconsistent_trees(self):
   node_1 = parser.parse_str(
       textwrap.dedent("""
     def f(a):
       return a + 1
   """))
   node_2 = parser.parse_str(
       textwrap.dedent("""
     def f(a):
       return a + (a * 2)
   """))
   node_3 = parser.parse_str(
       textwrap.dedent("""
     def f(a):
       return a + 2
   """))
   with self.assertRaises(ValueError):
     for _ in ast_util.parallel_walk(node_1, node_2):
       pass
   # There is not particular reason to reject trees that differ only in the
   # value of a constant.
   # TODO(mdan): This should probably be allowed.
   with self.assertRaises(ValueError):
     for _ in ast_util.parallel_walk(node_1, node_3):
       pass
示例#2
0
 def test_parallel_walk_inconsistent_trees(self):
     node_1 = parser.parse_str(
         textwrap.dedent("""
   def f(a):
     return a + 1
 """))
     node_2 = parser.parse_str(
         textwrap.dedent("""
   def f(a):
     return a + (a * 2)
 """))
     node_3 = parser.parse_str(
         textwrap.dedent("""
   def f(a):
     return a + 2
 """))
     with self.assertRaises(ValueError):
         for _ in ast_util.parallel_walk(node_1, node_2):
             pass
     # There is not particular reason to reject trees that differ only in the
     # value of a constant.
     # TODO(mdan): This should probably be allowed.
     with self.assertRaises(ValueError):
         for _ in ast_util.parallel_walk(node_1, node_3):
             pass
示例#3
0
 def test_parse_str(self):
     mod = parser.parse_str(
         textwrap.dedent("""
         def f(x):
           return x + 1
 """))
     self.assertEqual('f', mod.body[0].name)
示例#4
0
def _build_source_map(node, code):
  """Return the Python objects represented by given AST.

  Compiling the AST code this way ensures that the source code is readable by
  e.g. `pdb` or `inspect`.

  Args:
    node: An AST node of the original generated code, before the source code is
      generated.
    code: The string representation of the source code for the newly generated
      code.

  Returns:
    Dict[CodeLocation, OriginInfo], a mapping between the user and AutoGraph
    generated code.
  """
  # After we have the final generated code we reparse it to get the final line
  # numbers. Then we walk through the generated and original ASTs in parallel
  # to build the mapping between the user and generated code.
  new_node = parser.parse_str(code)
  origin_info.resolve(new_node, code)
  source_mapping = {}
  for before, after in ast_util.parallel_walk(node, new_node):
    # Need both checks because if origin information is ever copied over to new
    # nodes then we need to rely on the fact that only the original user code
    # has the origin annotation.
    if (anno.hasanno(before, anno.Basic.ORIGIN) and
        anno.hasanno(after, anno.Basic.ORIGIN)):
      source_info = anno.getanno(before, anno.Basic.ORIGIN)
      new_line_number = anno.getanno(after, anno.Basic.ORIGIN).line_number
      source_mapping[new_line_number] = source_info
  return source_mapping
示例#5
0
def replace(template, **replacements):
  """Replace placeholders in a Python template.

  AST Name and Tuple nodes always receive the context that inferred from
  the template. However, when replacing more complex nodes (that can potentially
  contain Name children), then the caller is responsible for setting the
  appropriate context.

  Args:
    template: A string representing Python code. Any symbol name can be used
        that appears in the template code can be used as placeholder.
    **replacements: A mapping from placeholder names to (lists of) AST nodes
        that these placeholders will be replaced by. String values are also
        supported as a shorthand for AST Name nodes with the respective ID.

  Returns:
    An AST node or list of AST nodes with the replacements made. If the
    template was a function, a list will be returned. If the template was a
    node, the same node will be returned. If the template was a string, an
    AST node will be returned (a `Module` node in the case of a multi-line
    string, an `Expr` node otherwise).

  Raises:
    ValueError: if the arguments are incorrect.
  """
  if not isinstance(template, str):
    raise ValueError('Expected string template, got %s' % type(template))
  tree = parser.parse_str(textwrap.dedent(template))
  for k in replacements:
    replacements[k] = _convert_to_ast(replacements[k])
  results = ReplaceTransformer(replacements).visit(tree).body
  if isinstance(results, list):
    return [qual_names.resolve(r) for r in results]
  return qual_names.resolve(results)
示例#6
0
def _build_source_map(node, code):
    """Return the Python objects represented by given AST.

  Compiling the AST code this way ensures that the source code is readable by
  e.g. `pdb` or `inspect`.

  Args:
    node: An AST node of the original generated code, before the source code is
      generated.
    code: The string representation of the source code for the newly generated
      code.

  Returns:
    Dict[CodeLocation, OriginInfo], a mapping between the user and AutoGraph
    generated code.
  """
    # After we have the final generated code we reparse it to get the final line
    # numbers. Then we walk through the generated and original ASTs in parallel
    # to build the mapping between the user and generated code.
    new_node = parser.parse_str(code)
    origin_info.resolve(new_node, code)
    source_mapping = {}
    for before, after in ast_util.parallel_walk(node, new_node):
        # Need both checks because if origin information is ever copied over to new
        # nodes then we need to rely on the fact that only the original user code
        # has the origin annotation.
        if (anno.hasanno(before, anno.Basic.ORIGIN)
                and anno.hasanno(after, anno.Basic.ORIGIN)):
            source_info = anno.getanno(before, anno.Basic.ORIGIN)
            new_line_number = anno.getanno(after,
                                           anno.Basic.ORIGIN).line_number
            source_mapping[new_line_number] = source_info
    return source_mapping
示例#7
0
def replace(template, **replacements):
    """Replace placeholders in a Python template.

  AST Name and Tuple nodes always receive the context that inferred from
  the template. However, when replacing more complex nodes (that can potentially
  contain Name children), then the caller is responsible for setting the
  appropriate context.

  Args:
    template: A string representing Python code. Any symbol name can be used
        that appears in the template code can be used as placeholder.
    **replacements: A mapping from placeholder names to (lists of) AST nodes
        that these placeholders will be replaced by. String values are also
        supported as a shorthand for AST Name nodes with the respective ID.

  Returns:
    An AST node or list of AST nodes with the replacements made. If the
    template was a function, a list will be returned. If the template was a
    node, the same node will be returned. If the template was a string, an
    AST node will be returned (a `Module` node in the case of a multi-line
    string, an `Expr` node otherwise).

  Raises:
    ValueError: if the arguments are incorrect.
  """
    if not isinstance(template, str):
        raise ValueError('Expected string template, got %s' % type(template))
    tree = parser.parse_str(textwrap.dedent(template))
    for k in replacements:
        replacements[k] = _convert_to_ast(replacements[k])
    results = ReplaceTransformer(replacements).visit(tree).body
    if isinstance(results, list):
        return [qual_names.resolve(r) for r in results]
    return qual_names.resolve(results)
    def test_subscript_resolve(self):
        samples = """
      x[i]
      x[i.b]
      a.b[c]
      a.b[x.y]
      a[z[c]]
      a[b[c[d]]]
      a[b].c
      a.b.c[d].e.f
      a.b[c[d]].e.f
      a.b[c[d.e.f].g].h
    """
        nodes = resolve(parser.parse_str(textwrap.dedent(samples)))
        nodes = tuple(n.value for n in nodes.body)

        self.assertQNStringIs(nodes[0], 'x[i]')
        self.assertQNStringIs(nodes[1], 'x[i.b]')
        self.assertQNStringIs(nodes[2], 'a.b[c]')
        self.assertQNStringIs(nodes[3], 'a.b[x.y]')
        self.assertQNStringIs(nodes[4], 'a[z[c]]')
        self.assertQNStringIs(nodes[5], 'a[b[c[d]]]')
        self.assertQNStringIs(nodes[6], 'a[b].c')
        self.assertQNStringIs(nodes[7], 'a.b.c[d].e.f')
        self.assertQNStringIs(nodes[8], 'a.b[c[d]].e.f')
        self.assertQNStringIs(nodes[9], 'a.b[c[d.e.f].g].h')
示例#9
0
 def test_parse_str(self):
   mod = parser.parse_str(
       textwrap.dedent("""
           def f(x):
             return x + 1
   """))
   self.assertEqual('f', mod.body[0].name)
示例#10
0
def to_graph(e,
             recursive=True,
             verbose=False,
             arg_values=None,
             arg_types=None,
             partial_types=None):
    """Compile a Python entity into equivalent TensorFlow code.

  Currently supported entities:
    * functions
    * classes

  Classes are handled by converting all their methods into a new class.

  Args:
    e: A Python entity.
    recursive: Whether to recusrively convert any functions that the decorator
        function may call.
    verbose: Whether to output the compiled code in the logs.
    arg_values: A dict containing value hints for symbols like function
        parameters.
    arg_types: A dict containing type hints for symbols like function
        parameters.
    partial_types: A set of types (e.g. classes) that will not be converted
        entirely. Calls to member functions for these types will be renamed
        independently.

  Returns:
    A function with a signature identical to `o`, but which when executed it
  creates TF a graph that has the same functionality as the original entity.
  """
    conversion_map = conversion.ConversionMap(
        recursive=recursive,
        nocompile_decorators=(convert, do_not_convert, converted_call),
        partial_types=partial_types,
        api_module=tf_inspect.getmodule(to_graph))
    _, name = conversion.entity_to_graph(e, conversion_map, arg_values,
                                         arg_types)

    module = gast.Module([])
    for import_line in config.COMPILED_IMPORT_STATEMENTS:
        module.body.extend(parser.parse_str(import_line).body)
    for dep in conversion_map.dependency_cache.values():
        module.body.append(dep)
    compiled_node, compiled_src = compiler.ast_to_object(module)

    # The compiled code should see everything the entry function saw.
    # TODO(mdan): This might not work well if the call tree spans modules?
    if tf_inspect.isfunction(e):
        for key, val in inspect_utils.getnamespace(e).items():
            # Avoid overwriting entities that have been transformed.
            if key not in compiled_node.__dict__:
                compiled_node.__dict__[key] = val
    compiled_fn = getattr(compiled_node, name)

    if verbose:
        logging.info('Compiled output of %s:\n\n%s\n', e, compiled_src)

    return compiled_fn
示例#11
0
 def test_parallel_walk(self):
   node = parser.parse_str(
       textwrap.dedent("""
     def f(a):
       return a + 1
   """))
   for child_a, child_b in ast_util.parallel_walk(node, node):
     self.assertEqual(child_a, child_b)
示例#12
0
 def test_parallel_walk(self):
     node = parser.parse_str(
         textwrap.dedent("""
   def f(a):
     return a + 1
 """))
     for child_a, child_b in ast_util.parallel_walk(node, node):
         self.assertEqual(child_a, child_b)
示例#13
0
def to_graph(e,
             recursive=True,
             verbose=False,
             arg_values=None,
             arg_types=None,
             partial_types=None):
  """Compile a Python entity into equivalent TensorFlow code.

  Currently supported entities:
    * functions
    * classes

  Classes are handled by converting all their methods into a new class.

  Args:
    e: A Python entity.
    recursive: Whether to recusrively convert any functions that the decorator
        function may call.
    verbose: Whether to output the compiled code in the logs.
    arg_values: A dict containing value hints for symbols like function
        parameters.
    arg_types: A dict containing type hints for symbols like function
        parameters.
    partial_types: A set of types (e.g. classes) that will not be converted
        entirely. Calls to member functions for these types will be renamed
        independently.

  Returns:
    A function with a signature identical to `o`, but which when executed it
  creates TF a graph that has the same functionality as the original entity.
  """
  conversion_map = conversion.ConversionMap(
      recursive=recursive,
      nocompile_decorators=(convert, do_not_convert, converted_call),
      partial_types=partial_types,
      api_module=tf_inspect.getmodule(to_graph))
  _, name = conversion.entity_to_graph(e, conversion_map, arg_values, arg_types)

  module = gast.Module([])
  for import_line in config.COMPILED_IMPORT_STATEMENTS:
    module.body.extend(parser.parse_str(import_line).body)
  for dep in conversion_map.dependency_cache.values():
    module.body.append(dep)
  compiled_node, compiled_src = compiler.ast_to_object(module)

  # The compiled code should see everything the entry function saw.
  # TODO(mdan): This might not work well if the call tree spans modules?
  if tf_inspect.isfunction(e):
    for key, val in inspect_utils.getnamespace(e).items():
      # Avoid overwriting entities that have been transformed.
      if key not in compiled_node.__dict__:
        compiled_node.__dict__[key] = val
  compiled_fn = getattr(compiled_node, name)

  if verbose:
    logging.info('Compiled output of %s:\n\n%s\n', e, compiled_src)

  return compiled_fn
示例#14
0
 def test_keywords_to_dict(self):
   keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
   d = ast_util.keywords_to_dict(keywords)
   # Make sure we generate a usable dict node by attaching it to a variable and
   # compiling everything.
   output = parser.parse_str('b = 3')
   output.body += (ast.Assign([ast.Name(id='d', ctx=ast.Store())], d),)
   result, _ = compiler.ast_to_object(output)
   self.assertDictEqual(result.d, {'a': 3, 'c': 1, 'd': 'e'})
示例#15
0
    def test_rename_symbols_attributes(self):
        node = parser.parse_str('b.c = b.c.d')
        node = qual_names.resolve(node)

        node = ast_util.rename_symbols(
            node, {qual_names.from_str('b.c'): qual_names.QN('renamed_b_c')})

        source, _ = compiler.ast_to_source(node)
        self.assertEqual(source.strip(), 'renamed_b_c = renamed_b_c.d')
示例#16
0
 def test_keywords_to_dict(self):
     keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
     d = ast_util.keywords_to_dict(keywords)
     # Make sure we generate a usable dict node by attaching it to a variable and
     # compiling everything.
     node = parser.parse_str('def f(b): pass').body[0]
     node.body.append(ast.Return(d))
     result, _ = compiler.ast_to_object(node)
     self.assertDictEqual(result.f(3), {'a': 3, 'c': 1, 'd': 'e'})
示例#17
0
 def test_keywords_to_dict(self):
   keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
   d = ast_util.keywords_to_dict(keywords)
   # Make sure we generate a usable dict node by attaching it to a variable and
   # compiling everything.
   node = parser.parse_str('def f(b): pass').body[0]
   node.body.append(ast.Return(d))
   result, _ = compiler.ast_to_object(node)
   self.assertDictEqual(result.f(3), {'a': 3, 'c': 1, 'd': 'e'})
示例#18
0
    def test_rename_symbols_basic(self):
        node = parser.parse_str('a + b')
        node = qual_names.resolve(node)

        node = ast_util.rename_symbols(
            node, {qual_names.QN('a'): qual_names.QN('renamed_a')})

        self.assertIsInstance(node.body[0].value.left.id, str)
        self.assertEqual(compiler.ast_to_source(node).strip(), 'renamed_a + b')
示例#19
0
  def test_rename_symbols_attributes(self):
    node = parser.parse_str('b.c = b.c.d')
    node = qual_names.resolve(node)

    node = ast_util.rename_symbols(
        node, {qual_names.from_str('b.c'): qual_names.QN('renamed_b_c')})

    source = compiler.ast_to_source(node)
    self.assertEqual(source.strip(), 'renamed_b_c = renamed_b_c.d')
示例#20
0
 def test_keywords_to_dict(self):
     keywords = parser.parse_expression('f(a=b, c=1, d=\'e\')').keywords
     d = ast_util.keywords_to_dict(keywords)
     # Make sure we generate a usable dict node by attaching it to a variable and
     # compiling everything.
     output = parser.parse_str('b = 3')
     output.body += (ast.Assign([ast.Name(id='d', ctx=ast.Store())], d), )
     result, _ = compiler.ast_to_object(output)
     self.assertDictEqual(result.d, {'a': 3, 'c': 1, 'd': 'e'})
示例#21
0
    def test_to_code_basic(self):
        def test_fn(x, s):
            while tf.reduce_sum(x) > s:
                x /= 2
            return x

        compiled_code = api.to_code(test_fn)

        # Just check that it is parseable Python code.
        self.assertIsNotNone(parser.parse_str(compiled_code))
示例#22
0
  def test_rename_symbols_basic(self):
    node = parser.parse_str('a + b')
    node = qual_names.resolve(node)

    node = ast_util.rename_symbols(
        node, {qual_names.QN('a'): qual_names.QN('renamed_a')})

    self.assertIsInstance(node.body[0].value.left.id, str)
    source = compiler.ast_to_source(node)
    self.assertEqual(source.strip(), 'renamed_a + b')
示例#23
0
    def test_rename_symbols_annotations(self):
        node = parser.parse_str('a[i]')
        node = qual_names.resolve(node)
        anno.setanno(node, 'foo', 'bar')
        orig_anno = anno.getanno(node, 'foo')

        node = ast_util.rename_symbols(
            node, {qual_names.QN('a'): qual_names.QN('b')})

        self.assertIs(anno.getanno(node, 'foo'), orig_anno)
示例#24
0
  def test_rename_symbols_annotations(self):
    node = parser.parse_str('a[i]')
    node = qual_names.resolve(node)
    anno.setanno(node, 'foo', 'bar')
    orig_anno = anno.getanno(node, 'foo')

    node = ast_util.rename_symbols(node,
                                   {qual_names.QN('a'): qual_names.QN('b')})

    self.assertIs(anno.getanno(node, 'foo'), orig_anno)
示例#25
0
 def test_apply_to_single_assignments_static_unpack(self):
     node = parser.parse_str('a, b, c = d, e, f')
     node = node.body[0]
     ast_util.apply_to_single_assignments(node.targets, node.value,
                                          self._mock_apply_fn)
     self.assertDictEqual(self._invocation_counts, {
         ('a', 'd'): 1,
         ('b', 'e'): 1,
         ('c', 'f'): 1,
     })
  def test_source_map(self):

    def test_fn(x):
      if x > 0:
        x += 1
      return x

    node, source = parser.parse_entity(test_fn)
    fn_node = node.body[0]
    origin_info.resolve(fn_node, source)

    # Insert a traced line.
    new_node = parser.parse_str('x = abs(x)').body[0]
    anno.copyanno(fn_node.body[0], new_node, anno.Basic.ORIGIN)
    fn_node.body.insert(0, new_node)

    # Insert an untraced line.
    fn_node.body.insert(0, parser.parse_str('x = 0').body[0])

    modified_source = compiler.ast_to_source(fn_node)

    source_map = origin_info.source_map(fn_node, modified_source,
                                        'test_filename', [0])

    loc = origin_info.LineLocation('test_filename', 1)
    origin = source_map[loc]
    self.assertEqual(origin.source_code_line, 'def test_fn(x):')
    self.assertEqual(origin.loc.lineno, 1)

    # The untraced line, inserted second.
    loc = origin_info.LineLocation('test_filename', 2)
    self.assertFalse(loc in source_map)

    # The traced line, inserted first.
    loc = origin_info.LineLocation('test_filename', 3)
    origin = source_map[loc]
    self.assertEqual(origin.source_code_line, '  if x > 0:')
    self.assertEqual(origin.loc.lineno, 2)

    loc = origin_info.LineLocation('test_filename', 4)
    origin = source_map[loc]
    self.assertEqual(origin.source_code_line, '  if x > 0:')
    self.assertEqual(origin.loc.lineno, 2)
示例#27
0
 def test_apply_to_single_assignments_static_unpack(self):
   node = parser.parse_str('a, b, c = d, e, f')
   node = node.body[0]
   ast_util.apply_to_single_assignments(node.targets, node.value,
                                        self._mock_apply_fn)
   self.assertDictEqual(self._invocation_counts, {
       ('a', 'd'): 1,
       ('b', 'e'): 1,
       ('c', 'f'): 1,
   })
示例#28
0
 def test_copy_clean_preserves_annotations(self):
     node = parser.parse_str(
         textwrap.dedent("""
   def f(a):
     return a + 1
 """))
     anno.setanno(node.body[0], 'foo', 'bar')
     anno.setanno(node.body[0], 'baz', 1)
     new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
     self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
     self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
示例#29
0
 def test_copy_clean(self):
     node = parser.parse_str(
         textwrap.dedent("""
   def f(a):
     return a + 1
 """))
     setattr(node.body[0], '__foo', 'bar')
     new_node = ast_util.copy_clean(node)
     self.assertIsNot(new_node, node)
     self.assertIsNot(new_node.body[0], node.body[0])
     self.assertFalse(hasattr(new_node.body[0], '__foo'))
示例#30
0
  def test_to_code_basic(self):

    def test_fn(x, s):
      while tf.reduce_sum(x) > s:
        x /= 2
      return x

    compiled_code = api.to_code(test_fn)

    # Just check that it is parseable Python code.
    self.assertIsNotNone(parser.parse_str(compiled_code))
示例#31
0
 def test_copy_clean_preserves_annotations(self):
   node = parser.parse_str(
       textwrap.dedent("""
     def f(a):
       return a + 1
   """))
   anno.setanno(node.body[0], 'foo', 'bar')
   anno.setanno(node.body[0], 'baz', 1)
   new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
   self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
   self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
示例#32
0
    def test_to_code_basic(self):
        def test_fn(x, s):
            while tf.reduce_sum(x) > s:
                x /= 2
            return x

        compiled_code = api.to_code(test_fn)

        # Just check for some key words and that it is parseable Python code.
        self.assertRegexpMatches(compiled_code, 'autograph_utils\\.run_while')
        self.assertIsNotNone(parser.parse_str(compiled_code))
示例#33
0
 def test_copy_clean(self):
   node = parser.parse_str(
       textwrap.dedent("""
     def f(a):
       return a + 1
   """))
   setattr(node.body[0], '__foo', 'bar')
   new_node = ast_util.copy_clean(node)
   self.assertIsNot(new_node, node)
   self.assertIsNot(new_node.body[0], node.body[0])
   self.assertFalse(hasattr(new_node.body[0], '__foo'))
def source_map(nodes, code, filename, indices_in_code):
    """Creates a source map between an annotated AST and the code it compiles to.

  Args:
    nodes: Iterable[ast.AST, ...]
    code: Text
    filename: Optional[Text]
    indices_in_code: Union[int, Iterable[int, ...]], the positions at which
        nodes appear in code. The parser always returns a module when parsing
        code. This argument indicates the position in that module's body at
        which the corresponding of node should appear.

  Returns:
    Dict[CodeLocation, OriginInfo], mapping locations in code to locations
    indicated by origin annotations in node.
  """
    reparsed_nodes = parser.parse_str(code)
    reparsed_nodes = [reparsed_nodes.body[i] for i in indices_in_code]

    resolve(reparsed_nodes, code)
    result = {}

    for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
        # Note: generated code might not be mapped back to its origin.
        # TODO(mdan): Generated code should always be mapped to something.
        origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
        final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
        if origin_info is None or final_info is None:
            continue

        line_loc = LineLocation(filename, final_info.loc.lineno)

        existing_origin = result.get(line_loc)
        if existing_origin is not None:
            # Overlaps may exist because of child nodes, but almost never to
            # different line locations. Exception make decorated functions, where
            # both lines are mapped to the same line in the AST.

            # Line overlaps: keep bottom node.
            if existing_origin.loc.line_loc == origin_info.loc.line_loc:
                if existing_origin.loc.lineno >= origin_info.loc.lineno:
                    continue

            # In case of overlaps, keep the leftmost node.
            if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
                continue

        result[line_loc] = origin_info

    return result
示例#35
0
def source_map(nodes, code, filename, indices_in_code):
  """Creates a source map between an annotated AST and the code it compiles to.

  Args:
    nodes: Iterable[ast.AST, ...]
    code: Text
    filename: Optional[Text]
    indices_in_code: Union[int, Iterable[int, ...]], the positions at which
        nodes appear in code. The parser always returns a module when parsing
        code. This argument indicates the position in that module's body at
        which the corresponding of node should appear.

  Returns:
    Dict[CodeLocation, OriginInfo], mapping locations in code to locations
    indicated by origin annotations in node.
  """
  reparsed_nodes = parser.parse_str(code)
  reparsed_nodes = [reparsed_nodes.body[i] for i in indices_in_code]

  resolve(reparsed_nodes, code)
  result = {}

  for before, after in ast_util.parallel_walk(nodes, reparsed_nodes):
    # Note: generated code might not be mapped back to its origin.
    # TODO(mdan): Generated code should always be mapped to something.
    origin_info = anno.getanno(before, anno.Basic.ORIGIN, default=None)
    final_info = anno.getanno(after, anno.Basic.ORIGIN, default=None)
    if origin_info is None or final_info is None:
      continue

    line_loc = LineLocation(filename, final_info.loc.lineno)

    existing_origin = result.get(line_loc)
    if existing_origin is not None:
      # Overlaps may exist because of child nodes, but almost never to
      # different line locations. Exception make decorated functions, where
      # both lines are mapped to the same line in the AST.

      # Line overlaps: keep bottom node.
      if existing_origin.loc.line_loc == origin_info.loc.line_loc:
        if existing_origin.loc.lineno >= origin_info.loc.lineno:
          continue

      # In case of overlaps, keep the leftmost node.
      if existing_origin.loc.col_offset <= origin_info.loc.col_offset:
        continue

    result[line_loc] = origin_info

  return result
 def test_function_calls(self):
     samples = """
   a.b
   a.b()
   a().b
   z[i]
   z[i]()
   z()[i]
 """
     nodes = resolve(parser.parse_str(textwrap.dedent(samples)))
     nodes = tuple(n.value for n in nodes.body)
     self.assertQNStringIs(nodes[0], 'a.b')
     self.assertQNStringIs(nodes[1].func, 'a.b')
     self.assertQNStringIs(nodes[2].value.func, 'a')
     self.assertQNStringIs(nodes[3], 'z[i]')
     self.assertQNStringIs(nodes[4].func, 'z[i]')
     self.assertQNStringIs(nodes[5].value.func, 'z')
    def test_resolve(self):
        samples = """
      a
      a.b
      (c, d.e)
      [f, (g.h.i)]
      j(k, l)
    """
        nodes = resolve(parser.parse_str(textwrap.dedent(samples)))
        nodes = tuple(n.value for n in nodes.body)

        self.assertQNStringIs(nodes[0], 'a')
        self.assertQNStringIs(nodes[1], 'a.b')
        self.assertQNStringIs(nodes[2].elts[0], 'c')
        self.assertQNStringIs(nodes[2].elts[1], 'd.e')
        self.assertQNStringIs(nodes[3].elts[0], 'f')
        self.assertQNStringIs(nodes[3].elts[1], 'g.h.i')
        self.assertQNStringIs(nodes[4].func, 'j')
        self.assertQNStringIs(nodes[4].args[0], 'k')
        self.assertQNStringIs(nodes[4].args[1], 'l')
示例#38
0
 def instantiate_list_node(self):
   return parser.parse_str('[]').body[0].value
示例#39
0
 def instantiate_list_node(self):
     return parser.parse_str('[]').body[0].value