Typing a decorator factory

Here’s an MRE:

import functools
from collections.abc import Callable
from typing import Literal, reveal_type

type Decorator[**P, T] = Callable[[Callable[P, T]], Callable[P, T]]

def decorator_factory[**P, T](text: str, method: Literal["upper", "lower"] = "upper") -> Decorator[P, T]:
    def decorator(func: Callable[P, T]) -> Callable[P, T]:
        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            match method:
                case "upper":
                    print(text.upper())
                case "lower":
                    print(text.lower())
            return func(*args, **kwargs)
        return wrapper
    return decorator

@decorator_factory("hello", method = "upper")
def foo(a: int, b: str) -> bool:
    return bool(a) and bool(b)

reveal_type(foo)

Type checkers disagree on this.

mypy doesn’t like it:

main.py:20: error: Argument 1 has incompatible type "Callable[[int, str], bool]"; expected "Callable[[VarArg(Never), KwArg(Never)], Never]"  [arg-type]
main.py:23: note: Revealed type is "def (*Never, **Never) -> Never"

pyright likes it:

Type of "foo" is "(a: int, b: str) -> bool"

pyrefly doesn’t like it:

INFO 23:12-17: revealed type: (ParamSpec(Unknown)) -> Unknown [reveal-type]

ty says @todo so I’m leaving it out.

This is a mypy bug apparently: Making a type alias of a Callable that describes a decorator discards all type information · Issue #18842 · python/mypy · GitHub.

Your mistake is binding **P and T to the decorator factory. You need to bind them to the inner decorator instead, making it polymorphic. To annotate the decorator factory, you can use a polymorphic Callback-Protocol, i.e. instead of

type Decorator[**P, T] = Callable[[Callable[P, T]], Callable[P, T]]

def decorator_factory[**P, T](...) -> Decorator[P, T]:
    def decorator(func: Callable[P, T]) -> Callable[P, T]: ...

you need

class Decorator(Protocol):  # <-- notice the lack of bound T-vars
    def __call__[**P, T](self, fn: Callable[P, T], /) -> Callable[P, T]: ...

def decorator_factory(...) -> Decorator:
    def decorator[**P, T](func: Callable[P, T]) -> Callable[P, T]: ...

Because you want your factory to return universal decorators.

mypy-playground

5 Likes

(+1 to Randolf. I wouldn’t consider this a bug in mypy and pyrefly. PEP 695 makes scoping explicit, but doesn’t have a convenient way to scope a type variable to just the return type. If you use old-style type variables and inline the alias, mypy will make its best guess about the intended scope and accept it)

Mypy implements some special-case handling for callable return types that have type variables that are bound to the def statement. Here’s another example that doesn’t involve ParamSpecs.

def decorator_factory[T]() -> Callable[[Callable[[T], None]], Callable[[T], None]]:
    def decorator(func: Callable[[T], None]) -> Callable[[T], None]: ...
    return decorator

@decorator_factory()
def foo(a: int) -> None: ...

# mypy and pyright reveal (int) -> None
# pyrefly reveals (Unknown) -> None
reveal_type(foo)

As Shantanu said, mypy includes some heuristics to guess the intent in cases where a Callable return type is used. Based on these heuristics, it sometimes re-scopes unsolved type variables to the Callable return type.

I implemented similar heuristics in pyright because I found that many libraries depend on mypy’s behavior here. There is no spec or documentation for these heuristics, and pyright’s implementation apparently differs in small ways from mypy’s. For example, pyright doesn’t depend on whether a type alias is used, but apparently mypy does. That explains the difference you’re seeing here.

I agree with Randolf that the correct way to handle this is through the use of a callback protocol.

I think that mypy’s and pyright’s behavior here is arguably non-conformant with the spec, although it’s a gray area because constraint solving behaviors are largely unspecified in the spec currently.

Unfortunately, it will probably be painful to change this behavior since many existing libraries and code bases depend upon it. We’d probably need to implement some way to ease that transition, such as backward compatibility flags.

It looks like pyrefly has not implemented any such heuristic for this case. It will be interesting to see whether they get user feedback that leads them to add such a heuristic.

I was also curious what ty does in this case, but it looks like its constraint solver isn’t far enough along to handle any of these code samples.

2 Likes