How to cast a typing.Union to one of its subtypes in Python?

Question

I’m using Python 3.6.1, mypy, and the typing module. I created two custom types, Foo and Bar, and then used them in a dict I return from a function. The dict is described as mapping str to a Union of Foo and Bar. Then I want to use values from this dict in a function that names only one argument each:

from typing import Dict, Union, NewType

Foo = NewType("Foo", str)
Bar = NewType("Bar", int)

def get_data() -> Dict[str, Union[Foo, Bar]]:
    return {"foo": Foo("one"), "bar": Bar(2)}

def process(foo_value: Foo, bar_value: Bar) -> None:
    pass

d = get_data()

I tried using the values as-is:

process(d["foo"], d["bar"])
# typing-union.py:15: error: Argument 1 to "process" has incompatible type "Union[Foo, Bar]"; expected "Foo"
# typing-union.py:15: error: Argument 2 to "process" has incompatible type "Union[Foo, Bar]"; expected "Bar"

Or using the types:

process(Foo(d["foo"]), Bar(d["bar"]))
# typing-union.py:20: error: Argument 1 to "Foo" has incompatible type "Union[Foo, Bar]"; expected "str"
# typing-union.py:20: error: Argument 1 to "Bar" has incompatible type "Union[Foo, Bar]"; expected "int"

How do I cast the Union to one of its subtypes?

Answer 1

You'd have to use cast():

process(cast(Foo, d["foo"]), cast(Bar, d["bar"]))

From the Casts section of PEP 484:

Occasionally the type checker may need a different kind of hint: the programmer may know that an expression is of a more constrained type than a type checker may be able to infer.

There is no way to spell what specific types of value go with what specific value of a dictionary key. You may want to consider returning a named tuple instead, which can be typed per key:

from typing import Dict, Union, NewType, NamedTuple

Foo = NewType("Foo", str)
Bar = NewType("Bar", int)

class FooBarData(NamedTuple):
    foo: Foo
    bar: Bar

def get_data() -> FooBarData:
    return FooBarData(foo=Foo("one"), bar=Bar(2))

Now the type hinter knows exactly what each attribute type is:

d = get_data()
process(d.foo, d.bar)

Or you could use a dataclass:

from dataclasses import dataclass

@dataclass
class FooBarData:
    foo: Foo
    bar: Bar

which makes it easier to add optional attributes as well as control other behaviour (such as equality testing or ordering).

I prefer either over typing.TypedDict, which is more meant to be used with legacy codebases and (JSON) serialisations.

Answer 2

Although I think casts are probably the right option to use in your case, I just briefly want to mention one additional option which might be applicable in similar scenarios, to round things out:

It's actually possible to type your dict more precisely using the new, experimental TypedDict feature, which is available latest versions of mypy (if you clone from the github repo) and will likely be available in the next pypi release.

In order to use TypedDict, you'll need to install the mypy_extensions from pypi by running pip install mypy_extensions.

TypedDict lets you assign individual types to each item in your dict:

from mypy_extensions import TypedDict

Foo = NewType("Foo", str)
Bar = NewType("Bar", int)

FooBarData = TypedDict('FooBarData', {
    'foo': Foo,
    'bar': Bar,
})

You can also define FooBarData using a class-based syntax in Python 3.6+:

from mypy_extensions import TypedDict

Foo = NewType("Foo", str)
Bar = NewType("Bar", int)

class FooBarData(TypedDict):
    foo: Foo
    bar: Bar

You also mentioned that your dict can have a dynamic number of elements. If it genuinely is dynamic, then TypedDict won't help for the same reasons that NamedTuple won't help, but if your TypedDict will ultimately have a finite and number of elements, and you're just progressively adding items to it instead of all at once, you can try using non-total TypedDicts, or try constructing TypeDicts that mix required and non-required items.

It's also worth noting that unlike pretty much every other type, TypedDicts are checked using structural typing, rather then nominal typing. This means that if you define a completely unrelated TypedDict named, say, QuxData that also has foo and bar fields with the same type as FooBarData, then QuxData will actually be a valid subtype of FooBarData. This may open up some interesting possibilities with a little bit of cleverness.