Skip to content

gh-118761: Improve import time of annotationlib #132028

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
Closed

gh-118761: Improve import time of annotationlib #132028

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f78c6ec
split annotationlib into two
DavidCEllis Apr 2, 2025
10f35b2
Also defer functools import
DavidCEllis Apr 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
233 changes: 26 additions & 207 deletions Lib/annotationlib.py → Lib/annotationlib/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
"""Helpers for introspecting and wrapping annotations."""

import ast
import builtins
import enum
import functools
import keyword
import sys
import types
Expand All @@ -20,6 +18,26 @@
]


class _LazyImporter:
def __getattr__(self, name):
if name == "ast":
import ast
self.ast = ast
return ast
elif name == "_Stringifier":
from ._stringifier import _Stringifier
self._Stringifier = _Stringifier
return _Stringifier
elif name == "functools":
import functools
self.functools = functools
return functools
else:
raise AttributeError(
f"{self.__class__.__name__!r} object has no attribute {name!r}"
)


class Format(enum.IntEnum):
VALUE = 1
VALUE_WITH_FAKE_GLOBALS = 2
Expand All @@ -29,6 +47,7 @@ class Format(enum.IntEnum):

_Union = None
_sentinel = object()
_laz = _LazyImporter()

# Slots shared by ForwardRef and _Stringifier. The __forward__ names must be
# preserved for compatibility with the old typing.ForwardRef class. The remaining
Expand Down Expand Up @@ -205,7 +224,7 @@ def __forward_arg__(self):
if self.__arg__ is not None:
return self.__arg__
if self.__ast_node__ is not None:
self.__arg__ = ast.unparse(self.__ast_node__)
self.__arg__ = _laz.ast.unparse(self.__ast_node__)
return self.__arg__
raise AssertionError(
"Attempted to access '__forward_arg__' on an uninitialized ForwardRef"
Expand Down Expand Up @@ -265,206 +284,6 @@ def __repr__(self):
return f"ForwardRef({self.__forward_arg__!r}{module_repr})"


class _Stringifier:
# Must match the slots on ForwardRef, so we can turn an instance of one into an
# instance of the other in place.
__slots__ = _SLOTS

def __init__(
self,
node,
globals=None,
owner=None,
is_class=False,
cell=None,
*,
stringifier_dict,
):
# Either an AST node or a simple str (for the common case where a ForwardRef
# represent a single name).
assert isinstance(node, (ast.AST, str))
self.__arg__ = None
self.__forward_evaluated__ = False
self.__forward_value__ = None
self.__forward_is_argument__ = False
self.__forward_is_class__ = is_class
self.__forward_module__ = None
self.__code__ = None
self.__ast_node__ = node
self.__globals__ = globals
self.__cell__ = cell
self.__owner__ = owner
self.__stringifier_dict__ = stringifier_dict

def __convert_to_ast(self, other):
if isinstance(other, _Stringifier):
if isinstance(other.__ast_node__, str):
return ast.Name(id=other.__ast_node__)
return other.__ast_node__
elif isinstance(other, slice):
return ast.Slice(
lower=(
self.__convert_to_ast(other.start)
if other.start is not None
else None
),
upper=(
self.__convert_to_ast(other.stop)
if other.stop is not None
else None
),
step=(
self.__convert_to_ast(other.step)
if other.step is not None
else None
),
)
else:
return ast.Constant(value=other)

def __get_ast(self):
node = self.__ast_node__
if isinstance(node, str):
return ast.Name(id=node)
return node

def __make_new(self, node):
stringifier = _Stringifier(
node,
self.__globals__,
self.__owner__,
self.__forward_is_class__,
stringifier_dict=self.__stringifier_dict__,
)
self.__stringifier_dict__.stringifiers.append(stringifier)
return stringifier

# Must implement this since we set __eq__. We hash by identity so that
# stringifiers in dict keys are kept separate.
def __hash__(self):
return id(self)

def __getitem__(self, other):
# Special case, to avoid stringifying references to class-scoped variables
# as '__classdict__["x"]'.
if self.__ast_node__ == "__classdict__":
raise KeyError
if isinstance(other, tuple):
elts = [self.__convert_to_ast(elt) for elt in other]
other = ast.Tuple(elts)
else:
other = self.__convert_to_ast(other)
assert isinstance(other, ast.AST), repr(other)
return self.__make_new(ast.Subscript(self.__get_ast(), other))

def __getattr__(self, attr):
return self.__make_new(ast.Attribute(self.__get_ast(), attr))

def __call__(self, *args, **kwargs):
return self.__make_new(
ast.Call(
self.__get_ast(),
[self.__convert_to_ast(arg) for arg in args],
[
ast.keyword(key, self.__convert_to_ast(value))
for key, value in kwargs.items()
],
)
)

def __iter__(self):
yield self.__make_new(ast.Starred(self.__get_ast()))

def __repr__(self):
if isinstance(self.__ast_node__, str):
return self.__ast_node__
return ast.unparse(self.__ast_node__)

def __format__(self, format_spec):
raise TypeError("Cannot stringify annotation containing string formatting")

def _make_binop(op: ast.AST):
def binop(self, other):
return self.__make_new(
ast.BinOp(self.__get_ast(), op, self.__convert_to_ast(other))
)

return binop

__add__ = _make_binop(ast.Add())
__sub__ = _make_binop(ast.Sub())
__mul__ = _make_binop(ast.Mult())
__matmul__ = _make_binop(ast.MatMult())
__truediv__ = _make_binop(ast.Div())
__mod__ = _make_binop(ast.Mod())
__lshift__ = _make_binop(ast.LShift())
__rshift__ = _make_binop(ast.RShift())
__or__ = _make_binop(ast.BitOr())
__xor__ = _make_binop(ast.BitXor())
__and__ = _make_binop(ast.BitAnd())
__floordiv__ = _make_binop(ast.FloorDiv())
__pow__ = _make_binop(ast.Pow())

del _make_binop

def _make_rbinop(op: ast.AST):
def rbinop(self, other):
return self.__make_new(
ast.BinOp(self.__convert_to_ast(other), op, self.__get_ast())
)

return rbinop

__radd__ = _make_rbinop(ast.Add())
__rsub__ = _make_rbinop(ast.Sub())
__rmul__ = _make_rbinop(ast.Mult())
__rmatmul__ = _make_rbinop(ast.MatMult())
__rtruediv__ = _make_rbinop(ast.Div())
__rmod__ = _make_rbinop(ast.Mod())
__rlshift__ = _make_rbinop(ast.LShift())
__rrshift__ = _make_rbinop(ast.RShift())
__ror__ = _make_rbinop(ast.BitOr())
__rxor__ = _make_rbinop(ast.BitXor())
__rand__ = _make_rbinop(ast.BitAnd())
__rfloordiv__ = _make_rbinop(ast.FloorDiv())
__rpow__ = _make_rbinop(ast.Pow())

del _make_rbinop

def _make_compare(op):
def compare(self, other):
return self.__make_new(
ast.Compare(
left=self.__get_ast(),
ops=[op],
comparators=[self.__convert_to_ast(other)],
)
)

return compare

__lt__ = _make_compare(ast.Lt())
__le__ = _make_compare(ast.LtE())
__eq__ = _make_compare(ast.Eq())
__ne__ = _make_compare(ast.NotEq())
__gt__ = _make_compare(ast.Gt())
__ge__ = _make_compare(ast.GtE())

del _make_compare

def _make_unary_op(op):
def unary_op(self):
return self.__make_new(ast.UnaryOp(op, self.__get_ast()))

return unary_op

__invert__ = _make_unary_op(ast.Invert())
__pos__ = _make_unary_op(ast.UAdd())
__neg__ = _make_unary_op(ast.USub())

del _make_unary_op


class _StringifierDict(dict):
def __init__(self, namespace, globals=None, owner=None, is_class=False):
super().__init__(namespace)
Expand All @@ -475,7 +294,7 @@ def __init__(self, namespace, globals=None, owner=None, is_class=False):
self.stringifiers = []

def __missing__(self, key):
fwdref = _Stringifier(
fwdref = _laz._Stringifier(
key,
globals=self.globals,
owner=self.owner,
Expand Down Expand Up @@ -537,7 +356,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
name = freevars[i]
else:
name = "__cell__"
fwdref = _Stringifier(name, stringifier_dict=globals)
fwdref = _laz._Stringifier(name, stringifier_dict=globals)
new_closure.append(types.CellType(fwdref))
closure = tuple(new_closure)
else:
Expand Down Expand Up @@ -588,7 +407,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
name = freevars[i]
else:
name = "__cell__"
fwdref = _Stringifier(
fwdref = _laz._Stringifier(
name,
cell=cell,
owner=owner,
Expand Down Expand Up @@ -772,7 +591,7 @@ def get_annotations(
if hasattr(unwrap, "__wrapped__"):
unwrap = unwrap.__wrapped__
continue
if isinstance(unwrap, functools.partial):
if isinstance(unwrap, _laz.functools.partial):
unwrap = unwrap.func
continue
break
Expand Down
Loading
Loading