Union and Optional types
We've seen that mypy says that this is an error:
var = 123
var = 'abc'
But what if you actually wanted to do this?
It might seem like you'd never want to do this, but consider this example:
chosen_number = "not found"
for number in [1, 5, 7, 4]:
if number % 2 == 0:
chosen_number = number
print("Chosen number:", chosen_number)
chosen_number
is going to end up as an integer in some cases, but if the list
doesn't have any even integers in it, it will stay as its default value of
"not found"
. And mypy is confused by this.
In this case, we intentionally want chosen_number
to be able to hold both the
string default, and the possible integer values. And to tell mypy that this is
our intention, we need to expand its knowledge on what the type of the variable
should be.
Due to the definition, it's assuming that chosen_number
should be of type
str
:
chosen_number = "not found"
reveal_type(chosen_number)
But we don't want that, we want the variable to be able to accept "either str
or int
values". And for that, you use the Union
type from the typing
module:
from typing import Union
chosen_number: Union[str, int] = "not found"
for number in [1, 5, 7, 4]:
if number % 2 == 0:
chosen_number = number
print("Chosen number:", chosen_number)
Now that mypy correctly knows that we want it to have either kind of values, it tells us that the rest of the code is good.
Note that Union
can take any number of types, not just two. For example, a set
that can contain strings, integers and floats would be typed as
set[Union[str, int, float]]
.
Some examples and caveats
Remember when we talked about a better way to make a collection type with multiple data types in it? The example we used was:
items = [1, 2, 3]
items.append(4)
items.append('world')
The old solution that we used was to use list[object]
, but that lends us into
this problem:
items = [1, 2, 3]
items.append(4)
items.append(False)
items.append({'this is': 'a dictionary'})
Yeah, with list[object]
you can suddenly append any object into the list. To
limit it to just ints and strings, just use a union:
from typing import Union
items: list[Union[int, str]] = [1, 2, 3]
items.append(4)
items.append('world')
# Now these don't work anymore!
items.append(False)
items.append({'this is': 'a dictionary'})
Cool right?
However, consider this:
from typing import Union
items: list[Union[int, str]] = [1, 2, 3]
items.append(4)
total = 0
for item in items:
total += item
print(total)
Although this code works, now mypy says we're doing something wrong. The error says:
Unsupported operand types for `+`: "int" and "str"
What this means is that in our addition operation, the left hand value is an
int
, and the right hand value is a str
:
...
total += item
Mypy is worried, that since we said items
can have strings in it, it's
possible that item
is of type str
. So it won't let you do this as it is code
that can possibly crash. To solve this, we need to do something called "type
narrowing".
Type narrowing
To be fair, there is a bug in our code. If we make the list this:
from typing import Union
items: list[Union[int, str]] = [1, 2, 3, "A string"]
items.append(4)
total = 0
for item in items:
total += item
print(total)
Python suddenly crashes with the same error that mypy was previously warning us about. And since we did tell mypy that we want to put strings in the list, it's completely understandable for mypy to give us this warning. Mypy always wants to make sure that our code doesn't crash at runtime.
So, how do we fix the code? Well, we could just check if each item is an int
before adding it to the total:
from typing import Union
items: list[Union[int, str]] = [1, 2, 3, "A string"]
items.append(4)
total = 0
for item in items:
if isinstance(item, int):
total += item
print(total)
... and now mypy correctly shows no errors! Great. Mypy understands that inside
the isinstance
block, item
can only be an integer. You can even confirm it
by doing reveal_type
:
from typing import Union
items = [1, 2, 3, "A string"]
items.append(4)
total = 0
for item in items:
reveal_type(item)
if isinstance(item, int):
reveal_type(item)
total += item
Here's what we learned: You can use isinstance
blocks that only run if the
data is of a certain type, to narrow down a union of types to the one that we
want. Here's another example:
from typing import Union
def buy(stuff: Union[str, list[str]]) -> None:
if isinstance(stuff, list):
# Buy every item in the list
for item in stuff:
print("Buying {item}")
else:
# Buy the one item
print("Buying {stuff}")
buy("Jam")
buy(["Milk", "Eggs", "Spam"])
Because we checked for list
s, Mypy understood that list[str]
is the only
possible type for stuff
. On the other hand, inside the else
block, stuff
is seen as a string:
from typing import Union
def buy(stuff: Union[str, list[str]]) -> None:
if isinstance(stuff, list):
reveal_type(stuff)
...
else:
reveal_type(stuff)
Optional
Take this function:
def pop_word(words):
"""Takes one word out of given words."""
if len(words) == 0:
print("Warning: empty list")
return None
word = words.pop()
return word
word = pop_word(['foo', 'bar'])
What should the type definition of the above code be?
If you answered:
def pop_word(words: list[str]) -> Union[str, None]:
...
Then you'd be correct! It's actually really common to see types that look like
Union[A, B, None]
, that the typing
module has a shortcut for it, called
Optional
. Any instance of Union[A, B, ..., None]
can be replaced with
Optional[A, B, ...]
.
For example, Union[int, None]
can be written as Optional[int]
, and
Union[str, bool, None]
as Optional[str, bool]
.
New syntax
Just like the use of list
instead of typing.List
, Union
has also become
unnecessary starting from Python 3.10. If you're going to run your code on 3.10
or higher, you can use this syntax instead:
items: list[int | str] = [1, 2, 3]
The |
operator can be used with as many types as you want to make the same
unions as you'd make with Union[...]
. int | str
means the same thing: Either
int
or str
.
Similarly, Optional[bool]
can be written as bool | None
, here's an example:
user = {
'name': 'Bryan',
'likes': ['parkour', 'dogs']
}
def likes_flowers() -> bool | None:
like_data = user.get('likes')
if like_data is None:
# We don't have likes data
return None
if 'flowers' in like_data:
return True
else:
return False
Again, I'll be using the new |
syntax for the rest of the course, for
simplicity's sake. Feel free to use either yourself.