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

  def test_replace_attribute(self):
    template = """
      def test_fn(a):
        return a.foo
    """

    node = templates.replace(template, foo='b')[0]
    result, _ = compiler.ast_to_object(node)
    mod = imp.new_module('test')
    mod.b = 3
    self.assertEquals(3, result.test_fn(mod))

    with self.assertRaises(ValueError):
      templates.replace(template, foo=1)

  def visit_For(self, node):
    self.generic_visit(node)
    body_scope = anno.getanno(node, 'body_scope')

    # TODO(mdan): Distinguish between `for i in n` and `for i in range(n)`
    # Or maybe we should replace range with tf.range?

    if anno.hasanno(node, 'extra_cond'):

      def template(loop_iter, target, body, i, n, extra_cond):  # pylint:disable=unused-argument
        i = 0
        n = len(loop_iter)  # pylint:disable=undefined-variable
        while i < n and extra_cond:
          # TODO(mdan): Use TensorListFromTensor(loop_iter) here.
          target = loop_iter[i]
          body  # pylint:disable=pointless-statement
          i += 1

      return templates.replace(
          template,
          loop_iter=node.iter,
          target=node.target,
          body=node.body,
          i=gast.Name(
              self.namer.new_symbol('i', body_scope.referenced), None, None),
          n=gast.Name(
              self.namer.new_symbol('n', body_scope.referenced), None, None),
          extra_cond=anno.getanno(node, 'extra_cond'))
    else:

      def template(loop_iter, target, body, i, n):  # pylint:disable=unused-argument
        i = 0
        n = len(loop_iter)  # pylint:disable=undefined-variable
        while i < n:
          # TODO(mdan): Use TensorListFromTensor(loop_iter) here.
          target = loop_iter[i]
          body  # pylint:disable=pointless-statement
          i += 1

      return templates.replace(
          template,
          loop_iter=node.iter,
          target=node.target,
          body=node.body,
          i=gast.Name(
              self.namer.new_symbol('i', body_scope.referenced), None, None),
          n=gast.Name(
              self.namer.new_symbol('n', body_scope.referenced), None, None))

 def _gate_symbols(self, guard_statement, guarded_args):
   template = """
     (args,) = (tf.identity(a) for a in (args,))
   """
   guards = templates.replace(template, args=tuple(guarded_args))
   guard_statement.body.extend(guards)
   return guard_statement

  def _convert_len(self, node):

    def template(args):
      tf.shape(args)[0]  # pylint:disable=undefined-variable,expression-not-assigned

    new_call = templates.replace(template, args=node.args)[0].value
    return new_call

  def _wrap_to_py_func_no_return(self, node):
    func_qn = anno.getanno(node.func, anno.Basic.QN)
    args_scope = anno.getanno(node, NodeAnno.ARGS_SCOPE)
    wrapper_name = self.context.namer.new_symbol(func_qn.ssf(),
                                                 args_scope.referenced)
    wrapper_args = []
    for arg in node.args:
      if anno.hasanno(arg, anno.Basic.QN):
        arg_qn = anno.getanno(arg, anno.Basic.QN)
      else:
        arg_qn = qual_names.QN('arg')
      wrapper_args.append(
          self.context.namer.new_symbol(arg_qn.ssf(), args_scope.referenced))
    # TODO(mdan): Properly handle varargs, kwargs, etc.
    # TODO(mdan): This is best handled as a dynamic dispatch.
    # That way we can separate tensors from non-tensor args.
    template = """
      def wrapper(wrapper_args):
        call(wrapper_args)
        return 1
      tf.py_func(wrapper, original_args, [tf.int64])
    """
    wrapper_def, call_expr = templates.replace(
        template,
        call=node.func,
        wrapper=wrapper_name,
        original_args=gast.List(elts=node.args, ctx=None),
        wrapper_args=wrapper_args)
    anno.setanno(wrapper_def, anno.Basic.SKIP_PROCESSING, True)

    return (wrapper_def, call_expr)

 def _insert_dynamic_conversion(self, node):
   """Inlines a dynamic conversion for a dynamic function."""
   # TODO(mdan): Pass information on the statically compiled functions.
   # Having access to the statically compiled functions can help avoid
   # unnecessary compilation.
   # For example, this would lead to function `a` being compiled twice:
   #
   #   def a():
   #     v = b
   #     b()
   #   def b():
   #     a()
   #
   # This is really a problem with recursive calls, which currently can
   # only be gated by a static condition, and should be rare.
   # TODO(mdan): It probably makes sense to use dynamic conversion every time.
   # Before we could convert all the time though, we'd need a reasonable
   # caching mechanism.
   template = """
     py2tf_api.converted_call(func, True, False, {}, original_args)
   """
   call_expr = templates.replace(
       template, func=node.func, original_args=node.args)
   new_call = call_expr[0].value
   # TODO(mdan): Improve the template mechanism to better support this.
   new_call.keywords = node.keywords
   return new_call

 def _create_break_trigger(self):
   template = """
     var_name = True
   """
   block = templates.replace(template, var_name=self.break_uses[-1][1])
   block.append(gast.Continue())
   return block

  def visit_While(self, node):
    self.generic_visit(node)

    body_scope = anno.getanno(node, 'body_scope')
    body_closure = tuple(body_scope.modified - body_scope.created)

    if len(body_closure) == 1:
      state = body_closure[0]
      state_ast_tuple = state
    else:
      state = tuple(body_closure)
      state_ast_tuple = gast.Tuple(
          tuple(gast.Name(n, None, None) for n in state), None)
    template = """
      def test_name(state):
        return test
      def body_name(state):
        body
        return state,
      state_ast_tuple = tf.while_loop(test_name, body_name, [state])
    """
    node = templates.replace(
        template,
        state=state,
        state_ast_tuple=state_ast_tuple,
        test_name=self.namer.new_symbol('loop_test', body_scope.referenced),
        test=node.test,
        body_name=self.namer.new_symbol('loop_body', body_scope.referenced),
        body=node.body)

    return node

  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.context.namer.compiled_class_name(
          target_fqn, live_entity=target_entity)
      do_rename = True
    else:
      owner_type = self._determine_function_owner(target_entity)
      new_name, do_rename = self.context.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 _create_continuation_trigger(self):
   template = """
     var_name = True
   """
   assign, = templates.replace(
       template, var_name=self.continuation_uses[-1][1])
   return assign

  def visit_Assert(self, node):
    self.generic_visit(node)

    # Note: The lone tf.Assert call will be wrapped with control_dependencies
    # by side_effect_guards.
    template = """
      tf.Assert(test, [tf.constant(msg)])
    """

    if node.msg is None:
      return templates.replace(
          template, test=node.test, msg=gast.Str('Assertion error'))
    elif isinstance(node.msg, gast.Str):
      return templates.replace(template, test=node.test, msg=node.msg)
    else:
      raise NotImplementedError('Can only convert string messages for now.')

  def _wrap_to_py_func_no_return(self, node):
    args_scope = anno.getanno(node, 'args_scope')
    # TODO(mdan): Properly handle varargs, kwargs, etc.
    args = tuple(gast.Name(n, gast.Load(), None) for n in args_scope.used)

    # pylint:disable=undefined-variable,unused-argument,function-redefined

    def template(call, wrapper, args):

      def wrapper(args):
        call(args)
        return 1

      tf.py_func(wrapper, [args], [tf.int64])

    # pylint:enable=undefined-variable,unused-argument,function-redefined

    wrapper_name = self.namer.compiled_function_name(node.func.id)
    wrapper_def, call_expr = templates.replace(
        template,
        call=node.func,
        wrapper=gast.Name(wrapper_name, gast.Load(), None),
        args=args)
    anno.setanno(call_expr.value, 'args_scope', args_scope)
    anno.setanno(wrapper_def, 'skip_processing', True)

    return (wrapper_def, call_expr)

 def _create_continuation_init(self):
   template = """
     var_name = False
   """
   assign, = templates.replace(
       template, var_name=self.continuation_uses[-1][1])
   return assign

  def test_replace_call_keyword(self):
    template = """
      def test_fn():
        def f(a, d, f):
          return a + d + f
        return f(1, kws=None)
    """

    source = parser.parse_expression('f(d=3, f=5)')
    node = templates.replace(template, kws=source.keywords)[0]
    result, _ = compiler.ast_to_object(node)
    self.assertEquals(9, result.test_fn())

    with self.assertRaises(ValueError):
      templates.replace(template, kws=[])
      templates.replace(template, kws=1)

  def visit_Expr(self, node):
    self.generic_visit(node)
    if isinstance(node.value, gast.Call):
      # Patterns of single function calls, like:
      #   opt.minimize(loss)
      # or:
      #   tf.py_func(...)
      template = """
        with py2tf_utils.control_dependency_on_returns(tf, call):
          # TODO(mdan): Also insert ops to re-fetch if variables are involved?
          pass  # Will be removed below.
      """
      # TODO(mdan): This is brittle. Reorganize the mechanism.
      statements = templates.replace(template, call=node.value)
      control_deps_guard = statements[-1]
      control_deps_guard.body = []

      # First, attempt to gate future evaluation of args. If that's not
      # possible, gate all remaining statements (and that may fail too, see
      # _visit_and_reindent.
      args_scope = anno.getanno(node.value, 'args_scope')
      guarded_args = tuple(args_scope.used & (args_scope.parent.modified
                                              | args_scope.parent.returned))
      if guarded_args:
        node = tuple(statements[:-1]) + (
            self._gate_symbols(control_deps_guard, guarded_args),)
      else:
        node = tuple(statements[:-1])
        # The mechanism will insert the guard statement later.
        self.indent_next = True
        self.next_indent_owner = control_deps_guard
    return node

 def visit_For(self, node):
   self.generic_visit(node)
   body_scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
   i_var = self.context.namer.new_symbol('i', body_scope.referenced)
   smart_loop_iter_var = self.context.namer.new_symbol('smart_loop_iter',
                                                       body_scope.referenced)
   cont_var = self.context.namer.new_symbol('cont', body_scope.referenced)
   # TODO(mdan): Use TensorListFromTensor(loop_iter) here.
   if anno.hasanno(node, 'extra_cond'):
     template = """
       i = 0
       smart_loop_iter = py2tf_utils.dynamic_dataset(loop_iter)
       cont, target = py2tf_utils.dynamic_for_cond(i, smart_loop_iter)
       while cont and extra_cond:
         body
         i += 1
         cont, target = py2tf_utils.dynamic_for_cond(i, smart_loop_iter)
     """
     return templates.replace(
         template,
         loop_iter=node.iter,
         target=node.target,
         body=node.body,
         i=i_var,
         smart_loop_iter=smart_loop_iter_var,
         cont=cont_var,
         extra_cond=anno.getanno(node, 'extra_cond'))
   else:
     template = """
       i = 0
       smart_loop_iter = py2tf_utils.dynamic_dataset(loop_iter)
       cont, target = py2tf_utils.dynamic_for_cond(i, smart_loop_iter)
       while cont:
         body
         i += 1
         cont, target = py2tf_utils.dynamic_for_cond(i, smart_loop_iter)
     """
     repl = templates.replace(
         template,
         loop_iter=node.iter,
         target=node.target,
         body=node.body,
         i=i_var,
         smart_loop_iter=smart_loop_iter_var,
         cont=cont_var)
     return repl

 def _create_cond_expr(self, results, test, body_name, orelse_name):
   if results is not None:
     template = """
       results = py2tf_utils.run_cond(test, body_name, orelse_name)
     """
     return templates.replace(
         template,
         test=test,
         results=results,
         body_name=body_name,
         orelse_name=orelse_name)
   else:
     template = """
       py2tf_utils.run_cond(test, body_name, orelse_name)
     """
     return templates.replace(
         template, test=test, body_name=body_name, orelse_name=orelse_name)

 def _create_continuation_check(self):
   template = """
     if not var_name:
       pass
   """
   cond, = templates.replace(template, var_name=self.continuation_uses[-1][1])
   cond.body = []
   return cond

  def _create_continuation_trigger(self):

    def template(var_name):  # pylint:disable=unused-argument
      var_name = True

    assign, = templates.replace(
        template, var_name=gast.Name(self.continuation_uses[-1][1], None, None))
    return assign