Tips on Type Narrowing

[padraic-mcatee]

Hello,

I'm a bit new to typing in python (coming from scala). I tend to write long chained method calls for processing pandas dataframes. I am working with an API that often returns union types of dataframes and iterators of dataframes. However, I rarely need to handle the case of iterators and would prefer to assume return type of simple DataFrame. For example:

def api_func() -> pd.DataFrame | Iterator[pd.DateFrame] :
    ...

result = api_func() \
    .df_method_1() \ # error with "Cannot access member for type Iterator[pd.DataFrame]
    .df_method_2() \ # same as above
    ...
    .df_method_n() # same as above

Does anyone have a nice typeguard / cast pattern for resolving this? Some method on the Union type that would return one of its constituent types?

[JelleZijlstra]

I would suggest an exclude_iterator function:

def exclude_iterator(df_or_iterator: pd.DataFrame | Iterator[pd.DataFrame]) -> pd.DataFrame:
    if isinstance(df_or_iterator, pd.DataFrame):
        return df_or_iterator
    raise TypeError("single dataframe expected")

result = exclude_iterator(api_func()).some_method()

[padraic-mcatee]

Thanks for the reply. Is there maybe a solution without so much boilerplate? I suppose this is as simple as it gets.

[erictraut]

In many cases, functions contain @overload directives that distinguish return types based on input parameters. In those cases, you shouldn't need to do any further work because the return type will not be a union. However, overloads aren't possible for all functions.

I like the solution @JelleZijlstra suggested because it provides both static type safety and runtime validation.

[padraic-mcatee]

Ok thank you very much for the quick responses here. I don't think @overload will work here. I will build something along the lines of @JelleZijlstra's solution.