Decorator with an optional parameter

decorator.py

"""Decorator with an optional parameter.

See also https://mypy.readthedocs.io/en/stable/generics.html#decorator-factories
"""

from collections.abc import Callable
import functools
from typing import Any, overload

FUNCS = {}


@overload
def register[F: Callable[..., Any]](func: F, /) -> F:
    """Return the decorated function when used as bare decorator."""

@overload
def register[F: Callable[..., Any]](*, name: str | None = ...) -> Callable[[F], F]:
    """Return a function-returning decorator when used as decorator factory."""

def register(func: Callable[..., Any] | None = None, /, *,
             name: str | None = None):
    """Register the decorated function under its name or the given one.""" 
    if func is None:
        return functools.partial(register, name=name)

    if name is None:
        name = func.__name__
    FUNCS[name] = func
    return func


@register
def spam():
    pass


@register(name='eggs')
def _eggs():
    pass


@register()
def bacon():
    pass


assert FUNCS == {'spam': spam, 'eggs': _eggs, 'bacon': bacon}