awf / awfutils Goto Github PK

# I suspect this fails hard, but am lazy and thus did not test.
@typecheck
def foo(x : List[int], y : int):
  z : List[int] = x * y
  w : float = z[0] * 3.2
  return w

foo([3, 2, 1], 1.3)

        node_assert = ast.Assert(
            test=ast.Call(
                ast.Name("isinstance", ctx=ast.Load()),
                [
                    ast.Name(
                        node.target.id, ctx=ast.Load()
                    ),  # Convert ctx from Store to Load
                    node.annotation,
                ],
                [],
            ),
            msg=ast.Constant(value=f"{node.target.id} not a {annot_str}", kind=None),
        )

class TypeCheckVisitor(ast.NodeTransformer):
    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, importing
        the :func:`beartype.abby.die_if_unbearable` runtime type-checker 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.
        import_beartype_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 import_beartype_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 an import
        # statement importing our beartype decorator 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 import_beartype_index != -1:
            # AST import child node importing our private
            # beartype._decor.decorcore.beartype_object_nonfatal() decorator for
            # subsequent use by the other visitor methods defined by this class.
            import_beartype = ImportFrom(
                module='beartype.abby',
                names=[alias('die_if_unbearable')],
            )

            # 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=import_beartype)

            # 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(import_beartype_index, import_beartype)
        # 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_AnnAssign(self, node):
        # An assignment with a type annotation.
        # node.target : single node and can be a Name, a Attribute or a Subscript.
        # node.annotation : annotation, such as a Str or Name node.
        # node.value : single optional node.
        # node.simple : True for a Name node in target that do not appear between
        #               parentheses and are hence pure names and not expressions.

        if not node.simple:
            return node

        assert isinstance(node.target, ast.Name)  # Should be guaranteed by node.simple

        node_typecheck = ast.Call(
            ast.Name('die_if_unbearable', ctx=ast.Load()),
            [
                ast.Name(
                    node.target.id, ctx=ast.Load()
                ),  # Convert ctx from Store to Load
                node.annotation,
            ],
            [],
        )

        node_typecheck = ast.copy_location(node_typecheck, node)
        ast.fix_missing_locations(node_typecheck)
        return [node, node_typecheck]