Typing a decorator with optional params

typed_decorator.py

from typing import Callable, Optional, TypeVar, overload, Union
from typing_extensions import ParamSpec, Literal
from functools import wraps

P = ParamSpec("P")
R = TypeVar("R")


@overload
def decorator(func: Callable[P, R], param: int = ...) -> Callable[P, R]:
    ...


@overload
def decorator(
    func: Literal[None] = None, param: int = ...
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    ...


def decorator(
    func: Optional[Callable[P, R]] = None, param: int = 1
) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]:
    """A decorator that takes optional parameters"""

    def _deco(_func: Callable[P, R]) -> Callable[P, R]:
        """The actual decorator, when we know _func is not None"""

        @wraps(_func)
        def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            """The function wrapper."""
            return _func(*args, **kwargs)

        return _wrapper

    return _deco if func is None else _deco(func)


@decorator
def f1(x: int, y: str) -> str:
    """docstring preserved by functools.wraps"""


@decorator(param=2)
def f2(x: int, y: str) -> str:
    """docstring preserved by functools.wraps"""


# reveal_type(f1)  # def (x: builtins.int, y: builtins.str) -> builtins.str*
# reveal_type(f2)  # def (x: builtins.int, y: builtins.str) -> builtins.int*