#!/usr/bin/env python3
# --------------------( LICENSE                            )--------------------
# Copyright (c) 2014-2023 Beartype authors.
# See "LICENSE" for further details.

'''
**Beartype all-at-once low-level abstract syntax tree (AST) transformation.**

This private submodule defines the low-level abstract syntax tree (AST)
transformation automatically decorating well-typed third-party packages and
modules with runtime type-checking dynamically generated by the
:func:`beartype.beartype` decorator.

This private submodule is *not* intended for importation by downstream callers.
'''

# ....................{ TODO                               }....................
#FIXME: Additionally define a new BeartypeNodeTransformer.visit_ClassDef()
#method modelled after the equivelent TypeguardTransformer.visit_ClassDef()
#method residing at:
#    https://github.com/agronholm/typeguard/blob/master/src/typeguard/importhook.py

#FIXME: *OMG.* See also the third-party "executing" Python package:
#    https://github.com/alexmojaki/executing
#
#IPython itself internally leverages "executing" via "stack_data" (i.e., a
#slightly higher-level third-party Python package that internally leverages
#"executing") to syntax-highlight the currently executing AST node. Indeed,
#"executing" sports an intense test suite (much like ours) effectively
#guaranteeing a one-to-one mapping between stack frames and AST nodes.
#
#So, what's the Big Idea here? The Big Idea here is that @beartype can
#internally (...possibly only optionally, but possibly mandatorily) leverage
#"stack_data" to begin performing full-blown static type-checking at runtime --
#especially of mission critical type hints like "typing.LiteralString" which can
#*ONLY* be type-checked via static analysis. :o
#
#So, what's the Little Idea here? The Little Idea here is that @beartype can
#generate type-checking wrappers that type-check parameters or returns annotated
#by "typing.LiteralString" by calling an internal private utility function --
#say, "_die_unless_literalstring(func: Callable, arg_name: str) -> None" -- with
#"func" as the current type-checking wrapper and "arg_name" as either the name
#of that parameter or "return". The _die_unless_literalstring() raiser then:
#* Dynamically searches up the call stack for the stack frame encapsulating an
#  external call to the passed "func" callable.
#* Passes that stack frame to the "executing" package.
#* "executing" then returns the AST node corresponding to that stack frame.
#* Introspects that node for the passed parameter whose name is "arg_name".
#* Raises an exception unless the value of that parameter is an AST node
#  corresponding to a string literal.
#
#Of course, that won't necessarily be fast -- but it will be accurate. Since
#security trumps speed, speed is significantly less of a concern insofar as
#"typing.LiteralString" is concerned. Of course, we should also employ
#significant caching... if we even can.
#FIXME: The above idea generalizes from "typing.LiteralString" to other
#fascinating topics as well. Indeed, given sufficient caching, one could begin
#to internally generate and cache a mypy-like graph network whose nodes are
#typed attributes and whose edges are relations between those typed attributes.

# ....................{ IMPORTS                            }....................
from ast import (
    AST,
    AnnAssign,
    Call,
    Expr,
    FunctionDef,
    ImportFrom,
    Load,
    Module,
    Name,
    NodeTransformer,
    Str,
    alias,
)
from beartype.typing import (
    List,
    Union,
)
from beartype._util.py.utilpyversion import IS_PYTHON_AT_LEAST_3_8

# ....................{ SUBCLASSES                         }....................
#FIXME: Implement us up, please.
#FIXME: Docstring us up, please.
#FIXME: Unit test us up, please.
class BeartypeNodeTransformer(NodeTransformer):
    '''
    **Beartype abstract syntax tree (AST) node transformer** (i.e., visitor
    pattern recursively transforming the AST tree passed to the :meth:`visit`
    method by decorating all typed callables and classes by the
    :func:`beartype.beartype` decorator).

    See Also
    ----------
    * The `comparable "typeguard.importhook" submodule <typeguard import
      hook_>`__ implemented by the incomparable `@agronholm (Alex Grönholm)
      <agronholm_>`__.

    .. _agronholm:
       https://github.com/agronholm
    .. _typeguard import hook:
       https://github.com/agronholm/typeguard/blob/master/src/typeguard/importhook.py
    '''

    # ..................{ VISITORS                           }..................
    def visit_Module(self, node: Module) -> Module:
        '''
        Add a new abstract syntax tree (AST) child node to the passed AST module
        parent node encapsulating the module currently being loaded by the
        :class:`beartype.claw._clawloader.BeartypeSourceFileLoader` object,
        importing our private
        :func:`beartype._decor.decorcore.beartype_object_nonfatal` decorator for
        subsequent use by the other visitor methods defined by this class.

        Parameters
        ----------
        node : Module
            AST module parent node to be transformed.

        Returns
        ----------
        Module
            That same AST module parent node.
        '''

        # 0-based index of the first safe position of the list of all AST child
        # nodes of this AST module parent node to insert an import statement
        # importing our beartype decorator, initialized to the erroneous index
        # "-1" to enable detection of empty modules (i.e., modules whose AST
        # module nodes containing *NO* child nodes) below.
        node_import_beartype_attrs_index = -1

        # AST child node of this AST module parent node immediately preceding
        # the AST import child node to be added below, defaulting to this AST
        # module parent node to ensure that the _copy_node_code_metadata()
        # function below *ALWAYS* copies from a valid AST node for sanity.
        module_child: AST = node

        # Efficiently find this index. Since, this iteration is guaranteed to
        # exhibit worst-case O(1) time complexity despite superficially
        # appearing to perform a linear search of all n child nodes of this
        # module parent node and thus exhibit worst-case O(n) time complexity.
        #
        # For the 0-based index and value of each direct AST child node of this
        # AST module parent node...
        for node_import_beartype_attrs_index, module_child in enumerate(
            node.body):
            # If this child node signifies either...
            if (
                # A module docstring...
                #
                # If that module defines a docstring, that docstring *MUST* be
                # the first expression of that module. That docstring *MUST* be
                # explicitly found and iterated past to ensure that the import
                # statement added below appears *AFTER* rather than *BEFORE* any
                # docstring. (The latter would destroy the semantics of that
                # docstring by reducing that docstring to an ignorable string.)
                (
                    isinstance(module_child, Expr) and
                    isinstance(module_child.value, Str)
                ) or
                # A future import (i.e., import of the form
                # "from __future__ ...") *OR*...
                #
                # If that module performs one or more future imports, these
                # imports *MUST* necessarily be the first non-docstring
                # statement of that module and thus appear *BEFORE* all import
                # statements that are actually imports -- including the import
                # statement added below.
                (
                    isinstance(module_child, ImportFrom) and
                    module_child.module == '__future__'
                )
            ):
                # Then continue past this child node to the next child node.
                continue

        # If the 0-based index of the first safe position of the list of all AST
        # child nodes of this AST module parent node to insert import
        # statements importing various beartype attributes is *NOT* the
        # erroneous index to which this index was initialized above, this module
        # contains one or more child nodes and is thus non-empty. In this
        # case...
        if node_import_beartype_attrs_index != -1:
            # Tuple of all module-scoped import nodes (i.e., child nodes to be
            # inserted under the parent node encapsulating the currently visited
            # bmodule in the AST for that module).
            nodes_import_beartype_attr = (
                # Our public beartype.door.die_if_unbearable() raiser,
                # intentionally imported from our private
                # "beartype.door._doorcheck" submodule rather than our public
                # "beartype.door" subpackage. Why? Because the former consumes
                # marginally less space and time to import than the latter.
                # Whereas the latter imports the full "TypeHint" hierarchy, the
                # former only imports multiple low-level utility functions.
                ImportFrom(
                    module='beartype.door._doorcheck',
                    names=[alias('die_if_unbearable')],
                ),
                # Our private
                # beartype._decor.decorcore.beartype_object_nonfatal()
                # decorator.
                ImportFrom(
                    module='beartype._decor.decorcore',
                    names=[alias('beartype_object_nonfatal')],
                ),
            )

            # For each module-scoped import node to be inserted...
            for node_import_beartype_attr in nodes_import_beartype_attr:
                # Copy all source code metadata from the AST child node of this
                # AST module parent node immediately preceding this AST import
                # child node onto this AST import child node.
                _copy_node_code_metadata(
                    node_src=node, node_trg=node_import_beartype_attr)

                # Insert this AST import child node at this safe position of the
                # list of all AST child nodes of this AST module parent node.
                node.body.insert(
                    node_import_beartype_attrs_index, node_import_beartype_attr)
        # Else, this module is empty. In this case, silently reduce to a noop.
        # Since this edge case is *EXTREMELY* uncommon, avoid optimizing for
        # this edge case (here or elsewhere).

        # Recursively transform *ALL* AST child nodes of this AST module node.
        self.generic_visit(node)

        # Return this AST module node as is.
        return node


    def visit_FunctionDef(self, node: FunctionDef) -> FunctionDef:
        '''
        Add a new abstract syntax tree (AST) child node to the passed AST
        callable parent node, decorating that callable by our private
        :func:`beartype._decor.decorcore.beartype_object_nonfatal` decorator if
        and only if that callable is **typed** (i.e., annotated by a return type
        hint and/or one or more parameter type hints).

        Parameters
        ----------
        node : FunctionDef
            AST callable parent node to be transformed.

        Returns
        ----------
        FunctionDef
            That same AST callable parent node.
        '''

        # True only if that callable is annotated by a return type hint,
        # trivially decided in O(1) time.
        is_return_typed = bool(node.returns)

        # True only if that callable is annotated by one or more parameter type
        # hints, non-trivially decided in O(n) time for n the number of
        # parameters accepted by that callable.
        is_args_typed = False

        # If that callable is *NOT* annotated by a return type hint, fallback to
        # deciding whether that callable is annotated by one or more parameter
        # type hints. Since doing is considerably more computationally
        # expensive, do so *ONLY* as needed.
        if not is_return_typed:
            for arg in node.args.args:
                if arg.annotation:
                    is_args_typed = True
                    break
        # Else, that callable is annotated by a return type hint. In this case,
        # do *NOT* spend useless time deciding whether that callable is
        # annotated by one or more parameter type hints.

        # If that callable is typed (i.e., annotated by a return type hint
        # and/or one or more parameter type hints)...
        #
        # Note that the former is intentionally tested *BEFORE* the latter, as
        # the detecting former is O(1) time complexity and thus trivial.
        if is_return_typed or is_args_typed:
            #FIXME: Additionally pass the current beartype configuration as a
            #keyword-only "conf={conf}" parameter to this decorator, please.

            # AST decoration child node decorating that callable by our
            # beartype._decor.decorcore.beartype_object_nonfatal() decorator. Note
            # that this syntax derives from the example for the ast.arg() class:
            #     https://docs.python.org/3/library/ast.html#ast.arg
            decorate_callable = Name(id='beartype_object_nonfatal', ctx=Load())

            # Copy all source code metadata from this AST callable parent node
            # onto this AST decoration child node.
            _copy_node_code_metadata(node_src=node, node_trg=decorate_callable)

            #FIXME: *INSUFFICIENT.* We need to additionally avoid redecorating
            #callables already explicitly decorated by @beartype, as that
            #redecoration would erroneously take precedence over the explicit
            #decoration; the latter should *ALWAYS* take precedence over the
            #former, however, due to "conf=BeartypeConf(...)" parametrization.
            #Happily, this should be trivial ala:
            #    #FIXME: Note that "decorator_node.id == 'beartype'" is probably
            #    #an insufficient test, as decorators can be trivially renamed
            #    #or imported under differing names.
            #    for decorator_node in node.decorator_list:
            #        if (
            #            isinstance(decorator_node, Name) and
            #            decorator_node.id == 'beartype'
            #        ):
            #            break
            #    else:
            #        node.decorator_list.append(decorate_callable)

            # Append this AST decoration child node to the end of the list of
            # all AST decoration child nodes for this AST callable parent node.
            # Since this list is "stored outermost first (i.e. the first in the
            # list will be applied last)", appending guarantees that our
            # decorator will be applied first (i.e., *BEFORE* all subsequent
            # decorators). This is *NOT* simply obsequious greed. The @beartype
            # decorator generally requires that it precede other decorators that
            # obfuscate the identity of the original callable, including:
            # * The builtin @property decorator.
            # * The builtin @classmethod decorator.
            # * The builtin @staticmethod decorator.
            node.decorator_list.append(decorate_callable)
        # Else, that callable is untyped. In this case, avoid needlessly
        # decorating that callable by @beartype for efficiency.

        # Recursively transform *ALL* AST child nodes of this AST callable node.
        self.generic_visit(node)

        # Return this AST callable node as is.
        return node


    def visit_AnnAssign(self, node: AnnAssign) -> Union[AST, List[AST]]:
        '''
        Add a new abstract syntax tree (AST) child node to the passed AST
        **annotated assignment** (i.e., assignment of an attribute annotated by
        a :pep:`562`-compliant type hint) parent node, inserting a subsequent
        statement following that annotated assignment that type-checks that
        attribute against that type hint by passing both to our public
        :func:`beartype.door.is_bearable` tester.

        Parameters
        ----------
        node : AnnAssign
            AST annotated assignment parent node to be transformed.

        Returns
        ----------
        Union[AST, List[AST]]
            Either:

            * If this annotated assignment parent node is *not* **simple**
              (i.e., the attribute being assigned to is embedded in parentheses
              and thus denotes a full-blown Python expression rather than a
              simple attribute name), that same parent node unmodified.
            * Else, a 2-list comprising both that parent node and a new adjacent
              :class:`Call` node performing this type-check.

        See Also
        ----------
        https://github.com/awf/awfutils
            Third-party Python package whose ``@awfutils.typecheck`` decorator
            implements statement-level :func:`isinstance`-based type-checking in
            a similar manner, strongly inspiring this implementation. Thanks so
            much to Cambridge researcher @awf (Andrew Fitzgibbon) for the
            phenomenal inspiration!
        '''

        # Note that "AnnAssign" node subclass defines these instance variables:
        # * "node.annotation", a child node describing the PEP-compliant type
        #   hint annotating this assignment, typically an instance of either:
        #   * "ast.Name".
        #   * "ast.Str".
        # * "node.simple", a boolean that is true only if "node.target" is an
        #   "ast.Name" node.
        # * "node.target", a child node describing the target attribute assigned
        #   to by this assignment, guaranteed to be an instance of either:
        #   * "ast.Name", in which case this assignment is denoted as "simple"
        #     via the "node.simple" instance variable. This is the common case
        #     in which the attribute being assigned to is *NOT* embedded in
        #     parentheses and thus denotes a simple attribute name rather than a
        #     full-blown Python expression.
        #   * "ast.Attribute".
        #   * "ast.Subscript".
        # * "node.value", an optional child node describing the source value
        #   being assigned to this target attribute.

        #FIXME: Can and/or should we also support "node.target" child nodes that
        #are instances of "ast.Attribute" and "ast.Subscript"?
        # If this assignment is *NOT* simple, this assignment is *NOT* assigning
        # to an attribute name. In this case, silently ignore this assignment.
        if not node.simple:
            return node
        # Else, this assignment is simple and assigning to an attribute name.

        # Validate this fact.
        assert isinstance(node.target, Name)

        #FIXME: Additionally pass the current beartype configuration as a
        #keyword-only "conf={conf}" parameter to this raiser, please.

        # Child node referencing the function performing this type-checking,
        # previously imported at module scope by visit_FunctionDef() above.
        node_typecheck_function = Name('die_if_unbearable', ctx=Load())

        # Child node passing the value newly assigned to this attribute by this
        # assignment as the first parameter to die_if_unbearable().
        node_typecheck_pith = Name(node.target.id, ctx=Load())

        # Adjacent node type-checking this newly assigned attribute against the
        # PEP-compliant type hint annotating this assignment by deferring to our
        # die_if_unbearable() raiser.
        node_typecheck = Call(
            node_typecheck_function,
            [
                # Child node passing the value newly assigned to this
                # attribute by this assignment as the first parameter.
                node_typecheck_pith,
                # Child node passing the type hint annotating this assignment as
                # the second parameter.
                node.annotation,
            ],
            [],
        )

        # Copy all source code metadata from this AST annotated assignment node
        # onto *ALL* AST nodes created above.
        _copy_node_code_metadata(
            node_src=node, node_trg=node_typecheck_function)
        _copy_node_code_metadata(node_src=node, node_trg=node_typecheck_pith)
        _copy_node_code_metadata(node_src=node, node_trg=node_typecheck)

        #FIXME: Can we replace this inefficient list with an efficient tuple?
        #Probably not. Let's avoid doing so for the moment, as the "ast" API is
        #obstruse enough as it is.
        # Return a list comprising these two adjacent nodes.
        return [node, node_typecheck]

# ....................{ PRIVATE ~ copiers                  }....................
def _copy_node_code_metadata(node_src: AST, node_trg: AST) -> None:
    '''
    Copy all **source code metadata** (i.e., beginning and ending line and
    column numbers) from the passed source abstract syntax tree (AST) node onto
    the passed target AST node.

    This function is an efficient alternative to:
    * The extremely inefficient (albeit still useful)
      :func:`ast.fix_missing_locations` function.
    * The mildly inefficient (and mostly useless) :func:`ast.copy_location`
      function.

    The tradeoffs are as follows:

    * :func:`ast.fix_missing_locations` is ``O(n)`` time complexity for ``n``
      the number of AST nodes across the entire AST tree, but requires only a
      single trivial call and is thus considerably more "plug-and-play" than
      this function.
    * This function is ``O(1)`` time complexity irrespective of the size of the
      AST tree, but requires one still mostly trivial call for each synthetic
      AST node inserted into the AST tree by the
      :class:`BeartypeNodeTransformer` above.

    Parameters
    ----------
    node_src: AST
        Source AST node to copy source code metadata from.
    node_trg: AST
        Target AST node to copy source code metadata onto.

    See Also
    ----------
    :func:`ast.copy_location`
        Less efficient analogue of this function running in ``O(k)`` time
        complexity for ``k`` the number of types of source code metadata.
        Typically, ``k == 4``.
    '''
    assert isinstance(node_src, AST), f'{repr(node_src)} not AST node.'
    assert isinstance(node_trg, AST), f'{repr(node_trg)} not AST node.'

    # Copy all source code metadata from this source to target AST node.
    node_trg.lineno     = node_src.lineno
    node_trg.col_offset = node_src.col_offset

    # If the active Python interpreter targets Python >= 3.8, then additionally
    # copy all source code metadata exposed by Python >= 3.8.
    if IS_PYTHON_AT_LEAST_3_8:
        node_trg.end_lineno     = node_src.end_lineno  # type: ignore[attr-defined]
        node_trg.end_col_offset = node_src.end_col_offset  # type: ignore[attr-defined]
