Expression type Union[Tuple[size1], Tuple[size2]] can't be cast to Tuple[size1]

[laszukdawid]

Describe the bug
Not sure if this is a bug or just looking for guidance how to deal with it. There are two tickets (#90 & #103) that are similar, yet potentially lower abstraction.

Depending on the arguments, a function intends to return a Tuple with either 2 or 3 elements. In the example below these are of custom type (Pandas DataFrame or Series, although a similar situation happens with str). Regardless whether I include the test_split it will return an error on assigning to a new variable that Tuple size mismatch: expected x but got y (where (x, y) are 2 or 3 and x != y for all permutations of test_split, direct output declaration, changing type...).

Changing the expected output to Tuple[Data, ...] works, i.e. no error is thrown, however I'd like to make clear that it's either going to be a tuple with size 2 or 3.

from typing import Union, Tuple
import pandas as pd

Data = Union[pd.Series, pd.DataFrame]
DateSplit = Tuple[str, str]

def split_data(ts: Data, train_split: DateSplit, validate_split: DateSplit, test_split: Optional[DateSplit]=None) -> Union[Tuple[Data, Data], Tuple[Data, Data, Data]]:
    train: Data = ts[(ts.index >= train_split[0]) & (ts.index < train_split[1])]
    validate: Data = ts[(ts.index >= validate_split[0]) & (ts.index < validate_split[1])]

    if test_split is not None:
        test: Data = ts[(ts.index >= test_split[0]) & (ts.index < test_split[1])]
        return (train, validate, test)

    return (train, validate)

Screenshots

VS Code extension or command-line
I'm running pyright v1.1.34 on VS Code v1.44.2 on Ubuntu 16.04 in a custom environment.

[erictraut]

This isn't specific to Tuples. A type checker will do the same thing for any Union type. Consider the following example:

def func(a: Optional[int] = None) -> Union[str, int]:
    if a is None:
        return "Nada"
    return a

foo1: int = func()  # error because func() might return a str
foo2: str = func() # error because func() might return an int

foo3 = func()
if isinstance(foo3, str):
    foo4: str = foo3   # this is fine because foo3 is definitely a str
else:
    foo5: int = foo3   # this is fine because foo3 is definitely an int

Unfortunately, there's no way to use isinstance to distinguish between two different tuples because they're the same type.

Here are a few potential workarounds.

First, if you can distinguish the return type based on input parameter types, you could use the @overload decorator to create two overloads for your split_data function. Here's how that would work with my simplified example above:

@overload
def func(a: None = None) -> str:
    raise NotImplementedError()

@overload
def func(a: int) -> int:
    raise NotImplementedError()

def func(a: Optional[int] = None) -> Union[str, int]:
    if a is None:
        return "Nada"
    return a

Now the type checker knows whether func returns a str or int based on the call arguments passed to it.

A second option is to use data classes (or subclasses of NamedTuple) rather than anonymous Tuple types.

class TrainValidateTest(NamedTuple):
    train: Data
    validate: Data
    test: Data

class TrainValidate(NamedTuple):
    train: Data
    validate: Data

Because these are separate classes, you can now use the isinstance trick I showed above.

[laszukdawid]

Thank you, yet again, for a quick response. Impressive timing :) (and maybe just a tad concerning)

Both approaches worked and I understand why my attempt failed. Ticket resolved.