Type hinting when multiple types are in input and output but in 1-to-1 relation between them

Question

I have the following code:

 AnyTypeEnum = TypeVar("AnyTypeEnum", Type[Enum], Type[IntEnum])
 AnyEnum = TypeVar("AnyEnum", Enum, IntEnum)

 def decode(value: str, enum_class: AnyTypeEnum) -> AnyEnum:
     return make_map(enum_class)[value]

 @functools.cache()
 def make_map(enum_class: AnyTypeEnum) -> Dict[str, AnyEnum]:
     return {x.name: x for x in enum_class}

If I pass an IntEnum class, I'll have in output IntEnum values, same with Enum. The mix from a class to the values of the other one is never possible.

Yet, mypy reports the erros:

1.Line 9: dictionary comprehension has incompatible type Enum, expected IntEnum

2.Line 9: dictionary comprehension has incompatible type IntEnum, expected Enum

I'm using TypeVar to solve the problem of "type mixing" between input and output that is present if I use Union instead, but apparently it's not completely fixed.

How do I solve this?

Answer 1

You have nothing binding your output and input types. At the moment they're entirely unrelated. Try this

AnyEnum = TypeVar("AnyEnum", Enum, IntEnum)

def decode(value: str, enum_class: Type[AnyEnum]) -> AnyEnum:
    return make_map(enum_class)[value]

@functools.cache()
def make_map(enum_class: Type[AnyEnum]) -> Dict[str, AnyEnum]:
    return {x.name: x for x in enum_class}

Alternatively, @overload should work.