CODEWITHSUNDEEP
pythontype-hintingmypy
Jim Hunziker
Jim Hunziker

Reputation: 15370

How can I type hint a variable whose Union type gets narrowed in Python?

I have a few helper functions that get passed a type converter and a value. Based on checks that happen later, I decide which helper function to call.

How can I correctly annotate the types to narrow the foo variable's type below so that it can pass a mypy check?

from typing import Type, Union


def do_something(
        typ: Type[Union[float, int]],
        bar: Union[float, int]
) -> Union[float, int]:
    return bar


foo: Type[Union[float, int, str]] = float

assert foo is float or foo is int

do_something(foo, 4.4)

Bonus points if the solution can ensure that typ is a converter to the type of bar!

Upvotes: 1

Views: 3367

Answers (1)

Josiah
Josiah

Reputation: 1364

The tool you want here is a TypeVar.

Essentially, a TypeVar let's you say "I don't know quite what type this is (although I may have some ideas), but it'll be the same one throughout its use in this function." (or in some cases throughout its use in a class)

For example, this ensures that each thing you had a Union for gets the same value on any given call to the function.

from typing import Type, TypeVar

# Define a type variable
# and list the things it is allowed to represent
NumberType = TypeVar("NumberType", int, float) 

def do_something(
        typ: Type[NumberType],
        bar: NumberType
) -> NumberType:
    return bar

This can be legally called with do_something(float, 2.5) in which case it will return a float, or it can be called with do_something(int, 2) in which case it will return an int. That is, it makes sure that all the things match.

Because you called it a type converter, I suspect that you may not actually want all the types to match. If you need to constrain on more than one Type Variable, you can do so with something more like

from typing import Callable, TypeVar

# Define a type variable
# and list the things it is allowed to represent
NumberTypeIn = TypeVar("NumberTypeIn", int, float)
NumberTypeOut = TypeVar("NumberTypeOut", int, float) 

def do_something(
        converter: Callable[[NumberTypeIn], NumberTypeOut],
        bar: NumberTypeIn
) -> NumberTypeOut:
    return type_(bar)

As to the original question of narrowing unions of Type[]s, as you've noticed is doesn't work. Instead you can use issubclass, as in

assert not issubclass(foo, str)

or

assert issubclass(foo, int) or issubclass(foo, float)

Upvotes: 3

Related Questions