[Feature Request] Generalize @typecheck to support PEP-compliant type hints

[leycec]

Extremely impressive. @awfutils.typecheck is the first practical attempt I've seen at performing static type-checking at runtime. Take all my thunderous clapping! 👏 👏

The current approach is outrageously useful, but appears to currently only support isinstance()-able classes rather than PEP-compliant type hints: e.g.,

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

Is that right? If so, that's still impressive tech for a several hundred-line decorator. Still, it'd be even more outrageously useful if we could generalize your @typecheck decorator to support arbitrary PEP-compliant type hints. Can we? Yes, we can.

You Know What @leycec Is Gonna Suggest Next...

That's right. We're talkin' @beartype, because of course we are. Implementing full-blown type-checking for all PEP standards past and present probably isn't somewhere you want to willingly go. Thankfully, @beartype already went there for you. 🥳

The core issue appears to be the current usage of the isinstance() builtin in the TypeCheckVisitor.visit_FunctionDef() method. Specifically, this AST transform:

        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),
        )

I'm fairly certain (...but technically uncertain, because lazy and thus untested) that replacing the above with the below should generalize @typecheck to support PEP-compliant type hints:

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]

The visit_Module() implementation is copied almost verbatim from a similar AST transform in the @beartype codebase itself. So, possibly working?

Regardless of where you choose to take this awesomeness, this has been a considerable inspiration. @beartype will probably end up assimilating this into itself, because everyone over at beartype/beartype#105 really wants this to happen.

In short, you're amazing.

[awf]

Thanks! Yes indeed @beartype is part of the plan, also for the O(1)-ness. I'll get on this just after I get the jaxtyping loveliness which triggered all this..

And of course I'm totally happy if this ends up in @beartype :)

[leycec]

Wunderbar! Take your languorous time and soak up those last rays of summer. They are precious – especially in Canada, where summer is a distant phenomena that mostly happens to other people.

And thanks for being so shockingly inventive. I never even knew about the visit_AnnAssign() method – which, in hindsight, kinda seems like it exists for the express purpose of supporting PEP 526 at runtime. We now give thanks. \o/