Type hinting a class that wraps function with __call__ method that conditionally modifies returned type

[soneill-pure]

I'm trying to type hint a class that essentially wraps a function and adds some additional method(s). The original function can return any one of the InitialOutputFrame types, while the class wrapping the function will have a __call__ method that outputs a subset of those types according to a type mapping (type_mapping). Anyone have any ideas how type hint this properly? If it seems like I'm doing something noobish it's probably because I am.

from abc import ABC, abstractmethod
from collections.abc import Callable, Generator, Iterable
from typing import Generic, TypeVar

import pandas
import polars

InitialOutputFrame = TypeVar(
    "InitialOutputFrame",
    polars.DataFrame,
    Iterable[polars.DataFrame],  # maps to polars.LazyFrame
    polars.LazyFrame,
    pandas.DataFrame,
    Iterable[pandas.DataFrame],  # maps to pandas.DataFrame
)
OutputFrame = TypeVar("OutputFrame", polars.DataFrame, polars.LazyFrame, pandas.DataFrame)

type_mapping: dict[type, type] = {
    polars.DataFrame: polars.DataFrame,
    Iterable[polars.DataFrame]: polars.LazyFrame,
    polars.LazyFrame: polars.LazyFrame,
    pandas.DataFrame: pandas.DataFrame,
    Iterable[pandas.DataFrame]: pandas.DataFrame,
}


class AbstractCachedFunction(ABC, Generic[InitialOutputFrame, OutputFrame]):
    """Abstract base class for cached functions."""

    @abstractmethod
    def __init__(self, func: Callable[..., InitialOutputFrame]): ...
    # TODO: how to map correct OutputFrame type to each InitialOutputFrame type?
    @abstractmethod
    def __call__(self, *args, show_progress_bar: bool = True, **kwargs) -> OutputFrame: ...
         # call `func` here and apply additional caching logic that leads to conditionally-modified return type
    @abstractmethod
    def get_path(self, *args, show_progress_bar: bool = True, **kwargs) -> Path:
        """Get the cached file path for the given arguments."""

[hauntsaninja]

You could try using overloads on __call__ where the overloads match on generic self

[soneill-pure]

Thanks for the tip. Based on that tip, this seems to works for me

class AbstractCachedFunction(ABC, Generic[InitialOutputFrame]):
    """Abstract base class for cached functions."""

    @abstractmethod
    def __init__(self, func: Callable[..., InitialOutputFrame]): ...

    # map correct OutputFrame type to each InitialOutputFrame type
    @overload
    @abstractmethod
    def __call__(
        self: "AbstractCachedFunction[OutputFrame]",
        *args,
        show_progress_bar: bool = True,
        **kwargs,
    ) -> OutputFrame: ...
    @overload
    @abstractmethod
    def __call__(
        self: "AbstractCachedFunction[Iterable[polars.DataFrame]]",
        *args,
        show_progress_bar: bool = True,
        **kwargs,
    ) -> polars.LazyFrame: ...
    @overload
    def __call__(
        self: "AbstractCachedFunction[Iterable[pandas.DataFrame]]",
        *args,
        show_progress_bar: bool = True,
        **kwargs,
    ) -> pandas.DataFrame: ...

    @abstractmethod
    def __call__(
        self,
        *args,
        show_progress_bar: bool = True,
        **kwargs,
    ) -> polars.DataFrame | polars.LazyFrame | pandas.DataFrame: ...

    @abstractmethod
    def get_path(
        self,
        *args,
        show_progress_bar: bool = True,
        **kwargs,
    ) -> Path:
        """Get the cached file path for the given arguments."""