Python anno.getanno函数代码示例

  def _process_variable_assignment(self, source, targets):
    # Special case: constructors.
    if isinstance(source, gast.Call):
      func = source.func
      if anno.hasanno(func, 'live_val'):
        func_obj = anno.getanno(func, 'live_val')
        if tf_inspect.isclass(func_obj):
          anno.setanno(source, 'is_constructor', True)
          anno.setanno(source, 'type', func_obj)
          anno.setanno(source, 'type_fqn', anno.getanno(func, 'fqn'))
          # TODO(mdan): Raise an error if constructor has side effects.
          # We can have a whitelist of no-side-effects constructors.
          # We can also step inside the constructor and further analyze.

    # Multiple targets mean multiple assignment.
    for target in targets:
      # Tuple target means unpacking.
      if isinstance(target, gast.Tuple):
        for i, target_item in enumerate(target.elts):
          # Two cases here:
          #   1. Static unpacking, e.g. a, b = c, d
          #   2. Dynamic unpacking, e.g. a, b = c
          # The former case is optimized away.
          if isinstance(source, (gast.Tuple, gast.List)):
            source_item = source.elts[i]
          else:
            source_item = gast.Subscript(source, gast.Index(i), ctx=None)
          self._process_variable_assignment(source_item, (target_item,))
      elif isinstance(target, (gast.Name, gast.Attribute)):
        target_symbol = anno.getanno(target, anno.Basic.QN)
        self.scope.setval(target_symbol, source)
      else:
        raise ValueError(
            'assignment target has unknown type: %s' % target_item)

  def test_if_attributes(self):

    def test_fn(a):
      if a > 0:
        a.b = -a.c
        d = 2 * a
      else:
        a.b = a.c
        d = 1
      return d

    node, _ = self._parse_and_analyze(test_fn)
    if_node = node.body[0].body[0]
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.BODY_SCOPE),
        ('a', 'a.c'),
        ('a.b', 'd'),
        ('d',),
    )
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.ORELSE_SCOPE),
        ('a', 'a.c'),
        ('a.b', 'd'),
        ('d',),
    )
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent,
        ('a', 'a.c', 'd'),
        ('a.b', 'd'),
        ('a', 'd'),
    )

  def _rename_compilable_function(self, node):
    assert anno.hasanno(node.func, 'live_val')
    assert anno.hasanno(node.func, 'fqn')
    target_entity = anno.getanno(node.func, 'live_val')
    target_fqn = anno.getanno(node.func, 'fqn')

    if not self._should_compile(node, target_fqn):
      return node

    if anno.hasanno(node, 'is_constructor'):
      new_name = self.ctx.namer.compiled_class_name(
          target_fqn, live_entity=target_entity)
      do_rename = True
    else:
      if anno.hasanno(node.func, 'parent_type'):
        owner_type = anno.getanno(node.func, 'parent_type')
      else:
        # Fallback - not reliable.
        owner_type = inspect_utils.getmethodclass(target_entity)
      new_name, do_rename = self.ctx.namer.compiled_function_name(
          target_fqn, live_entity=target_entity, owner_type=owner_type)

    if do_rename:
      if target_entity is not None:
        if tf_inspect.ismethod(target_entity):
          # The renaming process will transform it into a regular function.
          # TODO(mdan): Is this complete? How does it work with nested members?
          node.args = [node.func.value] + node.args
      node.func = templates.replace('func_name', func_name=new_name)[0]
    return node

 def visit_For(self, node):
   node.target = self.visit(node.target)
   node.body = self._process_block(
       anno.getanno(node, NodeAnno.BODY_SCOPE), node.body)
   node.orelse = self._process_block(
       anno.getanno(node, NodeAnno.ORELSE_SCOPE), node.orelse)
   return node

 def visit_Attribute(self, node):
   self.generic_visit(node)
   if anno.hasanno(node.value, 'live_val'):
     assert anno.hasanno(node.value, 'fqn')
     parent_object = anno.getanno(node.value, 'live_val')
     if not hasattr(parent_object, node.attr):
       raise AttributeError('%s has no attribute %s' % (parent_object,
                                                        node.attr))
     anno.setanno(node, 'parent_type', type(parent_object))
     anno.setanno(node, 'live_val', getattr(parent_object, node.attr))
     anno.setanno(node, 'fqn', anno.getanno(node.value, 'fqn') + (node.attr,))
   # TODO(mdan): Investigate the role built-in annotations can play here.
   elif anno.hasanno(node.value, 'type'):
     parent_type = anno.getanno(node.value, 'type')
     if hasattr(parent_type, node.attr):
       # This should hold for static members like methods.
       # This would not hold for dynamic members like function attributes.
       # For the dynamic case, we simply leave the node without an annotation,
       # and let downstream consumers figure out what to do.
       anno.setanno(node, 'parent_type', parent_type)
       anno.setanno(node, 'live_val', getattr(parent_type, node.attr))
       anno.setanno(node, 'fqn',
                    anno.getanno(node.value, 'type_fqn') + (node.attr,))
   elif isinstance(node.value, gast.Name):
     stem_name = node.value
     # All nonlocal symbols should be fully resolved.
     assert anno.hasanno(stem_name, NodeAnno.IS_LOCAL), stem_name
     # TODO(mdan): Figure out what to do when calling attribute on local object
     # Maybe just leave as-is?
   return node

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

  def _generate_pop_operation(self, original_call_node, pop_var_name):
    assert isinstance(original_call_node.func, gast.Attribute)

    if original_call_node.args:
      pop_element = original_call_node.args[0]
    else:
      pop_element = parser.parse_expression('None')
    # The call will be something like "target.pop()", and the dtype is hooked to
    # target, hence the func.value.
    dtype = anno.getanno(
        original_call_node.func.value,
        'element_type',
        default=templates.replace_as_expression('None'))
    shape = anno.getanno(
        original_call_node.func.value,
        'element_shape',
        default=templates.replace_as_expression('None'))

    template = """
      target, pop_var_name = ag__.list_pop(
          target, element,
          opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
    """
    return templates.replace(
        template,
        target=original_call_node.func.value,
        pop_var_name=pop_var_name,
        element=pop_element,
        dtype=dtype,
        shape=shape)

  def test_if_subscripts(self):

    def test_fn(a, b, c, e):
      if a > 0:
        a[b] = -a[c]
        d = 2 * a
      else:
        a[0] = e
        d = 1
      return d

    node, _ = self._parse_and_analyze(test_fn)
    if_node = node.body[0].body[0]
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.BODY_SCOPE),
        ('a', 'b', 'c', 'a[c]'),
        ('a', 'a[b]', 'd'),
        ('d',),
    )
    # TODO(mdan): Should subscript writes (a[0] = 1) be considered to read "a"?
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.ORELSE_SCOPE),
        ('a', 'e'),
        ('a', 'a[0]', 'd'),
        ('d',),
    )
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent,
        ('a', 'b', 'c', 'd', 'e', 'a[c]'),
        ('a', 'd', 'a[b]', 'a[0]'),
        ('a', 'b', 'c', 'd', 'e'),
    )

  def visit_Call(self, node):
    if anno.hasanno(node.func, 'live_val'):
      # Symbols targeted by the "set_type" marker function are assigned the data
      # type that it specified.
      if anno.getanno(node.func, 'live_val') is utils.set_element_type:

        if len(node.args) < 2 or len(node.args) > 3:
          raise ValueError('"%s" must have either two or three parameters'
                           % self.context.type_annotation_func)
        if len(node.args) == 2:
          target_arg, type_arg = node.args
          shape_arg = parser.parse_expression('None')
        else:
          target_arg, type_arg, shape_arg = node.args
        if not anno.hasanno(target_arg, anno.Basic.QN):
          raise ValueError('the first argument of "%s" must by a symbol' %
                           utils.set_element_type)
        # TODO(mdan): This is vulnerable to symbol renaming.
        element_type = type_arg
        element_shape = shape_arg

        target_symbol = anno.getanno(target_arg, anno.Basic.QN)
        # Find the definition of this symbol and annotate it with the given
        # data type. That in turn will cause future uses of the symbol
        # to receive the same type annotation.
        definition = self.scope.getval(target_symbol)
        anno.setanno(node, 'element_type', element_type)
        anno.setanno(node, 'element_shape', element_shape)
        anno.setanno(definition, 'element_type', element_type)
        anno.setanno(definition, 'element_shape', element_shape)
        # TODO(mdan): Should we update references between definition and here?
    return self.generic_visit(node)

 def visit_Call(self, node):
   if anno.hasanno(node.func, 'live_val'):
     # Symbols targeted by the "set_type" marker function are assigned the data
     # type that it specified.
     if (anno.getanno(node.func, 'live_val') is
         self.context.type_annotation_func):
       # Expecting the actual type to be the second argument.
       if len(node.args) != 2:
         raise ValueError('"%s" must have exactly two parameters'
                          % self.context.type_annotation_func)
       if not anno.hasanno(node.args[0], anno.Basic.QN):
         raise ValueError('the first argument of "%s" must by a symbol'
                          % self.context.type_annotation_func)
       if not anno.hasanno(node.args[1], 'live_val'):
         raise ValueError(
             'the second argument of "%s" must be statically resolvable' %
             self.context.type_annotation_func)
       target_symbol = anno.getanno(node.args[0], anno.Basic.QN)
       element_type = anno.getanno(node.args[1], 'live_val')
       # Find the definition of this symbol and annotate it with the given
       # data type. That in turn will cause future uses of the symbol
       # to receive the same type annotation.
       definition = self.scope.getval(target_symbol)
       anno.setanno(node, 'element_type', element_type)
       anno.setanno(definition, 'element_type', element_type)
       # TODO(mdan): Should we update references between definition and here?
   return self.generic_visit(node)

  def visit(self, node):
    """Depth-first walking the CFG, applying dataflow info propagation."""
    # node.value is None only for the exit CfgNode.
    if not node.value:
      return

    if anno.hasanno(node.value, self.out_label):
      before = hash(anno.getanno(node.value, self.out_label))
    else:
      before = None
    preds = [
        anno.getanno(pred.value, self.out_label)
        for pred in node.prev
        if anno.hasanno(pred.value, self.out_label)
    ]
    if preds:
      incoming = functools.reduce(self.transfer_fn, preds[1:], preds[0])
    else:
      incoming = frozenset()
    anno.setanno(node.value, self.in_label, incoming)
    gen, kill = self.get_gen_kill(node, incoming)
    anno.setanno(node.value, self.gen_label, gen)
    anno.setanno(node.value, self.kill_label, kill)
    anno.setanno(node.value, self.out_label, (incoming - kill) | gen)

    if hash(anno.getanno(node.value, self.out_label)) != before:
      for succ in node.next:
        self.visit(succ)

  def test_call_with_composite_names(self):

    def foo(*_):
      pass

    def test_fn(a):
      foo(a.b, a.c)
      if a > 0:
        a.b = 2
      else:
        d = 2
        d.e = a.c
        f = d.e + 1
        a.c = f

    node = self._parse_and_analyze(test_fn)
    call_node = node.body[0].body[0].value
    self.assertScopeIsRmc(
        anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'a.b', 'a.c'), (),
        ())
    if_node = node.body[0].body[1]
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('a',), ('a.b',), ())
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.ORELSE_SCOPE),
        ('a', 'a.c', 'd', 'd.e', 'f'), ('a.c', 'd', 'd.e', 'f'), ('d', 'f'))

  def test_nested_function(self):

    def test_fn(a):

      def f(x):
        y = x * x
        return y

      b = a
      for i in a:
        c = b
        b -= f(i)
      return b, c

    node, _ = self._parse_and_analyze(test_fn)
    fn_def_node = node.body[0].body[0]

    self.assertScopeIsRmc(
        anno.getanno(fn_def_node,
                     NodeAnno.BODY_SCOPE).parent, ('b', 'i', 'f', 'c', 'a'),
        ('f', 'b', 'c', 'i'), ('f', 'a', 'b', 'c', 'i'))
    self.assertScopeIsRmc(
        anno.getanno(fn_def_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('y',), (
            'x',
            'y',
        ))

  def test_if(self):

    def test_fn(x):
      if x > 0:
        x = -x
        y = 2 * x
        z = -y
      else:
        x = 2 * x
        y = -x
        u = -y
      return z, u

    node, _ = self._parse_and_analyze(test_fn)
    if_node = node.body[0].body[0]
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('x', 'y', 'z'),
        ('y', 'z'))
    # TODO(mdan): Double check: is it ok to not mark a local symbol as not read?
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent, ('x', 'z', 'u'),
        ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u'))
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), ('x', 'y'),
        ('x', 'y', 'u'), ('y', 'u'))
    self.assertScopeIsRmc(
        anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent, ('x', 'z', 'u'),
        ('x', 'y', 'z', 'u'), ('x', 'y', 'z', 'u'))

  def test_resolve(self):

    def test_fn(x):
      """Docstring."""
      return x  # comment

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

    origin = anno.getanno(fn_node, anno.Basic.ORIGIN)
    self.assertEqual(origin.loc.lineno, 1)
    self.assertEqual(origin.loc.col_offset, 0)
    self.assertEqual(origin.source_code_line, 'def test_fn(x):')
    self.assertIsNone(origin.comment)

    origin = anno.getanno(fn_node.body[0], anno.Basic.ORIGIN)
    self.assertEqual(origin.loc.lineno, 2)
    self.assertEqual(origin.loc.col_offset, 2)
    self.assertEqual(origin.source_code_line, '  """Docstring."""')
    self.assertIsNone(origin.comment)

    origin = anno.getanno(fn_node.body[1], anno.Basic.ORIGIN)
    self.assertEqual(origin.loc.lineno, 3)
    self.assertEqual(origin.loc.col_offset, 2)
    self.assertEqual(origin.source_code_line, '  return x  # comment')
    self.assertEqual(origin.comment, 'comment')