Pattern matching over nested `Union` types in Python

Question

Building a Python library, I am using type hints to guarantee consistency over certain data representation. In particular, I am making use of Union (sum types) in a nested fashion to represent the different "flavor" a datum can take.

What I end up with so far is similar to the following example:

from typing import Union

MyNumberT = Union[float,int]
MyDataT = Union[str,MyNumber]

def my_data_to_string(datum: MyDataT) -> str:
    if isinstance(datum, float):
        return _my_number_to_string(datum)
    elif isinstance(datum, int):
        return _my_number_to_string(datum)
    elif isinstance(datum, str):
        return datum
    # assert_never omitted for simplicity

def _my_number_to_string(number: MyNumberT) -> str:
    return "%s" % number

Which type-checks fine using mypy.

Now, my real code is a bit more complex, and I need to perform some common operations on variables that are of type MyNumberT. In the example, this is simply highlighted by adapting the import and replacing my_data_to_string as in the following:

from typing import get_args, Union

[...]

def my_data_to_string(datum: MyDataT) -> str:
    if isinstance(datum, get_args(MyNumberT)):
        return _my_number_to_string(datum)
    elif isinstance(datum, str):
        return datum
    # assert_never omitted for simplicity

[...]

On which the type-checking of mypy fails: Argument 1 to "_my_number_to_string" has incompatible type "Union[str, Union[float, int]]"; expected "Union[float, int]" .

I expected mypy to "realise" that in the first branch, datum could only be of type float or int, but the error message indicates it's not the case...

How could I achieve some pattern matching over "parts" of such nested types?

question from:https://stackoverflow.com/questions/65912706/pattern-matching-over-nested-union-types-in-python

Answer 1

Your use-case is a great example to use a utility provided by functools called singledispatch. It allows you to define multiple functionality to a single function based on type of input.

from functools import singledispatch

# This class defines the function with
# a base case if the input type doesn't match
@singledispatch
def my_data_to_string(datum) -> str:
    raise TypeError(f"unsupported format: {type(datum)}")

# Registering for type str using type hint
@my_data_to_string.register
def _(datum: str):
    return datum

# Registering for multiple 
# types using decorator
@my_data_to_string.register(float)
@my_data_to_string.register(int)
def _(datum):
    return "<%s>" % datum


print(my_data_to_string("a"))    # a
print(my_data_to_string(1))      # <1>
print(my_data_to_string(1.5))    # <1.5>
print(my_data_to_string([1, 2])) # TypeError

It is extensible, readable and doesn't generate error in linters/formatters. Docs link.