# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""Overloads all variable read operations."""

import gast

from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import templates


class VariableAccessTransformer(converter.Base):
  """Rewrites basic symbol reads.

  This transformer rewrites variable reads with a "read" operator which allows
  tracking activity.

  Example:

  For a basic statement:

      a = b + c

  This is translated to:

      a = ld(b) + ld(c)

  Augmented assignment operations also introduce a `ld` operator:

      a += b

  The assignment target also receives an operator to properly represent the
  read:

      a = ld(a)
      a += ld(b)
  """

  def visit_Name(self, node):
    # Only the loads which existed in the original code are overloaded.
    if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
      return node
    if isinstance(node.ctx, gast.Load):
      node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
    return node

  def visit_Delete(self, node):
    node = self.generic_visit(node)

    rewrite_targets = []
    for tgt in node.targets:
      # Don't rewrite composites like `del a[0]`.
      if isinstance(tgt, gast.Name):
        rewrite_targets.append(tgt)

    if not rewrite_targets:
      return node

    results = []
    for tgt in rewrite_targets:
      template = """
        var_ = ag__.Undefined(var_name)
      """
      results.extend(templates.replace(
          template, var_=tgt, var_name=gast.Constant(tgt.id, kind=None)))
    remaining_targets = [n for n in node.targets if n not in rewrite_targets]
    if remaining_targets:
      results.append(gast.Delete(targets=remaining_targets))

    return results

  def visit_AugAssign(self, node):
    if isinstance(node.target, gast.Name):
      template = """
        var_ = ag__.ld(var_)
        original
      """
      node = templates.replace(template, var_=node.target, original=node)
    else:
      node = self.generic_visit(node)
    return node


def transform(node, ctx):
  return VariableAccessTransformer(ctx).visit(node)

335ae5fbc2ed1154e3bbd83e6a7c4508c0bb9fed
<mask token>


class VariableAccessTransformer(converter.Base):
    <mask token>

    def visit_Name(self, node):
        if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
            return node
        if isinstance(node.ctx, gast.Load):
            node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
        return node

    def visit_Delete(self, node):
        node = self.generic_visit(node)
        rewrite_targets = []
        for tgt in node.targets:
            if isinstance(tgt, gast.Name):
                rewrite_targets.append(tgt)
        if not rewrite_targets:
            return node
        results = []
        for tgt in rewrite_targets:
            template = '\n        var_ = ag__.Undefined(var_name)\n      '
            results.extend(templates.replace(template, var_=tgt, var_name=
                gast.Constant(tgt.id, kind=None)))
        remaining_targets = [n for n in node.targets if n not in
            rewrite_targets]
        if remaining_targets:
            results.append(gast.Delete(targets=remaining_targets))
        return results

    def visit_AugAssign(self, node):
        if isinstance(node.target, gast.Name):
            template = (
                '\n        var_ = ag__.ld(var_)\n        original\n      ')
            node = templates.replace(template, var_=node.target, original=node)
        else:
            node = self.generic_visit(node)
        return node


<mask token>

<mask token>


class VariableAccessTransformer(converter.Base):
    """Rewrites basic symbol reads.

  This transformer rewrites variable reads with a "read" operator which allows
  tracking activity.

  Example:

  For a basic statement:

      a = b + c

  This is translated to:

      a = ld(b) + ld(c)

  Augmented assignment operations also introduce a `ld` operator:

      a += b

  The assignment target also receives an operator to properly represent the
  read:

      a = ld(a)
      a += ld(b)
  """

    def visit_Name(self, node):
        if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
            return node
        if isinstance(node.ctx, gast.Load):
            node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
        return node

    def visit_Delete(self, node):
        node = self.generic_visit(node)
        rewrite_targets = []
        for tgt in node.targets:
            if isinstance(tgt, gast.Name):
                rewrite_targets.append(tgt)
        if not rewrite_targets:
            return node
        results = []
        for tgt in rewrite_targets:
            template = '\n        var_ = ag__.Undefined(var_name)\n      '
            results.extend(templates.replace(template, var_=tgt, var_name=
                gast.Constant(tgt.id, kind=None)))
        remaining_targets = [n for n in node.targets if n not in
            rewrite_targets]
        if remaining_targets:
            results.append(gast.Delete(targets=remaining_targets))
        return results

    def visit_AugAssign(self, node):
        if isinstance(node.target, gast.Name):
            template = (
                '\n        var_ = ag__.ld(var_)\n        original\n      ')
            node = templates.replace(template, var_=node.target, original=node)
        else:
            node = self.generic_visit(node)
        return node


<mask token>

<mask token>


class VariableAccessTransformer(converter.Base):
    """Rewrites basic symbol reads.

  This transformer rewrites variable reads with a "read" operator which allows
  tracking activity.

  Example:

  For a basic statement:

      a = b + c

  This is translated to:

      a = ld(b) + ld(c)

  Augmented assignment operations also introduce a `ld` operator:

      a += b

  The assignment target also receives an operator to properly represent the
  read:

      a = ld(a)
      a += ld(b)
  """

    def visit_Name(self, node):
        if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
            return node
        if isinstance(node.ctx, gast.Load):
            node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
        return node

    def visit_Delete(self, node):
        node = self.generic_visit(node)
        rewrite_targets = []
        for tgt in node.targets:
            if isinstance(tgt, gast.Name):
                rewrite_targets.append(tgt)
        if not rewrite_targets:
            return node
        results = []
        for tgt in rewrite_targets:
            template = '\n        var_ = ag__.Undefined(var_name)\n      '
            results.extend(templates.replace(template, var_=tgt, var_name=
                gast.Constant(tgt.id, kind=None)))
        remaining_targets = [n for n in node.targets if n not in
            rewrite_targets]
        if remaining_targets:
            results.append(gast.Delete(targets=remaining_targets))
        return results

    def visit_AugAssign(self, node):
        if isinstance(node.target, gast.Name):
            template = (
                '\n        var_ = ag__.ld(var_)\n        original\n      ')
            node = templates.replace(template, var_=node.target, original=node)
        else:
            node = self.generic_visit(node)
        return node


def transform(node, ctx):
    return VariableAccessTransformer(ctx).visit(node)

<mask token>
import gast
from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import templates


class VariableAccessTransformer(converter.Base):
    """Rewrites basic symbol reads.

  This transformer rewrites variable reads with a "read" operator which allows
  tracking activity.

  Example:

  For a basic statement:

      a = b + c

  This is translated to:

      a = ld(b) + ld(c)

  Augmented assignment operations also introduce a `ld` operator:

      a += b

  The assignment target also receives an operator to properly represent the
  read:

      a = ld(a)
      a += ld(b)
  """

    def visit_Name(self, node):
        if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
            return node
        if isinstance(node.ctx, gast.Load):
            node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
        return node

    def visit_Delete(self, node):
        node = self.generic_visit(node)
        rewrite_targets = []
        for tgt in node.targets:
            if isinstance(tgt, gast.Name):
                rewrite_targets.append(tgt)
        if not rewrite_targets:
            return node
        results = []
        for tgt in rewrite_targets:
            template = '\n        var_ = ag__.Undefined(var_name)\n      '
            results.extend(templates.replace(template, var_=tgt, var_name=
                gast.Constant(tgt.id, kind=None)))
        remaining_targets = [n for n in node.targets if n not in
            rewrite_targets]
        if remaining_targets:
            results.append(gast.Delete(targets=remaining_targets))
        return results

    def visit_AugAssign(self, node):
        if isinstance(node.target, gast.Name):
            template = (
                '\n        var_ = ag__.ld(var_)\n        original\n      ')
            node = templates.replace(template, var_=node.target, original=node)
        else:
            node = self.generic_visit(node)
        return node


def transform(node, ctx):
    return VariableAccessTransformer(ctx).visit(node)

# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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
#
#     http://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.
# ==============================================================================
"""Overloads all variable read operations."""

import gast

from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import templates


class VariableAccessTransformer(converter.Base):
  """Rewrites basic symbol reads.

  This transformer rewrites variable reads with a "read" operator which allows
  tracking activity.

  Example:

  For a basic statement:

      a = b + c

  This is translated to:

      a = ld(b) + ld(c)

  Augmented assignment operations also introduce a `ld` operator:

      a += b

  The assignment target also receives an operator to properly represent the
  read:

      a = ld(a)
      a += ld(b)
  """

  def visit_Name(self, node):
    # Only the loads which existed in the original code are overloaded.
    if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
      return node
    if isinstance(node.ctx, gast.Load):
      node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
    return node

  def visit_Delete(self, node):
    node = self.generic_visit(node)

    rewrite_targets = []
    for tgt in node.targets:
      # Don't rewrite composites like `del a[0]`.
      if isinstance(tgt, gast.Name):
        rewrite_targets.append(tgt)

    if not rewrite_targets:
      return node

    results = []
    for tgt in rewrite_targets:
      template = """
        var_ = ag__.Undefined(var_name)
      """
      results.extend(templates.replace(
          template, var_=tgt, var_name=gast.Constant(tgt.id, kind=None)))
    remaining_targets = [n for n in node.targets if n not in rewrite_targets]
    if remaining_targets:
      results.append(gast.Delete(targets=remaining_targets))

    return results

  def visit_AugAssign(self, node):
    if isinstance(node.target, gast.Name):
      template = """
        var_ = ag__.ld(var_)
        original
      """
      node = templates.replace(template, var_=node.target, original=node)
    else:
      node = self.generic_visit(node)
    return node


def transform(node, ctx):
  return VariableAccessTransformer(ctx).visit(node)
