[Feature Request] Generalize @typecheck to support PEP-compliant type hints #2
Description
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.