gh-89263: Add typing.get_overloads

Doc/library/functools.rst

@@ -535,6 +535,8 @@ The :mod:`functools` module defines the following functions:
    .. versionchanged:: 3.7
       The :func:`register` attribute now supports using type annotations.
 
+   .. versionchanged:: 3.11
+      Implementation functions are now registered using :func:`register_variant`.
 
 .. class:: singledispatchmethod(func)
 
@@ -587,6 +589,9 @@ The :mod:`functools` module defines the following functions:
 
    .. versionadded:: 3.8
 
+   .. versionchanged:: 3.11
+      Implementation functions are now registered using :func:`register_variant`.
+
 
 .. function:: update_wrapper(wrapper, wrapped, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES)
 
@@ -664,6 +669,43 @@ The :mod:`functools` module defines the following functions:
    would have been ``'wrapper'``, and the docstring of the original :func:`example`
    would have been lost.
 
+.. function:: get_variants(key)
+
+   Return all registered function variants for this key. Function variants are
+   objects that represent some subset of the functionality of a function, for
+   example overloads decorated with :func:`typing.overload` or :func:`singledispatch`
+   implementation functions.
+
+   Variants are registered by calling :func:`register_variant`.
+   The *key* argument is a string that uniquely identifies the function and its
+   variants. It should be the result of a call to :func:`get_key_for_callable`.
+
+   .. versionadded:: 3.11
+
+.. function:: register_variant(key, variant)
+
+   Register a function variant that can later be retrieved using
+   :func:`get_variants`. The key should be the result of a call to
+   :func:`get_key_for_callable`.
+
+   .. versionadded:: 3.11
+
+.. function:: clear_variants(key=None)
+
+   Clear all registered variants with the given *key*. If *key* is None, clear
+   all variants.
+
+   .. versionadded:: 3.11
+
+.. function:: get_key_for_callable(func)
+
+   Return a string key that can be used with :func:`get_variants` and
+   :func:`register_variant`. *func* must be a :class:`function`,
+   :class:`classmethod`, :class:`staticmethod`, or similar callable.
+   If no key can be computed, the function returns None.
+
+   .. versionadded:: 3.11
+
 
 .. _partial-objects:
 

Doc/library/typing.rst

@@ -2110,6 +2110,9 @@ Functions and decorators
 
    See :pep:`484` for details and comparison with other typing semantics.
 
+   .. versionchanged:: 3.11
+      Overloaded functions are now registered using :func:`functools.register_variant`.
+
 .. decorator:: final
 
    A decorator to indicate to type checkers that the decorated method

Lib/functools.py

@@ -653,6 +653,52 @@ def cache(user_function, /):
     return lru_cache(maxsize=None)(user_function)
 
 
+################################################################################
+### Function variant registry
+################################################################################
+
+# {key: [variant]}
+_variant_registry = {}
+
+
+def register_variant(func, variant):
+    """Register a function variant."""
+    key = _get_key_for_callable(func)
+    _variant_registry.setdefault(key, []).append(variant)
+
+
+def get_variants(func):
+    """Get all function variants for the given function."""
+    key = _get_key_for_callable(func)
+    return _variant_registry.get(key, [])
+
+
+def clear_variants(func=None):
+    """Clear all variants for the given function (or all functions)."""
+    if func is None:
+        _variant_registry.clear()
+    else:
+        key = _get_key_for_callable(func)
+        _variant_registry.pop(key, None)
+
+
+def _get_key_for_callable(func):
+    """Return a key for the given callable.
+
+    This key can be used to register the callable in the variant registry
+    with register_variant() or to get variants for this callable with get_variants().
+
+    If no key can be created (because the object is not of a supported type), raise
+    AttributeError.
+    """
+    # classmethod and staticmethod
+    func = getattr(func, "__func__", func)
+    try:
+        return f"{func.__module__}.{func.__qualname__}"
+    except AttributeError:
+        return None
+
+
 ################################################################################
 ### singledispatch() - single-dispatch generic function decorator
 ################################################################################
@@ -814,6 +860,7 @@ def singledispatch(func):
     registry = {}
     dispatch_cache = weakref.WeakKeyDictionary()
     cache_token = None
+    outer_func = func
 
     def dispatch(cls):
         """generic_func.dispatch(cls) -> <function implementation>
@@ -890,6 +937,11 @@ def register(cls, func=None):
                         f"{cls!r} is not a class."
                     )
 
+        try:
+            register_variant(outer_func, func)
+        except AttributeError:
+            pass
+
         if _is_union_type(cls):
             from typing import get_args
 

Lib/test/test_functools.py

@@ -1890,6 +1890,70 @@ def cached_staticmeth(x, y):
         return 3 * x + y
 
 
+class MethodHolder:
+    @classmethod
+    def clsmethod(cls): ...
+    @staticmethod
+    def stmethod(): ...
+    def method(self): ...
+
+
+class TestVariantRegistry(unittest.TestCase):
+    def test_get_key_for_callable(self):
+        self.assertEqual(functools._get_key_for_callable(len),
+                         "builtins.len")
+        self.assertEqual(functools._get_key_for_callable(py_cached_func),
+                         f"{__name__}.py_cached_func")
+        self.assertEqual(functools._get_key_for_callable(MethodHolder.clsmethod),
+                         f"{__name__}.MethodHolder.clsmethod")
+        self.assertEqual(functools._get_key_for_callable(MethodHolder.stmethod),
+                         f"{__name__}.MethodHolder.stmethod")
+        self.assertEqual(functools._get_key_for_callable(MethodHolder.method),
+                         f"{__name__}.MethodHolder.method")
+
+    def test_get_variants(self):
+        def func1(): pass
+        def func2(): pass
+        obj1 = object()
+        obj2 = object()
+        self.assertEqual(functools.get_variants(func1), [])
+        self.assertEqual(functools.get_variants(func2), [])
+
+        functools.register_variant(func1, obj1)
+        self.assertEqual(functools.get_variants(func1), [obj1])
+        self.assertEqual(functools.get_variants(func2), [])
+
+        functools.register_variant(func1, obj2)
+        self.assertEqual(functools.get_variants(func1), [obj1, obj2])
+        self.assertEqual(functools.get_variants(func2), [])
+
+    def test_clear_variants(self):
+        def func1(): pass
+        def func2(): pass
+        obj1 = object()
+
+        functools.register_variant(func1, obj1)
+        self.assertEqual(functools.get_variants(func1), [obj1])
+        self.assertEqual(functools.get_variants(func2), [])
+
+        functools.clear_variants(func2)
+        self.assertEqual(functools.get_variants(func1), [obj1])
+        self.assertEqual(functools.get_variants(func2), [])
+
+        functools.clear_variants(func1)
+        self.assertEqual(functools.get_variants(func1), [])
+        self.assertEqual(functools.get_variants(func2), [])
+
+        functools.register_variant(func1, obj1)
+        functools.register_variant(func2, obj1)
+        self.assertEqual(functools.get_variants(func1), [obj1])
+        self.assertEqual(functools.get_variants(func2), [obj1])
+
+        functools.clear_variants()
+        self.assertEqual(functools.get_variants(func1), [])
+        self.assertEqual(functools.get_variants(func2), [])
+
+
 class TestSingleDispatch(unittest.TestCase):
     def test_simple_overloads(self):
         @functools.singledispatch

Lib/test/test_typing.py

@@ -1,6 +1,7 @@
 import contextlib
 import collections
 from functools import lru_cache
+import functools
 import inspect
 import pickle
 import re
@@ -9,7 +10,7 @@
 from unittest import TestCase, main, skipUnless, skip
 from copy import copy, deepcopy
 
-from typing import Any, NoReturn, Never, assert_never
+from typing import Any, NoReturn, Never, assert_never, overload
 from typing import TypeVar, TypeVarTuple, Unpack, AnyStr
 from typing import T, KT, VT  # Not in __all__.
 from typing import Union, Optional, Literal
@@ -3773,6 +3774,37 @@ def blah():
 
         blah()
 
+    def test_get_variants(self):
+        def blah():
+            pass
+
+        overload1 = blah
+        overload(blah)
+
+        def blah():
+            pass
+
+        overload2 = blah
+        overload(blah)
+
+        def blah():
+            pass
+
+        self.assertEqual(functools.get_variants(blah), [overload1, overload2])
+
+    def test_get_variants_repeated(self):
+        for _ in range(2):
+            def blah():
+                pass
+
+            overload_func = blah
+            overload(blah)
+
+            def blah():
+                pass
+
+            self.assertEqual(functools.get_variants(blah), [overload_func])
+
 
 # Definitions needed for features introduced in Python 3.6
 

Lib/typing.py

@@ -2382,10 +2382,40 @@ def utf8(value: bytes) -> bytes: ...
       def utf8(value: str) -> bytes: ...
       def utf8(value):
           # implementation goes here
+
+    Each overload is registered with functools.register_variant and can be
+    retrieved using functools.get_variants.
     """
+    try:
+        existing = functools.get_variants(func)
+    except AttributeError:
+        # Not a normal function; ignore.
+        pass
+    else:
+        if existing:
+            # If we are registering a variant with a lineno below or equal to that of the
+            # most recent existing variant, we're probably re-creating overloads for a
+            # function that already exists. In that case, we clear the existing variants
+            # to avoid leaking memory.
+            firstlineno = _get_firstlineno(func)
+            if firstlineno is not None:
+                existing_lineno = _get_firstlineno(existing[-1])
+                if existing_lineno is not None and firstlineno <= existing_lineno:
+                    functools.clear_variants(func)
+
+        functools.register_variant(func, func)
     return _overload_dummy
 
 
+def _get_firstlineno(func):
+    # staticmethod, classmethod
+    if hasattr(func, "__func__"):
+        func = func.__func__
+    if not hasattr(func, '__code__'):
+        return None
+    return func.__code__.co_firstlineno
+
+
 def final(f):
     """A decorator to indicate final methods and final classes.
 

Misc/NEWS.d/next/Library/2022-03-06-18-15-32.bpo-45100.B_lHu0.rst

@@ -0,0 +1,2 @@
+Add a mechanism to register function variants, such as overloads and
+singledispatch implementation functions. Patch by Jelle Zijlstra.