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

'''
**Beartype abstract syntax tree (AST) transformers** (i.e., low-level classes
instrumenting well-typed third-party 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
#  equivalent TypeguardTransformer.visit_ClassDef() method.
#* Generator transformers. The idea here is that @beartype *CAN* actually
#  automatically type-check generator yields, sends, and returns at runtime.
#  How? By automatically injecting appropriate die_if_unbearable() calls
#  type-checking the values to be yielded, sent, and returned against the
#  appropriate type hints of the current generator factory *BEFORE* yielding,
#  sending, and returning those values. Shockingly, typeguard already does this
#  -- which is all manner of impressive. See the
#  TypeguardTransformer._use_memo() context manager for working code. Wow!
#
#See also:
#    https://github.com/agronholm/typeguard/blob/master/src/typeguard/_transformer.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
#"executing" 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: Actually, while demonstrably awesome, even the above fails to suffice to
#to statically type-check "typing.LiteralString". We failed to fully read PEP
#675, which contains a section on inference. In the worst case, nothing less
#than a complete graph of the entire app and all transitive dependencies thereof
#suffices to decide whether a parameter satisfies "typing.LiteralString".
#
#Thankfully, 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,
    ImportFrom,
    Load,
    Module,
    Name,
    NodeTransformer,
    Str,
    alias,
)
from beartype.claw._ast._clawastmunge import copy_node_code_metadata
from beartype.claw._ast._clawasttest import is_node_callable_typed
from beartype.claw._clawtyping import NodeCallable
from beartype.typing import (
    List,
    Optional,
    Union,
)
from beartype._conf.confcls import BeartypeConf

# ....................{ SUBCLASSES                         }....................
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# CAUTION: To improve forward compatibility with the superclass API over which
# we have *NO* control, avoid accidental conflicts by suffixing *ALL* private
# and public attributes of this subclass by "_beartype".
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

#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).

    Attributes
    ----------
    conf : BeartypeConf
        **Beartype configuration** (i.e., dataclass configuring the
        :mod:`beartype.beartype` decorator for *all* decoratable objects
        recursively decorated by this node transformer).

    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
    '''

    # ..................{ INITIALIZERS                       }..................
    def __init__(
        self,

        # Mandatory keyword-only parameters.
        *,
        conf_beartype: BeartypeConf,
    ) -> None:
        '''
        Initialize this node transformer.

        Parameters
        ----------
        conf : BeartypeConf
            **Beartype configuration** (i.e., dataclass configuring the
            :mod:`beartype.beartype` decorator for *all* decoratable objects
            recursively decorated by this node transformer).
        '''
        assert isinstance(conf_beartype, BeartypeConf), (
            f'{repr(conf_beartype)} not beartype configuration.')

        # Initialize our superclass.
        super().__init__()

        # Classify all passed parameters.
        self._conf_beartype = conf_beartype

    # ..................{ VISITORS ~ module                  }..................
    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._importlib._clawimpload.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 in the list of all child
        # nodes of this parent module 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 module
        # nodes containing *NO* child nodes) below.
        node_import_beartype_attrs_index = -1

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

        # For the 0-based index and value of each direct child node of this
        # parent module node...
        #
        # This iteration efficiently finds "node_import_beartype_attrs_index"
        # (i.e., the 0-based index of the first safe position in the list of all
        # child nodes of this parent module node to insert an import statement
        # importing our beartype decorator). 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, this iteration
        # is guaranteed to exhibit worst-case O(1) time complexity. \o/
        #
        # Note that the "body" instance variable for module nodes is a list of
        # all child nodes of this parent module node.
        for node_import_beartype_attrs_index, node_import_prev in enumerate(
            node.body):
            # If it is *NOT* the case that this child node signifies either...
            if not (
                # 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(node_import_prev, Expr) and
                    isinstance(node_import_prev.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(node_import_prev, ImportFrom) and
                    node_import_prev.module == '__future__'
                )
            # Then immediately halt iteration, guaranteeing O(1) runtime.
            ):
                break
            # Else, this child node signifies either a module docstring of
            # future import. In this case, implicitly skip past this child node
            # to the next child node.
            #
        # "node_import_beartype_attrs_index" is now the index of the first safe
        # position in this list to insert output child import nodes below.

        # If this 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 >= 0:
            # Module-scoped import nodes (i.e., child nodes to be inserted under
            # the parent node encapsulating the currently visited submodule in
            # the AST for that module).
            #
            # Note that these nodes are intentionally *NOT* generalized into
            # global constants. In theory, doing so would reduce space and time
            # complexity by enabling efficient reuse here. In practice, doing so
            # would also be fundamentally wrong; these nodes are subsequently
            # modified to respect the source code metadata (e.g., line numbers)
            # of this AST module parent node, which prevents such trivial reuse.
            # Although we could further attempt to circumvent that by shallowly
            # or deeply copying from global constants, both the copy() and
            # deepcopy() functions defined by the standard "copy" module are
            # pure-Python and thus shockingly slow -- which defeats the purpose.

            # Our private beartype._decor.decorcore.beartype_object_nonfatal()
            # decorator.
            node_import_decorator = ImportFrom(
                module='beartype._decor.decorcore',

                #FIXME: Globalize this list as a new private constant, please.
                names=[alias('beartype_object_nonfatal')],
            )

            # 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 low-level utility functions.
            node_import_raiser = ImportFrom(
                module='beartype.door._doorcheck',

                #FIXME: Globalize this list as a new private constant, please.
                names=[alias('die_if_unbearable')],
            )

            # Copy all source code metadata (e.g., line numbers) from:
            # * For the first output import child node (in arbitrary order, as
            #   import statements are necessarily idempotent), from the input
            #   child node of this parent module node directly preceding this
            #   output node onto this output node.
            # * For all subsequent output import child nodes, from the preceding
            #   output node onto the current output node.
            copy_node_code_metadata(
                node_src=node_import_prev, node_trg=node_import_decorator)
            copy_node_code_metadata(
                node_src=node_import_prev, node_trg=node_import_raiser)

            # Insert these output import child nodes at this safe position of
            # the list of all child nodes of this parent module node.
            #
            # Note that this syntax efficiently (albeit unreadably) inserts
            # these output import child nodes at the desired index (in this
            # arbitrary order) of this parent module node.
            node.body[node_import_beartype_attrs_index:0] = (
                node_import_decorator, node_import_raiser)
        # 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

    # ..................{ VISITORS ~ callable                }..................
    def visit_FunctionDef(self, node: NodeCallable) -> Optional[NodeCallable]:
        '''
        Add a new child node to the passed **callable node** (i.e., node
        signifying the definition of a pure-Python function or method)
        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 : NodeCallable
            Callable node to be transformed.

        Returns
        ----------
        Optional[NodeCallable]
            This same callable node.
        '''

        # If that callable is annotated by one or more type hints and thus
        # unignorable...
        if is_node_callable_typed(node):
            #FIXME: Additionally pass the current beartype configuration as a
            #keyword-only "conf={conf}" parameter to this decorator, please.

            #FIXME: [CACHE] Consider generalizing the
            #BeartypeNodeTransformer.__new__() class method to internally cache
            #and return "BeartypeNodeTransformer" instances depending on the
            #passed "conf_beartype" parameter. In general, most codebases will
            #only leverage a single @beartype configuration (if any @beartype
            #configuration at all); ergo, caching improves everything by
            #enabling us to reuse the same "BeartypeNodeTransformer" instance
            #for every hooked module. Score @beartype!
            #
            #See the BeartypeConf.__new__() method for relevant logic. \o/
            #FIXME: Oh, wait. We probably do *NOT* want to cache -- at least,
            #not within defining a comparable reinit() method as we do for
            #"BeartypeCall". After retrieving a cached "BeartypeNodeTransformer"
            #instance, we'll need to immediately call
            #BeartypeNodeTransformer.reinit() to reinitialize that instance.
            #
            #This is all feasible, of course -- but let's just roll with the
            #naive implementation for now, 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=_NODE_CONTEXT_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)

            # Prepend this child decoration node to the beginning of the list of
            # all child decoration nodes for this parent callable node. Since
            # this list is "stored outermost first (i.e. the first in the list
            # will be applied last)", prepending guarantees that our decorator
            # will be applied last (i.e., *AFTER* all subsequent decorators).
            # This ensures that explicitly configured @beartype decorations
            # (e.g., "beartype(conf=BeartypeConf(...))") assume precedence over
            # implicitly configured @beartype decorations inserted by this hook.
            node.decorator_list.insert(0, decorate_callable)
        # Else, that callable is ignorable. 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

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

        Parameters
        ----------
        node : AnnAssign
            Annotated assignment node to be transformed.

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

            * If this annotated assignment 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 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=_NODE_CONTEXT_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=_NODE_CONTEXT_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 ~ constants                }....................
#FIXME: Shift into a new private "_clawastmagic" submodule, please.
_NODE_CONTEXT_LOAD = Load()
'''
**Node context load singleton** (i.e., object suitable for passing as the
``ctx`` keyword parameter accepted by the ``__init__()`` method of various
abstract syntax tree (AST) node classes).
'''
