Type hinting function with two use cases

Question

I am trying to fully type hint a function that ensures that an element is in a given dictionary, and then checks that the element type is what the user expects it to be. My initial implementation works correctly, and is shown below

T = TypeVar("T")

def check_and_validate_element_in_dict(
    element_name: str, dictionary: Dict[str, Any], element_type: Type[T]
) -> T:
    assert element_name in dictionary
    element = dictionary[element_name]
    assert isinstance(element, element_type)
    return element

it allows me to replace this

assert "key1" in _dict
key1 = _dict["key1"]
assert isinstance(key1, type1)

assert "key2" in _dict
key2 = _dict["key2"]
assert isinstance(key2, type2)

with this

key1 = check_and_validate_element_in_dict("key1", _dict, type1)
key2 = check_and_validate_element_in_dict("key2", _dict, type2)

now, this only works if the element type to test is only one, like int, str, etc. I also want to be able to test multuple different types in my function, like

isinstance(element, (int, dict))
isinstance(element, (float, type(None)))

the issue here is type hinting the function in order to make it understand that if element_type is a single value T, the return value is T, but if element_type is one of e.g. two types T and U, the return value will be either T or U.

I guess it is possible, but since I'm still a newbie in the type hinting area I'll need some help!

Edit:

I tried making the function support either a single type or a tuple of two different types as base case, so I updated element_type to be

element_type: Union[Type[T], Tuple[Type[T], Type[T]]]

now the return element statement gets flagged by mypy with the error:

Returning Any from function declared to return "T"

this also raises a question: do I need to indicate each different input type as a new TypeVar? In such case, the element_type definition becomes

# using U = TypeVar("U")
def ...(..., element_type: Union[Type[T], Tuple[Type[T], Type[U]]]) -> Union[T, U]:

In this case the issues keeps being

Returning Any from function declared to return "T"

Answer 1

You can use typing.overload, which allows you to register multiple different signatures for one function. A function decorated with @overload is ignored by python at runtime, so you can leave the body of these functions empty by just putting an ellipsis ... in the body. These implementations are just for the type-checker — you have to make sure that there is at least one "true" implementation of the function that is not decorated with overload.

from typing import TypeVar, overload, Any, Union, Dict, Type, Tuple

t0 = TypeVar('t0')
t1 = TypeVar('t1')


@overload
def check_and_validate_element_in_dict(
        element_name: str,
        dictionary: Dict[str, Any],
        element_type: Type[t0]
) -> t0:
    """
    Signature of the function if exactly one type
    is supplied to argument element_type
    """
    ...


@overload
def check_and_validate_element_in_dict(
        element_name: str,
        dictionary: Dict[str, Any],
        element_type: Tuple[Type[t0], Type[t1]]
) -> Union[t0, t1]:
    """
    Signature of the function if a tuple of exactly two types
    is supplied to argument element_type
    """     
    ...


def check_and_validate_element_in_dict(
        element_name: str,
        dictionary: Dict[str, Any],
        element_type: Any
) -> Any:
    """Concrete implementation of the function"""
    
    assert element_name in dictionary
    element = dictionary[element_name]
    assert isinstance(element, element_type)
    return element

This feels like a deeply imperfect solution, however, as it doesn't provide a solution for a tuple of arbitrary length being passed to your element_type argument. It only works if you know the length of your tuple will be one of (for example) 2, 3 or 4 -- you can then provide an overloaded implementation for each of those situations. Would definitely be interested if anybody can think of a better solution.