Don't single-dispatch on return types

Author: bswck

Lib/functools.py

@@ -935,6 +935,31 @@ def _is_valid_dispatch_type(cls):
         return (isinstance(cls, UnionType) and
                 all(isinstance(arg, type) for arg in cls.__args__))
 
+    def _get_type_hints(func):
+        ann = getattr(func, '__annotate__', None)
+        if ann is None:
+            raise TypeError(
+                f"Invalid first argument to `register()`: {func!r}. "
+                f"Use either `@register(some_class)` or plain `@register` "
+                f"on an annotated function."
+            )
+
+        # only import typing if annotation parsing is necessary
+        from typing import get_type_hints
+        from annotationlib import Format
+
+        type_hints = get_type_hints(func, format=Format.FORWARDREF)
+        type_hints.pop("return", None)  # don't dispatch on return types
+
+        if not type_hints:
+            raise TypeError(
+                f"Invalid first argument to `register()`: {func!r}. "
+                f"Use either `@register(some_class)` or plain `@register` "
+                f"on a function with annotated parameters."
+            )
+
+        return type_hints
+
     def register(cls, func=None):
         """generic_func.register(cls, func) -> func
 
@@ -951,20 +976,14 @@ def register(cls, func=None):
                     f"Invalid first argument to `register()`. "
                     f"{cls!r} is not a class or union type."
                 )
-            ann = getattr(cls, '__annotate__', None)
-            if ann is None:
-                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
+            type_hints = _get_type_hints(func)
+
+            argname, cls = next(iter(type_hints.items()))
 
-            # only import typing if annotation parsing is necessary
-            from typing import get_type_hints
-            from annotationlib import Format, ForwardRef
-            argname, cls = next(iter(get_type_hints(func, format=Format.FORWARDREF).items()))
             if not _is_valid_dispatch_type(cls):
+                from annotationlib import ForwardRef
+
                 if isinstance(cls, UnionType):
                     raise TypeError(
                         f"Invalid annotation for {argname!r}. "

Lib/test/test_functools.py

@@ -3180,6 +3180,18 @@ def _(arg):
         )
         self.assertEndsWith(str(exc.exception), msg_suffix)
 
+        with self.assertRaises(TypeError) as exc:
+            @i.register
+            def _(arg) -> str:
+                return "I only have a return type annotation"
+        self.assertStartsWith(str(exc.exception), msg_prefix +
+            "<function TestSingleDispatch.test_invalid_registrations.<locals>._"
+        )
+        self.assertEndsWith(str(exc.exception),
+            ". Use either `@register(some_class)` or plain `@register` on "
+            "a function with annotated parameters."
+        )
+
         with self.assertRaises(TypeError) as exc:
             @i.register
             def _(arg: typing.Iterable[str]):