How to write a function that accepts either a TypeVar or a list of that TypeVar?

I’ve got a function where an argument is a dict[str, int | list[int]], and I want to replace int with a type variable. In other words, dict holds non-sequence objects or lists of those objects, but all basal elements are of the same type.

V = TypeVar("V")  # if we had type negation, would be bound to !list
def baz(val: Union[list[V], V]) -> V:
    if isinstance(val, list):
        return val[0]
    return val

baz(1)
baz([1])

mypy & pyright both point out that, in baz([1]),

debugme.py:7: note: Revealed type is "Union[builtins.list[V`-1], builtins.list[Any]]"
debugme.py:12: error: Need type annotation for "out2"  [var-annotated]
debugme.py:12: error: Argument 1 to "baz" has incompatible type "list[int]"; expected "list[Never]"  [arg-type]

I assume what’s going on is that, in baz([1]), the type checker chooses the second element of the Union, solving V as list. If I bound/constrain V, the errors clear, but that leaves the function less generic. Is there another way to get V to be solved as int, such as something like an ordered Union?

Also confusing about this:

  • Why is mypy expecting list[Never]? Wouldn’t that also be an error, since Never would get returned?
  • Am I using the term “solve” correctly? Is there a better term to use?

The issue here is that the signature is ambiguous, if you pass in a list[int] should it bind as whole to T or list[T], both are valid answers, because if it binds T to list[int] then the other possible type to pass in would be list[list[int]].

The most reliable way to help the type checker out in this case is to use overloads instead of an union, generally you should avoid generic unions in function signatures (outside of the return type), since it introduces ambiguities.

@overload
def baz(val: list[V]) -> V: ...
@overload
def baz(val: V) -> V: ...
1 Like

Thank you, so that’s what @overload is for! So how does the type checker know which overload to use? Without constraints on V, a variable of type list[int] could be val in either of those overloads. Does the type checker care about order here?

And is “bind” the right word where I was using “solve” above?

overload is one of the areas that are currently underspecified, so you are correct that in this case the two overloads are still overlapping, so it could pick either one. Different type checkers apply different heuristics to overload resolution. Order does matter, but it’s not the only factor, the real rules are much more complicated.

2 Likes

As @Daverball mentioned, there are multiple ways that a constraint solver could solve for V in this example. However, there’s almost always multiple valid answers that satisfy constraints. The typing spec doesn’t provide any specific guidance to type checkers on how to solve type variables.

Heuristics in pyright’s constraint solver attempt to produce the “simplest” example. You can read more about it here. That’s why pyright is able to produce the answer most people would expect in your example above.

2 Likes

Thanks guys! Really useful links. I also ran into the type widening difference b/w mypy and pyright in the same block of code.