Conditional Python Typing for Generic Functions

Question

Assuming I have a class A and a derived class B:

from typing import Type, TypeVar

class A:
    type_field: Type[int] = int

class B(A):
    type_field: Type[str] = str

I want to have a function

T = TypeVar('T')

def func(object: T, some_field): 
    ...

A want some_field to be annotated (e.g. to show type hints) with type of T.type_field, assuming T derives from A or A itself. In other words, I want call func(A(), ... to suggest me passing int next, while call func(B(), ... should suggest passing string. Is there any way to achieve it in Python 3?

I want to develop a proper typing for function def func so developers get a proper type hint depending on the type of object passed first.

Answer 1

This inheritance hierarchy is broken from the start. All instances of A are supposed to have a type_field attribute of static type type[int], but instances of B are instances of A that do not satisfy this property. mypy will already complain, even before you try to write your func:

from typing import Type, TypeVar

class A:
    type_field: Type[int] = int

class B(A):
    type_field: Type[str] = str

Output:

main.py:7: error: Incompatible types in assignment (expression has type "type[str]", base class "A" defined the type as "type[int]")  [assignment]
Found 1 error in 1 file (checked 1 source file)

That said, you can sort of write a function like what you want, with a potential caveat:

from typing import Protocol, TypeVar

T = TypeVar('T')

class FieldHaver(Protocol[T]):
    type_field: type[T]

def func(a: FieldHaver[T], b: type[T]):
    pass

The caveat is that type is covariant, so this will also accept any subtype of the type that type_field is annotated with. So these calls will work, like you want:

func(A(), int)
func(B(), str)

and this will fail (with a less helpful error message than you'd probably like):

func(B(), int)

but this will pass, which it sounds like you don't want:

func(A(), bool)  # bool is a subclass of int

I don't think there's any way around the covariance... although, if you don't want covariance, there's probably no point taking the second argument in the first place.