How to use a type stored in a class variable as type hint for a parameter of a method of the same class?

Question

Consider the following code:

from typing import Protocol, ClassVar, Any


class Tool(Protocol):
    t: ClassVar

    # def f(self, params: self.t) -> Any:
    # NameError: name 'self' is not defined
    def f(self, params) -> Any:
        pass


class Sum:
    t: ClassVar = list[int]

    def f(self, params: list[int]) -> int:
        return sum(params)


def use_tool(tool: Tool, params: Any) -> Any:
    return tool.f(params)


s = Sum()
print(use_tool(s, [1, 2]))

I would like to define a protocol Tool such that classes like Sum could implement it with a specific t, which stores a type/class object, so that in methods such as Sum.f(), params would be constrained have the type of t in type checking. How may I do that?

I tried to use self.t to refer to t in type annotation, which does not work, nor does replacing self with typing.Self.

I think storing a type for the parameter for f() would be useful in the case that, suppose t stores a pydantic.BaseModel (which means that t is of type ModelMetaclass), then in functions such as use_tool, I can look at the schema of t via t.model_dump_json(). Also, use_tool() may be able to take a JSON string as an input, and parse it into the appropriate model via t.model_validate_json() before calling Tool.f.

Answer 1

To be used as a type hint, t must be a TypeAlias, but if so it must be initialized at the same time it is defined:

_{(playgrounds: Mypy, Pyright)}

class Tool(Protocol):
    t: TypeAlias = list[int]  # pyright => fine 
                              # mypy    => error: Type aliases are prohibited in protocol bodies
    def f(self, params: 'Tool.t') -> Any:  # both    => fine
        reveal_type(params)                # both    => list[int]

(Whether type aliases are allowed within Protocols is unclear; the specification says nothing about this.)

Instead, make Tool generic:

_{(playgrounds: Mypy, Pyright)}

class Tool[T](Protocol):
    t: type[T]

    def f(self, params: T) -> Any: ...

It can then be used as:

class Sum:
    # "type[list[int]]" is necessary for Mypy
    # to understand that it is not a class-level type alias
    t: type[list[int]] = list[int]

    def f(self, params: list[int]) -> int: ...

def use_tool[T](tool: Tool[T], params: T) -> Any: ...

s = Sum()
use_tool(s, [1, 2])  # fine

This has another advantage: Tool can also be generic over the return type of .f().

_{(playgrounds: Mypy, Pyright)}

class Tool[T, R](Protocol):
    t: type[T]

    def f(self, params: T) -> R: ...

...

def use_tool[T, R](tool: Tool[T, R], params: T) -> R: ...

s = Sum()
reveal_type(use_tool(s, [1, 2]))  # int

For the NameError part (which would not be a problem with this solution), see this question.