Search code examples
pythondecoratormypytyping

Python 3.10 type hinting for decorator to be used in a method


I'm trying to use typing.Concatenate alongside typing.ParamSpec to type hint a decorator to be used by the methods of a class. The decorator simply receives flags and only runs if the class has that flag as a member. Code shown below:

import enum
from typing import Callable, ParamSpec, Concatenate

P = ParamSpec("P")
Wrappable = Callable[Concatenate["Foo", P], None]

class Flag(enum.Enum):
    FLAG_1 = enum.auto()
    FLAG_2 = enum.auto()


def requires_flags(*flags: Flag) -> Callable[[Wrappable], Wrappable]:
    def wrap(func: Wrappable) -> Wrappable:
        def wrapped_f(foo: "Foo", *args: P.args, **kwargs: P.kwargs) -> None:
            if set(flags).issubset(foo.flags):
                func(foo, *args, **kwargs)
        return wrapped_f
    return wrap


class Foo:
    def __init__(self, flags: set[Flag] | None = None) -> None:
        self.flags: set[Flag] = flags or set()
        super().__init__()

    @requires_flags(Flag.FLAG_1)
    def some_conditional_method(self, some_int: int):
        print(f"Number given: {some_int}")


Foo({Flag.FLAG_1}).some_conditional_method(1)  # prints "Number given: 1"
Foo({Flag.FLAG_2}).some_conditional_method(2)  # does not print anything

The point of using Concatenate here is that the first parameter of the decorated function must be an instance of Foo, which aligns with methods of Foo (for which the first parameter is self, an instance of Foo). The rest of the parameters of the decorated function can be anything at all, hence allowing *args and **kwargs

mypy is failing the above code with the following:

error: Argument 1 has incompatible type "Callable[[Foo, int], Any]"; expected "Callable[[Foo, VarArg(Any), KwArg(Any)], None]"  [arg-type]
note: This is likely because "some_conditional_method of Foo" has named arguments: "self". Consider marking them positional-only

It's having an issue with the fact that at the call site, I'm not explicitly passing in an instance of Foo as the first argument (as I'm calling it as a method). Is there a way that I can type this strictly and correctly? Does the wrapper itself need to be defined within the class somehow so that it has access to self directly?

Note that if line 5 is updated to Wrappable = Callable[P, None] then mypy passes, but this is not as strict as it could be, as I'm trying to enforce in the type that it can only be used on methods of Foo (or free functions which receive a Foo as their first parameter).

Similarly, if I update some_conditional_method to be a free function rather than a method on Foo, then mypy also passes (this aligns with the linked SO question below). In this case it is achieving the strictness that I'm after, but I really want to be able to apply this to methods, not just free functions (in fact, it doesn't need to apply to free functions at all).

This question is somewhat of an extension to Python 3 type hinting for decorator but has the nuanced difference of the decorator needing to be used in a method.

To be clear, the difference between this and that question is that the following (as described in the linked question) works perfectly:

@requires_flags(Flag.FLAG_1)
def some_conditional_free_function(foo: Foo, some_int: int):
    print(f"Number given: {some_int}")

some_conditional_free_function(Foo({Flag.FLAG_1}), 1)  # prints "Number given: 1"

Solution

  • You have two major issues here, and more detailed warnings would have been given if you had strict on:

    1. Your type alias, Wrappable = Callable[Concatenate["Foo", P], None], has a type variable (here, the ParamSpec P), but you're not providing the type variable when you're using the alias. This means you've lost all signature information after decorating with @requires_flags(...). You can fix this by explicitly parameterising Wrappable with P upon usage:

      def requires_flags(*flags: Flag) -> Callable[[Wrappable[P]], Wrappable[P]]:
          def wrap(func: Wrappable[P]) -> Wrappable[P]:
      

      This is caught with either strict mode or the mypy configuration setting disallow_any_generics = True.

    2. The arguments to Concatenate consist of any number of positional-only parameters followed by a ParamSpec. The declaration

      def wrapped_f(foo, *args, **kwargs): ...
      
      >>> wrapped_f(foo=Foo(), some_int=1)  # Works at runtime
      

      has now replaced some_conditional_method like so:

      class Foo:
          @requires_flags(Flag.FLAG_1)
          def some_conditional_method(self, some_int) -> None: ...
      
      >>> Foo.some_conditional_method(foo=Foo(), some_int=1)  # Works at runtime
      

      This isn't great, especially if some_conditional_method is part of the public API. The type checker is warning you to write wrapped_f properly, by using a positional-only marker:

      def wrapped_f(foo: "Foo", /, *args: P.args, **kwargs: P.kwargs) -> None: ...
      
      >>> Foo.some_conditional_method(foo=Foo(), some_int=1)  # Fails at runtime
      

      This is caught with either strict mode or the mypy configuration setting extra_checks = True.

    The remaining minor issue is that def some_conditional_method(self, some_int: int): isn't annotated with a return type; this is a source of typing bugs if you're going to be using inheritance. I suggest getting into the habit of writing -> None for functions which don't return anything.


    The fixed code can be re-run here: mypy-playground