gh-141388: Fully support non-function callables as annotate functions

Lib/annotationlib.py

@@ -728,13 +728,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
  annotate, owner, is_class, globals, allow_evaluation=False
  )
  func = types.FunctionType(
- annotate.__code__,
+_get_annotate_attr(annotate, "__code__"),
  globals,
  closure=closure,
- argdefs=annotate.__defaults__,
- kwdefaults=annotate.__kwdefaults__,
+ argdefs=_get_annotate_attr(annotate, "__defaults__", None),
+ kwdefaults=_get_annotate_attr(annotate, "__kwdefaults__", None),
  )
- annos = func(Format.VALUE_WITH_FAKE_GLOBALS)
+ annos = _direct_call_annotate(func, annotate, Format.VALUE_WITH_FAKE_GLOBALS)
  if _is_evaluate:
  return _stringify_single(annos)
  return{
@@ -759,11 +759,21 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
  # reconstruct the source. But in the dictionary that we eventually return, we
  # want to return objects with more user-friendly behavior, such as an __eq__
  # that returns a bool and an defined set of attributes.
- namespace ={**annotate.__builtins__, **annotate.__globals__}
+
+ # Grab and store all the annotate function attributes that we might need to access
+ # multiple times as variables, as this could be a bit expensive for non-functions.
+ annotate_globals = _get_annotate_attr(annotate, "__globals__",{})
+ annotate_code = _get_annotate_attr(annotate, "__code__")
+ annotate_defaults = _get_annotate_attr(annotate, "__defaults__", None)
+ annotate_kwdefaults = _get_annotate_attr(annotate, "__kwdefaults__", None)
+ namespace ={
+ **_get_annotate_attr(annotate, "__builtins__",{}),
+ **annotate_globals
+ }
  is_class = isinstance(owner, type)
  globals = _StringifierDict(
  namespace,
- globals=annotate.__globals__,
+ globals=annotate_globals,
  owner=owner,
  is_class=is_class,
  format=format,
@@ -772,14 +782,14 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
  annotate, owner, is_class, globals, allow_evaluation=True
  )
  func = types.FunctionType(
-annotate.__code__,
+annotate_code,
  globals,
  closure=closure,
- argdefs=annotate.__defaults__,
- kwdefaults=annotate.__kwdefaults__,
+ argdefs=annotate_defaults,
+ kwdefaults=annotate_kwdefaults,
  )
  try:
- result = func(Format.VALUE_WITH_FAKE_GLOBALS)
+ result = _direct_call_annotate(func, annotate, Format.VALUE_WITH_FAKE_GLOBALS)
  except NotImplementedError:
  # FORWARDREF and VALUE_WITH_FAKE_GLOBALS not supported, fall back to VALUE
  return annotate(Format.VALUE)
@@ -793,7 +803,7 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
  # a value in certain cases where an exception gets raised during evaluation.
  globals = _StringifierDict(
 {},
- globals=annotate.__globals__,
+ globals=annotate_globals,
  owner=owner,
  is_class=is_class,
  format=format,
@@ -802,13 +812,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
  annotate, owner, is_class, globals, allow_evaluation=False
  )
  func = types.FunctionType(
-annotate.__code__,
+annotate_code,
  globals,
  closure=closure,
- argdefs=annotate.__defaults__,
- kwdefaults=annotate.__kwdefaults__,
+ argdefs=annotate_defaults,
+ kwdefaults=annotate_kwdefaults,
  )
- result = func(Format.VALUE_WITH_FAKE_GLOBALS)
+ result = _direct_call_annotate(func, annotate, Format.VALUE_WITH_FAKE_GLOBALS)
  globals.transmogrify(cell_dict)
  if _is_evaluate:
  if isinstance(result, ForwardRef):
@@ -833,12 +843,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
 
 
 def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluation):
- if not annotate.__closure__:
+ closure = _get_annotate_attr(annotate, "__closure__", None)
+ if not closure:
  return None, None
- freevars = annotate.__code__.co_freevars
+ freevars = _get_annotate_attr(annotate, "__code__", None).co_freevars
  new_closure = []
  cell_dict ={}
- for i, cell in enumerate(annotate.__closure__):
+ for i, cell in enumerate(closure):
  if i < len(freevars):
  name = freevars[i]
  else:
@@ -857,7 +868,7 @@ def _build_closure(annotate, owner, is_class, stringifier_dict, *, allow_evaluat
  name,
  cell=cell,
  owner=owner,
- globals=annotate.__globals__,
+ globals=_get_annotate_attr(annotate, "__globals__",{}),
  is_class=is_class,
  stringifier_dict=stringifier_dict,
  )
@@ -879,6 +890,141 @@ def _stringify_single(anno):
  return repr(anno)
 
 
+def _get_annotate_attr(annotate, attr, default=_sentinel):
+ # Try to get the attr on the annotate function. If it doesn't exist, we might
+ # need to look in other places on the object. If all of those fail, we can
+ # return the default at the end.
+ if hasattr(annotate, attr):
+ return getattr(annotate, attr)
+
+ # Redirect method attribute access to the underlying function. The C code
+ # verifies that the __func__ attribute is some kind of callable, so we need
+ # to look for attributes recursively.
+ if isinstance(annotate, types.MethodType):
+ return _get_annotate_attr(annotate.__func__, attr, default)
+
+ # Python generics are callable. Usually, the __init__ method sets attributes.
+ # However, typing._BaseGenericAlias overrides the __init__ method, so we need
+ # to use the original class method for fake globals and the like.
+ # _BaseGenericAlias also override __call__, so let's handle this earlier than
+ # other class construction.
+ if (
+ (typing := sys.modules.get("typing", None))
+ and isinstance(annotate, typing._BaseGenericAlias)
+ ):
+ return _get_annotate_attr(annotate.__origin__.__init__, attr, default)
+
+ # If annotate is a class instance, its __call__ is the relevant function.
+ # However, __call__ Could be a method, a function descriptor, or any other callable.
+ # Normal functions have a __call__ property which is a useless method wrapper,
+ # ignore these.
+ if (
+ (call := getattr(annotate, "__call__", None)) and
+ not isinstance(call, types.MethodWrapperType)
+ ):
+ return _get_annotate_attr(annotate.__call__, attr, default)
+
+ # Classes and generics are callable. Usually the __init__ method sets attributes,
+ # so let's access this method for fake globals and the like.
+ # Technically __init__ can be any callable object, so we recurse.
+ if isinstance(annotate, type) or isinstance(annotate, types.GenericAlias):
+ return _get_annotate_attr(annotate.__init__, attr, default)
+
+ # Most 'wrapped' functions, including functools.cache and staticmethod, need us
+ # to manually, recursively unwrap. For partial.update_wrapper functions, the
+ # attribute is accessible on the function itself, so we never get this far.
+ if hasattr(annotate, "__wrapped__"):
+ return _get_annotate_attr(annotate.__wrapped__, attr, default)
+
+ # Partial functions and methods both store their underlying function as a
+ # func attribute. They can wrap any callable, so we need to recursively unwrap.
+ if (
+ (functools := sys.modules.get("functools", None))
+ and isinstance(annotate, functools.partial)
+ ):
+ return _get_annotate_attr(annotate.func, attr, default)
+
+ if default is _sentinel:
+ raise TypeError(f"annotate function missing{attr!r} attribute")
+ return default
+
+def _direct_call_annotate(func, annotate, *args):
+ # If annotate is a method, we need to pass self as the first param.
+ if (
+ hasattr(annotate, "__func__") and
+ (self := getattr(annotate, "__self__", None))
+ ):
+ # We don't know what type of callable will be in the __func__ attribute,
+ # so let's try again with knowledge of that type, including self as the first
+ # argument.
+ return _direct_call_annotate(func, annotate.__func__, self, *args)
+
+ # Python generics (typing._BaseGenericAlias) override __call__, so let's handle
+ # them earlier than other class construction.
+ if (
+ (typing := sys.modules.get("typing", None))
+ and isinstance(annotate, typing._BaseGenericAlias)
+ ):
+ inst = annotate.__new__(annotate.__origin__)
+ func(inst, *args)
+ # Try to set the original class on the instance, if possible.
+ # This is the same logic used in typing for custom generics.
+ try:
+ inst.__orig_class__ = annotate
+ except Exception:
+ pass
+ return inst
+
+ # If annotate is a class instance, its __call__ is the function.
+ # __call__ Could be a method, a function descriptor, or any other callable.
+ # Normal functions have a __call__ property which is a useless method wrapper,
+ # ignore these.
+ if (
+ (call := getattr(annotate, "__call__", None)) and
+ not isinstance(call, types.MethodWrapperType)
+ ):
+ return _direct_call_annotate(func, annotate.__call__, *args)
+
+ # If annotate is a class, `func` is the __init__ method, so we still need to call
+ # __new__() to create the instance
+ if isinstance(annotate, type):
+ inst = annotate.__new__(annotate)
+ # func might refer to some non-function object.
+ _direct_call_annotate(func, annotate.__init__, inst, *args)
+ return inst
+
+ # Generic instantiation is slightly different. Since we want to give
+ # __call__ priority, the custom logic for builtin generics is here.
+ if isinstance(annotate, types.GenericAlias):
+ inst = annotate.__new__(annotate.__origin__)
+ # func might refer to some non-function object.
+ _direct_call_annotate(func, annotate.__init__, inst, *args)
+ # Try to set the original class on the instance, if possible.
+ # This is the same logic used in typing for custom generics.
+ try:
+ inst.__orig_class__ = annotate
+ except Exception:
+ pass
+ return inst
+
+ if functools := sys.modules.get("functools", None):
+ # If annotate is a partial function, re-create it with the new function object.
+ # We could call the function directly, but then we'd have to handle placeholders,
+ # and this way should be more robust for future changes.
+ if isinstance(annotate, functools.partial):
+ # Partial methods
+ if self := getattr(annotate, "__self__", None):
+ return functools.partial(func, self, *annotate.args, **annotate.keywords)(*args)
+ return functools.partial(func, *annotate.args, **annotate.keywords)(*args)
+
+ # If annotate is a cached function, we've now updated the function data, so
+ # let's not use the old cache. Furthermore, we're about to call the function
+ # and never use it again, so let's not bother trying to cache it.
+
+ # Or, if it's a normal function or unsupported callable, we should just call it.
+ return func(*args)
+
+
 def get_annotate_from_class_namespace(obj):
  """Retrieve the annotate function from a class namespace dictionary.