Specifying a type to be a List of numbers (ints and/or floats)?

How do I specific a function can take a list of numbers which can be ints or floats?

I tried making a new type using Union like so:

num = Union[int, float]

def quick_sort(arr: List[num]) -> List[num]:
    ...

However, mypy didn't like this:

 quickSortLomutoFirst.py:32: error: Argument 1 to "quickSortOuter" has
 incompatible type List[int]; expected List[Union[int, float]]  

Is there a Type that encompasses ints and floats?

Solution 1:

The short answer to your question is you should use either TypeVars or Sequence -- using List[Union[int, float]] would actually potentially introduce a bug into your code!

In short, the problem is that Lists are invariant according to the PEP 484 type system (and in many other typesystems -- e.g. Java, C#...). You're attempting to use that list as if it were covariant instead. You can learn more about covariance and invariance here and here, but perhaps an example of why your code is potentially un-typesafe might be useful.

Consider the following code:

from typing import Union, List

Num = Union[int, float]

def quick_sort(arr: List[Num]) -> List[Num]:
    arr.append(3.14)  # We deliberately append a float
    return arr

foo = [1, 2, 3, 4]  # type: List[int]

quick_sort(foo)

# Danger!!!
# Previously, `foo` was of type List[int], but now
# it contains a float!? 

If this code were permitted to typecheck, we just broke our code! Any code that relies on foo being of exactly type List[int] would now break.

Or more precisely, even though int is a legitimate subtype of Union[int, float], that doesn't mean that List[int] is a subtype of List[Union[int, float]], or vice versa.

If we're ok with this behavior (we're ok with quick_sort deciding to inject arbitrary ints or floats into the input array), the fix is to manually annotate foo with List[Union[int, float]]:

foo = [1, 2, 3, 4]  # type: List[Union[int, float]]

# Or, in Python 3.6+
foo: List[Union[int, float]] = [1, 2, 3, 4]

That is, declare up-front that foo, despite only containing ints, is also meant to contain floats as well. This prevents us from incorrectly using the list after quick_sort is called, sidestepping the issue altogether.

In some contexts, this may be what you want to do. For this method though, probably not.

If we're not ok with this behavior, and want quick_sort to preserve whatever types were originally in the list, two solutions come to mind:

The first is to use a covariant type instead of list -- for example, Sequence:

from typing import Union, Sequence

Num = Union[int, float]

def quick_sort(arr: Sequence[Num]) -> Sequence[Num]:
    return arr

It turns out Sequence is more or less like List, except that it's immutable (or more precisely, Sequence's API doesn't contain any way of letting you mutate the list). This lets us safely sidestep the bug we had up above.

The second solution is to type your array more precisely, and insist that it must contain either all ints or all floats, disallowing a mixture of the two. We can do so using TypeVars with value restrictions:

from typing import Union, List, TypeVar 

# Note: The informal convention is to prefix all typevars with
# either 'T' or '_T' -- so 'TNum' or '_TNum'.
TNum = TypeVar('TNum', int, float)

def quick_sort(arr: List[TNum]) -> List[TNum]:
    return arr

foo = [1, 2, 3, 4]  # type: List[int]
quick_sort(foo)

bar = [1.0, 2.0, 3.0, 4.0]  # type: List[float]
quick_sort(foo)

This will also prevent us from accidentally "mixing" types like we had up above.

I would recommend using the second approach -- it's a bit more precise, and will prevent you from losing information about the exact type a list contains as you pass it through your quicksort function.

Solution 2:

From PEP 484, which proposed type hints:

Rather than requiring that users write import numbers and then use numbers.Float etc., this PEP proposes a straightforward shortcut that is almost as effective: when an argument is annotated as having type float, an argument of type int is acceptable...

Don't bother with the Unions. Just stick to Sequence[float].

Edit: Thanks to Michael for catching the difference between List and Sequence.