Mypy: Typing two list of int or str to be added together

Question

I have a function that merge two list of either int or str. There cannot be a situation when the two list are being of different types.

it is defined by the following code:

AddableList = ...


def add_arrays(array: AddableList, array2: AddableList) -> AddableList:
    if len(array) != len(array2):
        raise ValueError

    return [a + b for a, b in zip(array, array2)]

When typing AddableList, using List[int] mypy: Success: no issues

When typing AddableList, using List[str] mypy: Success: no issues

However, mypy will return the following error

error: Unsupported operand types for + ("int" and "str")
error: Unsupported operand types for + ("str" and "int")
note: Both left and right operands are unions
Found 2 errors in 1 file (checked 333 source files)

when using typing correctly the lists with AddableList = List[Union[int, str]]

Finally, when trying to type AddableList to Union[List[int], List[str]], mypy errors with:

error: Unsupported left operand type for + ("object")

What typing should i use to resolve this issue?

Answer 1

Use a TypeVar that will resolve to one type or the other (but not both at once within the same context):

from typing import List, TypeVar

Addable = TypeVar("Addable", str, int)


def add_arrays(array: List[Addable], array2: List[Addable]) -> List[Addable]:
    if len(array) != len(array2):
        raise ValueError

    return [a + b for a, b in zip(array, array2)]

reveal_type(add_arrays([1, 2, 3], [2, 3, 4]))    # List[int]
reveal_type(add_arrays(["a", "b"], ["c", "d"]))  # List[str]
reveal_type(add_arrays([1, 2, 3], ["a", "b", "c"]))  # error

The final line throws an error because there's no type that Addable can resolve to:

test.py:14: error: Cannot infer type argument 1 of "add_arrays"