I want to type a function that can take either a class type, or an instance of the class, and then return an instance of that same class. For example:
from typing import Type, TypeVar, Union
DT = TypeVar("DT")
def tmap(dest_type_or_obj: Union[DT, Type[DT]]) -> DT:
if isinstance(dest_type_or_obj, type):
dest_obj = dest_type_or_obj()
else:
dest_obj = dest_type_or_obj
return dest_obj
class Dest:
pass
instance_dest = tmap(Dest()) # Works fine
type_dest = tmap(Dest) # [arg-type] mypy(error)
# Argument 2 to "map" of "ObjectMapper" has incompatible type "Type[Dest]"; expected "Type[<nothing>]"
While technically equivalent, MyPy requires to separate both cases using overload
:
@overload
def tmap(dest_type_or_obj: Type[DT]) -> DT:
...
@overload
def tmap(dest_type_or_obj: DT) -> DT:
...
This enables MyPy to infer the type variable properly in both the instance and type case.
reveal_type(tmap(Dest())) # note: Revealed type is "__main__.Dest*"
reveal_type(tmap(Dest)) # note: Revealed type is "__main__.Dest*"