Implementation, tests, and docs

JelleZijlstra · Apr 16, 2022 · Mar 6, 2022 · Mar 7, 2022 · Mar 7, 2022 · Mar 8, 2022
commit 831b5650936ea83d3531e6f92bc9719161e8e875
diff --git a/Doc/library/functools.rst b/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:
 

@@ -2106,6 +2106,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

diff --git a/Lib/functools.py b/Lib/functools.py
@@ -661,7 +661,7 @@ def cache(user_function, /):
 _variant_registry = {}
 
 
-def register_variant(key, variant, *, test) -> None:
+def register_variant(key, variant):
     """Register a function variant."""
     _variant_registry.setdefault(key, []).append(variant)
 
@@ -671,6 +671,14 @@ def get_variants(key):
     return _variant_registry.get(key, [])
 
 
+def clear_variants(key=None):
+    """Clear all variants for the given key (or all keys)."""
+    if key is None:
+        _variant_registry.clear()
+    else:
+        _variant_registry.pop(key, None)
+
+
 def get_key_for_callable(func):
     """Return a key for the given callable.
 
@@ -679,6 +687,8 @@ def get_key_for_callable(func):
 
     If no key can be created, return None.
     """
+    # classmethod and staticmethod
+    func = getattr(func, "__func__", func)
     try:
         return f"{func.__module__}.{func.__qualname__}"
     except AttributeError:
@@ -922,6 +932,10 @@ def register(cls, func=None):
                         f"{cls!r} is not a class."
                     )
 
+        key = get_key_for_callable(func)
+        if key is not None:
+            register_variant(key, func)
+
         if _is_union_type(cls):
             from typing import get_args
 

diff --git a/Lib/test/test_functools.py b/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):
+        key1 = "key1"
+        key2 = "key2"
+        obj1 = object()
+        obj2 = object()
+        self.assertEqual(functools.get_variants(key1), [])
+        self.assertEqual(functools.get_variants(key2), [])
+
+        functools.register_variant(key1, obj1)
+        self.assertEqual(functools.get_variants(key1), [obj1])
+        self.assertEqual(functools.get_variants(key2), [])
+
+        functools.register_variant(key1, obj2)
+        self.assertEqual(functools.get_variants(key1), [obj1, obj2])
+        self.assertEqual(functools.get_variants(key2), [])
+
+    def test_clear_variants(self):
+        key1 = "key1"
+        key2 = "key2"
+        obj1 = object()
+
+        functools.register_variant(key1, obj1)
+        self.assertEqual(functools.get_variants(key1), [obj1])
+        self.assertEqual(functools.get_variants(key2), [])
+
+        functools.clear_variants(key2)
+        self.assertEqual(functools.get_variants(key1), [obj1])
+        self.assertEqual(functools.get_variants(key2), [])
+
+        functools.clear_variants(key1)
+        self.assertEqual(functools.get_variants(key1), [])
+        self.assertEqual(functools.get_variants(key2), [])
+
+        functools.register_variant(key1, obj1)
+        functools.register_variant(key2, obj1)
+        self.assertEqual(functools.get_variants(key1), [obj1])
+        self.assertEqual(functools.get_variants(key2), [obj1])
+
+        functools.clear_variants()
+        self.assertEqual(functools.get_variants(key1), [])
+        self.assertEqual(functools.get_variants(key2), [])
+
+
 class TestSingleDispatch(unittest.TestCase):
     def test_simple_overloads(self):
         @functools.singledispatch

@@ -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, AnyStr
 from typing import T, KT, VT  # Not in __all__.
 from typing import Union, Optional, Literal
@@ -2930,8 +2931,8 @@ def fun(x: a):
 
     def test_forward_repr(self):
         self.assertEqual(repr(List['int']), "typing.List[ForwardRef('int')]")
-        self.assertEqual(repr(List[ForwardRef('int', module='mod')]), 
-                         "typing.List[ForwardRef('int', module='mod')]")	
+        self.assertEqual(repr(List[ForwardRef('int', module='mod')]),
+                         "typing.List[ForwardRef('int', module='mod')]")
 
     def test_union_forward(self):
 
@@ -3303,6 +3304,39 @@ 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
+
+        key = functools.get_key_for_callable(blah)
+        self.assertEqual(functools.get_variants(key), [overload1, overload2])
+
+    def test_get_variants_repeated(self):
+        for _ in range(2):
+            def blah():
+                pass
+
+            overload_func = blah
+            overload(blah)
+
+            def blah():
+                pass
+
+            key = functools.get_key_for_callable(blah)
+            self.assertEqual(functools.get_variants(key), [overload_func])
+
 
 # Definitions needed for features introduced in Python 3.6
 

@@ -2183,13 +2183,38 @@ 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.
     """
     key = functools.get_key_for_callable(func)
     if key is not None:
+
+        # 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 = functools.get_variants(key)
+            if existing:
+                existing_lineno = _get_firstlineno(existing[-1])
+                if existing_lineno is not None and firstlineno <= existing_lineno:
+                    functools.clear_variants(key)
+
         functools.register_variant(key, 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.
 

diff --git a/Misc/NEWS.d/next/Library/2022-03-06-18-15-32.bpo-45100.B_lHu0.rst b/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.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Add a mechanism to register function variants, such as overloads and
		singledispatch implementation functions. Patch by Jelle Zijlstra.