Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 31 additions & 11 deletions Lib/functools.py
Original file line number Diff line number Diff line change
Expand Up @@ -935,6 +935,32 @@ def _is_valid_dispatch_type(cls):
return (isinstance(cls, UnionType) and
all(isinstance(arg, type) for arg in cls.__args__))

def _get_func_type_hints(func):
"""Called when type hints are needed to choose the first argument to dispatch on."""
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

Expand All @@ -951,20 +977,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_func_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}. "
Expand Down
12 changes: 12 additions & 0 deletions Lib/test/test_functools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
A :exc:`TypeError` is raised by :py:func:`functools.singledispatch`
if it is attempted to register function that only annotates its return type.

Contributed by Bartosz Sławecki in :gh:`84644`.
Loading