Python templates.replace_as_expression函数代码示例

 def test_replace_as_expression_restrictions(self):
   template = """
     foo(a)
     bar(b)
   """
   with self.assertRaises(ValueError):
     templates.replace_as_expression(template)

  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 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_as_expression(
          'func_name', func_name=new_name)
    return node

  def visit_Call(self, node):
    # TODO(mdan): Refactor converted_call as a 'Call' operator.

    # Calls to the internal 'ag__' module are never converted (though their
    # arguments might be).
    full_name = str(anno.getanno(node.func, anno.Basic.QN, default=''))
    if full_name.startswith('ag__.'):
      return self.generic_visit(node)
    if (full_name == 'print' and
        not self.ctx.program.options.uses(converter.Feature.BUILTIN_FUNCTIONS)):
      return self.generic_visit(node)

    template = """
      ag__.converted_call(func, owner, options, args)
    """
    if isinstance(node.func, gast.Attribute):
      func = gast.Str(node.func.attr)
      owner = node.func.value
    else:
      func = node.func
      owner = parser.parse_expression('None')

    new_call = templates.replace_as_expression(
        template,
        func=func,
        owner=owner,
        options=self.ctx.program.options.to_ast(
            self.ctx,
            internal_convert_user_code=self.ctx.program.options.recursive),
        args=node.args)
    # TODO(mdan): Improve the template mechanism to better support this.
    new_call.keywords = node.keywords

    return new_call

 def _as_function(self, func_name, args):
   template = """
     func_name(args)
   """
   replacement = templates.replace_as_expression(
       template, func_name=parser.parse_expression(func_name), args=args)
   anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True)
   return replacement

  def test_replace_as_expression(self):
    template = """
      foo(a)
    """

    node = templates.replace_as_expression(template, foo='bar', a='baz')
    self.assertIsInstance(node, gast.Call)
    self.assertEqual(node.func.id, 'bar')
    self.assertEqual(node.args[0].id, 'baz')

 def _convert_builtin(self, f, args, as_expression):
   template = """
     ag__.func(args)
   """
   if as_expression:
     return templates.replace_as_expression(
         template, func=py_builtins.overload_of(f).__name__, args=args)
   else:
     return templates.replace(
         template, func=py_builtins.overload_of(f).__name__, args=args)

 def _replace_stack_call(self, node):
   assert len(node.args) == 1
   dtype = self.get_definition_directive(
       node.args[0],
       directives.set_element_type,
       'dtype',
       default=templates.replace_as_expression('None'))
   template = """
     ag__.list_stack(
         target,
         opts=ag__.ListStackOpts(
             element_dtype=dtype,
             original_call=orig_call))
   """
   return templates.replace_as_expression(
       template,
       dtype=dtype,
       target=node.args[0],
       orig_call=node.func)

 def _wrap_to_py_func_single_return(self, node, dtype):
   # TODO(mdan): Properly handle varargs, etc.
   template = """
     ag__.utils.wrap_py_func(func, dtype, (args,), kwargs, False)
   """
   return templates.replace_as_expression(
       template,
       func=node.func,
       dtype=parser.parse_expression(dtype),
       args=node.args,
       kwargs=ast_util.keywords_to_dict(node.keywords))

  def visit_For(self, node):
    node.iter = self.visit(node.iter)
    node.target = self.visit(node.target)

    # Add the check for return to the loop condition.
    node.body = self._visit_statement_block(node, node.body)
    if self.state[_Return].used:
      extra_test = anno.getanno(node, 'extra_test', default=None)
      if extra_test is not None:
        extra_test = templates.replace_as_expression(
            'ag__.and_(lambda: ag__.not_(control_var), lambda: extra_test)',
            extra_test=extra_test,
            control_var=self.state[_Function].do_return_var_name)
      else:
        extra_test = templates.replace_as_expression(
            'ag__.not_(control_var)',
            control_var=self.state[_Function].do_return_var_name)
      anno.setanno(node, 'extra_test', extra_test)

    node.orelse = self._visit_statement_block(node, node.orelse)
    return node

  def visit_While(self, node):
    node.test = self.visit(node.test)

    # Add the check for return to the loop condition.
    node.body = self._visit_statement_block(node, node.body)
    if self.state[_Return].used:
      node.test = templates.replace_as_expression(
          'ag__.and_(lambda: ag__.not_(control_var), lambda: test)',
          test=node.test,
          control_var=self.state[_Function].do_return_var_name)

    node.orelse = self._visit_statement_block(node, node.orelse)
    return node

  def _as_function(self, func_name, args, args_as_lambda=False):
    if args_as_lambda:
      args_as_lambda = []
      for arg in args:
        template = """
          lambda: arg
        """
        args_as_lambda.append(
            templates.replace_as_expression(template, arg=arg))
      args = args_as_lambda

    if not args:
      template = """
        func_name()
      """
      replacement = templates.replace_as_expression(
          template, func_name=parser.parse_expression(func_name))
    elif len(args) == 1:
      template = """
        func_name(arg)
      """
      replacement = templates.replace_as_expression(
          template, func_name=parser.parse_expression(func_name), arg=args[0])
    elif len(args) == 2:
      template = """
        func_name(arg1, arg2)
      """
      replacement = templates.replace_as_expression(
          template,
          func_name=parser.parse_expression(func_name),
          arg1=args[0],
          arg2=args[1])
    else:
      raise NotImplementedError('{} arguments for {}'.format(
          len(args), func_name))

    anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True)
    return replacement

  def visit_IfExp(self, node):
    if anno.hasanno(node.test, anno.Basic.QN):
      name_root = anno.getanno(node.test, anno.Basic.QN).ssf()
    else:
      name_root = 'ifexp'

    true_fn_name = self._create_branch(node.body, '%s_true' % name_root)
    false_fn_name = self._create_branch(node.orelse, '%s_false' % name_root)

    return templates.replace_as_expression(
        'ag__.utils.run_cond(test, true_fn_name, false_fn_name)',
        test=node.test,
        true_fn_name=true_fn_name,
        false_fn_name=false_fn_name)

  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.
    # TODO(mdan): For lists of lists, this won't work.
    # The reason why it won't work is because it's unclear how to annotate
    # the list as a "list of lists with a certain element type" when using
    # operations like `l.pop().pop()`.
    dtype = self.get_definition_directive(
        original_call_node.func.value,
        directives.set_element_type,
        'dtype',
        default=templates.replace_as_expression('None'))
    shape = self.get_definition_directive(
        original_call_node.func.value,
        directives.set_element_type,
        '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 visit_Print(self, node):
    node = self.generic_visit(node)
    args = node.values
    # Following is the case when calling print(a, b)
    if len(args) == 1 and isinstance(args[0], gast.Tuple):
      args = args[0].elts

    template = """
      ag__.converted_call(func, None, options, args, {})
    """
    return templates.replace_as_expression(
        template,
        func='print',
        options=self.ctx.program.options.to_ast(),
        args=args)

  def visit_Subscript(self, node):
    node = self.generic_visit(node)
    if not isinstance(node.slice, gast.Index):
      return node

    if not isinstance(node.ctx, gast.Load):
      # Index writes are handled at a higher level, one at which the rvalue is
      # also available.
      return node

    dtype = self.get_definition_directive(
        node.value,
        directives.set_element_type,
        'dtype',
        default=templates.replace_as_expression('None'))

    template = """
      ag__.get_item(
          target,
          key,
          opts=ag__.GetItemOpts(element_dtype=dtype))
    """
    return templates.replace_as_expression(
        template, target=node.value, key=node.slice.value, dtype=dtype)