Python, MyPy: how can I do an "exclusive OR" of types in a generic's restrictions?

Viewed 107

I want to use generics so that I can avoid the "invariant Xyz" problem, and have a single source of truth for a couple of large type structure implementations, without having to Union the internals of a generic class and the externals as well. My code has a heavily nested series of aliases for a custom data structure, which is basically a big JSON of dicts and lists where the "leaves" are of one of two possible types (call them A and B), and only ever "A exclusively", or "B exclusively" for a given tree. However, some substructures of the overall data model (like List[A] or List[B]) can happen in different layers of the dictionary tree, and I don't want to have a to maintain a specific type for each layer, when their logic is exactly the same.

I'll be using the invariant List as an example. I want to have an alias C = List[A | B] recognize List[A] and List[B] as valid subtypes, without having to do C = List[A] | List[B] everywhere this happens, since Unions would need discrimination to function properly in my utility functions. I want A and B to be replaced by a GenericAorB when I can, to auto-infer if I have a List[A] or List[B] in my utilities.

I took inspiration from here, since it seemed like what I needed: https://mypy.readthedocs.io/en/stable/generics.html#type-variables-with-value-restriction

However it doesn't seem to work.

from typing import TypeVar, List

GenericStringOrFloat = TypeVar("GenericStringOrFloat", str, float)
ListStringOrFloat    = List[GenericStringOrFloat]
my_list : ListStringOrFloat
my_list = [12.0, 13.0]       # Fine
my_list = ["hello", "world"] # Fine
my_list = ["hello", 13.0]    # Also fine... what ? shouldn't mypy return an error for this ?

I don't understand why the last line doesn't return an error. To me it looks like

GenericStringOrFloat = TypeVar("GenericStringOrFloat", str, float)

is synonymous to

GenericStringOrFloat = TypeVar("GenericStringOrFloat", str | float)

which is thus completely pointless. Is this maybe a MyPy bug ? the documentation looks like this syntax is precisely meant to avoid this issue, and have the generic not be a dumb Union.

I want an exclusive OR, not inclusive, between the types given to my generic type restrictions. Phrased mathematically (though not rigorously), I want a sort of functor which can do something like ListUnionMapper : (A | B) -> List[A] | List[B] (and for Dicts as well).

Any idea how to do that ?

=========

EDIT:

After some more experimentation, I found an almost workable solution, but it's still far from perfect. The TLDR is that you need to give the generic variable a second time to allow mypy to do the automatic inference/discrimination; but there is still a bunch of weirdness with generic type restrictions.

When declaring a variable, give it the specific type being used (ie, str or float), and when declaring a function signature, give it the generic type variable.

The below code exemplifies what I mean by this.

from typing import TypeVar, List

GenericStringOrFloat  = TypeVar("GenericStringOrFloat", str, float)
ListFloat             = List[float]
ListString            = List[str]
ListStringOrFloat     = List[GenericStringOrFloat]
ListStringOrListFloat = ListString | ListFloat
my_list1A : ListStringOrFloat[      float] = [12.0,    13.0   ]  # Good as expected
my_list1B : ListStringOrFloat[      float] = ["hello", "world"]  # Error as expected: not float
my_list2A : ListStringOrFloat[str        ] = ["hello", "world"]  # Good as expected
my_list2B : ListStringOrFloat[str        ] = [12.0,    13.0   ]  # Error as expected: not str
my_list3  : ListStringOrFloat[str | float] = ["hello", 13.0   ]  # Good as expected
my_list4A : ListStringOrListFloat          = ["hello", "world"]  # Good as expected
my_list4B : ListStringOrListFloat          = [12.0,    13.0   ]  # Good as expected
my_list4C : ListStringOrListFloat          = ["hello", 13.0   ]  # Error as expected: not List[str] or List[float]
my_list4B                                  = ["hello", "world"]  # NB: doesn't infer/discriminate between List[str] and List[float], since union, so can be reassigned
my_list5A  : ListStringOrFloat[int]        = ["hello", 13.0   ]  # Error as expected: neither List[str} nor List[float]
my_list5B  : ListStringOrFloat[int]        = [1,       2      ]  # UNEXPECTED GOOD: this means that the type restriction basically doesn't work...

my_list_extra : ListStringOrFloat[GenericStringOrFloat] = ["hello", "world"]  # Error: Type variable GenericStringOrFloat is unbound


def return_first_element(lst: ListStringOrFloat[GenericStringOrFloat]) -> GenericStringOrFloat:
    assert len(lst) > 0
    return lst[0]


my_float_A : float = return_first_element(my_list1A)  # Good as expected
my_float_B : str   = return_first_element(my_list1A)  # Error: correctly infers that the function's return type is float
my_str_A   : float = return_first_element(my_list2A)  # Error: correctly infers that the function's return type is str
my_str_B   : str   = return_first_element(my_list2A)  # Good as expected
my_int_A   : float = return_first_element(my_list5B)  # Error: expects List[float] (and not List[Generic_Value]...)
my_int_B   : str   = return_first_element(my_list5B)  # Error: expression has type float + expects List[float] (and not List[Generic_Value]...)
my_int_C   : int   = return_first_element(my_list5B)  # Error: expression has type float + expects List[float] (and not List[Generic_Value]...)
1 Answers

The problem isn't in the TypeVar, but how mypy interprets my_list : ListStringOrFloat. Since it has no element type specified, it uses object. For example, your code also passes with ListStringOrFloat[object], though I agree that appears to contradict the meaning of GenericStringOrFloat. Generally speaking, solitary TypeVars don't really make sense, so I'm not too surprised it does something weird.

If you're only working with List[A] | List[B], you could alias ListX = List[A] | List[B] then use ListX[str, float], which does raise an error for ["hello", 13.0].

Related