Type hinting decorator with bounded arguments

Question

An example of what I’d like to do is as follows:

@dataclass
class Arg1:
    x = field()


@dataclass
class Arg2:
    x = field()


@dataclass
class Obj:
    y = field()

    T = TypeVar("T")
    R = TypeVar("R", bound=(Arg1 | Arg2))
    C = TypeVar("C", bound=Callable[[Self, R], T]) 

    @staticmethod
    def deco(f: C) -> C:
        def wrapper(self: Self, arg: Obj.R) -> Obj.T:
            print(arg.x)
            print(self.y)   
            return f(self, arg)

        return wrapper

    @deco
    def f1(self: Self, arg: Arg1) -> str:
        return "ok"  

    @deco
    def f2(self: Self, arg: Arg2) -> int:
        return 3

I’m getting the following two errors from mypy:

test.py:22: error: A function returning TypeVar should receive at least one argument containing the same TypeVar  [type-var]
test.py:26: error: Incompatible return value type (got "Callable[[Self, Obj.R], Obj.T]", expected "C")  [return-value]

If I remove T and replace uses of it with Any, I can fix the first error (at the expense of weakening my type assertions) but still get the second error. This one is particularly puzzling to me because mypy appears to be giving me back the definition of C in complaining that I’m not returning a C. What should I be doing to convince it of the type equivalence here?

Regarding the first error, I expect that I could solve it by unpacking the definition of C in def deco(f: C) -> C but then I would lose the restriction that the type of the argument being passed to f is the same as the type of the argument to the return value of deco. Is there anything I can do about this?

Edit: after discussion in comments I updated the definition of deco to be def deco(f: Callable[[Self, R], T]) -> Callable[[Self, R], T]: and now get the following error instead:

test.py:21: error: Static methods cannot use Self type  [misc]

This seems like progress but I’m not sure what to do about this either.

Asked By: danben

||

Source

Answer 1

I would move the decorator outside the class. It’s not really specific to your class, in that all the type information you need can be specified with other type variables.

This appears to be sufficient (at least with my mongrel setup using Python 3.9 and mypy 1.0.1; Self is provided by typing_extensions rather than the standard library typing module).

T = TypeVar("T")
S = TypeVar("S", bound='Obj')
R = TypeVar("R", Arg1, Arg2)  # constrained variable, not a bound variable

def deco(f: Callable[[S, R], T]) -> Callable[[S, R], T]:
    def wrapper(self: S, arg: R) -> T:
        print(arg.x)
        print(self.y)
        return f(self, arg)
    return wrapper

class Obj:
    @deco
    def f1(self: Self, arg: Arg1) -> str:
        return "ok"

    @deco
    def f2(self: Self, arg: Arg2) -> int:
        return 3

We make R a constrained variable so that we’ll know if arg is an instance of Arg1 or Arg2 in the body of the function, rather than an unknown instance of one of those two types.

Answered By: chepner

Answer 2

Here’s something that typechecks in mypy 1.1.1, I hope it’s what you need:

from dataclasses import dataclass, field
from typing import TypeVar, Callable

@dataclass
class Arg1:
    x = field()


@dataclass
class Arg2:
    x = field()

T = TypeVar("T")
R = TypeVar("R", Arg1, Arg2)
ObjSelf = TypeVar('ObjSelf', bound='Obj')

def deco(f: Callable[[ObjSelf, R], T]) -> Callable[[ObjSelf, R], T]:
    def wrapper(self: ObjSelf, arg: R) -> T:
        print(arg.x)
        print(self.y)
        return f(self, arg)

    return wrapper


@dataclass
class Obj:
    y = field()

    @deco
    def f1(self, arg: Arg1) -> str:
        return "ok"  

    @deco
    def f2(self, arg: Arg2) -> int:
        return 3

To explain: Self is a bit of magic that doesn’t play well with the kind of non-trivial value juggling you’re doing, so I reverted to doing what we had to do in the bad old days: a type variable with an explicit bound.

C had to go, because not every R is the same, just like not every T is the same, so we need to have them listed explicitly in the type signature of deco, so mypy knows what scope the type variable has. I pulled the type vars and deco outside of Obj because the complexity of it all gave me a headache.

Edit: chepner was right, a constrained typevar for R is better than a bound one in this case.

Answered By: Jasmijn

Type hinting decorator with bounded arguments

Question:

Answers: