bpo-32380 Fix singledispatch interaction with methods

Doc/library/functools.rst

@@ -331,7 +331,7 @@ The :mod:`functools` module defines the following functions:
      >>> fun_num is fun
      False
 
-   When called, the generic function dispatches on the type of the first
+   When called, the generic function by default dispatches on the type of the first
    argument::
 
      >>> fun("Hello, world.")
@@ -357,6 +357,23 @@ The :mod:`functools` module defines the following functions:
    for the base ``object`` type, which means it is used if no better
    implementation is found.
 
+   You can also tell ``@singeldispatch`` to dispatch on another argument index
+   with the ``arg`` keyword. This is especially useful for methods::
+
+    class Negater:
+        """neg(a) negates a for bools and ints"""
+        @singledispatch(arg=1)
+        def neg(self, a):
+            raise NotImplementedError("Cannot negate a")
+
+        @neg.register(int)
+        def _(self, a):
+            return -a
+
+        @neg.register(bool)
+        def _(self, a):
+            return not a
+
    To check which implementation will the generic function choose for
    a given type, use the ``dispatch()`` attribute::
 

Lib/functools.py

@@ -744,7 +744,7 @@ def _find_impl(cls, registry):
             match = t
     return registry.get(match)
 
-def singledispatch(func):
+def singledispatch(func=None, arg=0):
     """Single-dispatch generic function decorator.
 
     Transforms a function into a generic function, which can have different
@@ -753,76 +753,78 @@ def singledispatch(func):
     implementations can be registered using the register() attribute of the
     generic function.
     """
-    # There are many programs that use functools without singledispatch, so we
-    # trade-off making singledispatch marginally slower for the benefit of
-    # making start-up of such applications slightly faster.
-    import types, weakref
-
-    registry = {}
-    dispatch_cache = weakref.WeakKeyDictionary()
-    cache_token = None
-
-    def dispatch(cls):
-        """generic_func.dispatch(cls) -> <function implementation>
-
-        Runs the dispatch algorithm to return the best available implementation
-        for the given *cls* registered on *generic_func*.
-
-        """
-        nonlocal cache_token
-        if cache_token is not None:
-            current_token = get_cache_token()
-            if cache_token != current_token:
-                dispatch_cache.clear()
-                cache_token = current_token
-        try:
-            impl = dispatch_cache[cls]
-        except KeyError:
+    if func is None:
+        return partial(singledispatch, arg=arg)
+    else:
+        # There are many programs that use functools without singledispatch, so we
+        # trade-off making singledispatch marginally slower for the benefit of
+        # making start-up of such applications slightly faster.
+        import types, weakref
+        registry = {}
+        dispatch_cache = weakref.WeakKeyDictionary()
+        cache_token = None
+
+        def dispatch(cls):
+            """generic_func.dispatch(cls) -> <function implementation>
+
+            Runs the dispatch algorithm to return the best available implementation
+            for the given *cls* registered on *generic_func*.
+
+            """
+            nonlocal cache_token
+            if cache_token is not None:
+                current_token = get_cache_token()
+                if cache_token != current_token:
+                    dispatch_cache.clear()
+                    cache_token = current_token
             try:
-                impl = registry[cls]
+                impl = dispatch_cache[cls]
             except KeyError:
-                impl = _find_impl(cls, registry)
-            dispatch_cache[cls] = impl
-        return impl
-
-    def register(cls, func=None):
-        """generic_func.register(cls, func) -> func
-
-        Registers a new implementation for the given *cls* on a *generic_func*.
-
-        """
-        nonlocal cache_token
-        if func is None:
-            if isinstance(cls, type):
-                return lambda f: register(cls, f)
-            ann = getattr(cls, '__annotations__', {})
-            if not ann:
-                raise TypeError(
-                    f"Invalid first argument to `register()`: {cls!r}. "
-                    f"Use either `@register(some_class)` or plain `@register` "
-                    f"on an annotated function."
+                try:
+                    impl = registry[cls]
+                except KeyError:
+                    impl = _find_impl(cls, registry)
+                dispatch_cache[cls] = impl
+            return impl
+
+        def register(cls, func=None):
+            """generic_func.register(cls, func) -> func
+
+            Registers a new implementation for the given *cls* on a *generic_func*.
+
+            """
+            nonlocal cache_token
+            if func is None:
+                if isinstance(cls, type):
+                    return lambda f: register(cls, f)
+                ann = getattr(cls, '__annotations__', {})
+                if not ann:
+                    raise TypeError(
+                        f"Invalid first argument to `register()`: {cls!r}. "
+                        f"Use either `@register(some_class)` or plain `@register` "
+                        f"on an annotated function."
+                    )
+                func = cls
+
+                # only import typing if annotation parsing is necessary
+                from typing import get_type_hints
+                argname, cls = next(iter(get_type_hints(func).items()))
+                assert isinstance(cls, type), (
+                    f"Invalid annotation for {argname!r}. {cls!r} is not a class."
                 )
-            func = cls
-
-            # only import typing if annotation parsing is necessary
-            from typing import get_type_hints
-            argname, cls = next(iter(get_type_hints(func).items()))
-            assert isinstance(cls, type), (
-                f"Invalid annotation for {argname!r}. {cls!r} is not a class."
-            )
-        registry[cls] = func
-        if cache_token is None and hasattr(cls, '__abstractmethods__'):
-            cache_token = get_cache_token()
-        dispatch_cache.clear()
-        return func
-
-    def wrapper(*args, **kw):
-        return dispatch(args[0].__class__)(*args, **kw)
-
-    registry[object] = func
-    wrapper.register = register
-    wrapper.dispatch = dispatch
-    wrapper.registry = types.MappingProxyType(registry)
-    wrapper._clear_cache = dispatch_cache.clear
-    update_wrapper(wrapper, func)
-    return wrapper
+            registry[cls] = func
+            if cache_token is None and hasattr(cls, '__abstractmethods__'):
+                cache_token = get_cache_token()
+            dispatch_cache.clear()
+            return func
+
+        def wrapper(*args, **kw):
+            return dispatch(args[arg].__class__)(*args, **kw)
+
+        registry[object] = func
+        wrapper.register = register
+        wrapper.dispatch = dispatch
+        wrapper.registry = types.MappingProxyType(registry)
+        wrapper._clear_cache = dispatch_cache.clear
+        update_wrapper(wrapper, func)
+        return wrapper

Lib/test/test_functools.py

@@ -2147,6 +2147,43 @@ def __eq__(self, other):
                 return self.arg == other
         self.assertEqual(i("str"), "str")
 
+    def test_method_register(self):
+        class A:
+            @functools.singledispatch(arg=1)
+            def t(self, arg):
+                return "base"
+            @t.register(int)
+            def _(self, arg):
+                return "int"
+            @t.register(str)
+            def _(self, arg):
+                return "str"
+        a = A()
+
+        self.assertEqual(a.t(0), "int")
+        self.assertEqual(a.t(''), "str")
+        self.assertEqual(a.t(0.0), "base")
+
+        self.assertEqual(A.t(a,0), "int")
+        self.assertEqual(A.t(a,''), "str")
+        self.assertEqual(A.t(a,0.0), "base")
+
+    def test_method_register_noargs(self):
+        msg = "tuple index out of range"
+        class B:
+            @functools.singledispatch(arg=1)
+            def fun():
+                return 'base'
+            @fun.register(str)
+            def _():
+                return "str"
+            @fun.register(float)
+            def _():
+                return "float"
+        with self.assertRaises(IndexError) as exc:
+            B.fun()
+        self.assertEqual(str(exc.exception), msg)
+
     def test_invalid_registrations(self):
         msg_prefix = "Invalid first argument to `register()`: "
         msg_suffix = (

Misc/ACKS

@@ -1484,6 +1484,7 @@ Václav Šmilauer
 Allen W. Smith
 Christopher Smith
 Eric V. Smith
+Ethan Smith
 Gregory P. Smith
 Mark Smith
 Roy Smith

Misc/NEWS.d/next/Library/2017-12-22-22-31-46.bpo-32380.vHdQW7.rst

@@ -0,0 +1,2 @@
+``functools.singledispatch`` now supports decorating and registering with
+methods.