# Copyright 2017 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.
# ==============================================================================
"""AST node annotation support.

Adapted from Tangent.
"""

import enum

# pylint:disable=g-bad-import-order

import gast
# pylint:enable=g-bad-import-order


# TODO(mdan): Shorten the names.
# These names are heavily used, and anno.blaa
# TODO(mdan): Replace the attr-dict mechanism with a more typed solution.


class NoValue(enum.Enum):
  """Base class for different types of AST annotations."""

  def of(self, node, default=None):
    return getanno(node, self, default=default)

  def add_to(self, node, value):
    setanno(node, self, value)

  def exists(self, node):
    return hasanno(node, self)

  def __repr__(self):
    return str(self.name)


class Basic(NoValue):
  """Container for basic annotation keys.

  The enum values are used strictly for documentation purposes.
  """

  QN = 'Qualified name, as it appeared in the code. See qual_names.py.'
  SKIP_PROCESSING = (
      'This node should be preserved as is and not processed any further.')
  INDENT_BLOCK_REMAINDER = (
      'When a node is annotated with this, the remainder of the block should'
      ' be indented below it. The annotation contains a tuple'
      ' (new_body, name_map), where `new_body` is the new indented block and'
      ' `name_map` allows renaming symbols.')
  ORIGIN = ('Information about the source code that converted code originated'
            ' from. See origin_information.py.')
  DIRECTIVES = ('User directives associated with a statement or a variable.'
                ' Typically, they affect the immediately-enclosing statement.')

  EXTRA_LOOP_TEST = (
      'A special annotation containing additional test code to be executed in'
      ' for loops.')


class Static(NoValue):
  """Container for static analysis annotation keys.

  The enum values are used strictly for documentation purposes.
  """

  # Symbols
  # These flags are boolean.
  IS_PARAM = 'Symbol is a parameter to the function being analyzed.'

  # Scopes
  # Scopes are represented by objects of type activity.Scope.
  SCOPE = 'The scope for the annotated node. See activity.py.'
  # TODO(mdan): Drop these in favor of accessing the child's SCOPE.
  ARGS_SCOPE = 'The scope for the argument list of a function call.'
  COND_SCOPE = 'The scope for the test node of a conditional statement.'
  BODY_SCOPE = (
      'The scope for the main body of a statement (True branch for if '
      'statements, main body for loops).')
  ORELSE_SCOPE = (
      'The scope for the orelse body of a statement (False branch for if '
      'statements, orelse body for loops).')

  # Static analysis annotations.
  DEFINITIONS = (
      'Reaching definition information. See reaching_definitions.py.')
  ORIG_DEFINITIONS = (
      'The value of DEFINITIONS that applied to the original code before any'
      ' conversion.')
  DEFINED_FNS_IN = (
      'Local function definitions that may exist when exiting the node. See'
      ' reaching_fndefs.py')
  DEFINED_VARS_IN = (
      'Symbols defined when entering the node. See reaching_definitions.py.')
  LIVE_VARS_OUT = ('Symbols live when exiting the node. See liveness.py.')
  LIVE_VARS_IN = ('Symbols live when entering the node. See liveness.py.')
  TYPES = 'Static type information. See type_inference.py.'
  CLOSURE_TYPES = 'Types of closure symbols at each detected call site.'
  VALUE = 'Static value information. See type_inference.py.'


FAIL = object()


def keys(node, field_name='___pyct_anno'):
  if not hasattr(node, field_name):
    return frozenset()
  return frozenset(getattr(node, field_name).keys())


def getanno(node, key, default=FAIL, field_name='___pyct_anno'):
  if (default is FAIL or (hasattr(node, field_name) and
                          (key in getattr(node, field_name)))):
    return getattr(node, field_name)[key]
  return default


def hasanno(node, key, field_name='___pyct_anno'):
  return hasattr(node, field_name) and key in getattr(node, field_name)


def setanno(node, key, value, field_name='___pyct_anno'):
  annotations = getattr(node, field_name, {})
  setattr(node, field_name, annotations)
  annotations[key] = value

  # So that the annotations survive gast_to_ast() and ast_to_gast()
  if field_name not in node._fields:
    node._fields += (field_name,)


def delanno(node, key, field_name='___pyct_anno'):
  annotations = getattr(node, field_name)
  del annotations[key]
  if not annotations:
    delattr(node, field_name)
    node._fields = tuple(f for f in node._fields if f != field_name)


def copyanno(from_node, to_node, key, field_name='___pyct_anno'):
  if hasanno(from_node, key, field_name=field_name):
    setanno(
        to_node,
        key,
        getanno(from_node, key, field_name=field_name),
        field_name=field_name)


def dup(node, copy_map, field_name='___pyct_anno'):
  """Recursively copies annotations in an AST tree.

  Args:
    node: ast.AST
    copy_map: Dict[Hashable, Hashable], maps a source anno key to a destination
        key. All annotations with the source key will be copied to identical
        annotations with the destination key.
    field_name: str
  """
  for n in gast.walk(node):
    for k in copy_map:
      if hasanno(n, k, field_name):
        setanno(n, copy_map[k], getanno(n, k, field_name), field_name)