How do I annotate a function that is simultaneously generic for both a collection type and a member object type at the same time?

[gandhis1]

Code is below. It's kind of a jumbled mess because I was trying different things, and I don't even know if it actually runs. But what I'm trying to emphasize is the polymorphism embedded in the decorator, in which you can specify both an object type to serialize to and a collection type to containerize the objects in.

For instance, the errors indicate that you can't use a TypeVar in the bound of a TypeVar. But it's not clear to me what the correct way is.

from collections.abc import Mapping, Sequence
from functools import wraps
from typing import Any, Callable, Hashable, Optional, Type, TypeVar, Union, overload
from typing_extensions import ParamSpec
from dataclasses import dataclass

import pandas as pd


P = ParamSpec("P")
ObjectT = TypeVar("ObjectT")
SequenceT = TypeVar("SequenceT", bound=Sequence[ObjectT])
HashableT = TypeVar("HashableT", bound=Hashable)
MappingT = TypeVar("MappingT", bound=Mapping[HashableT, ObjectT])

@overload
def df_to_types(
    class_type: Type[ObjectT],
    container_type: Type[SequenceT],
) -> Callable[[Callable[P, pd.DataFrame]], Callable[P, SequenceT[ObjectT]]]:
    ...

@overload
def df_to_types(
    class_type: Type[ObjectT],
    container_type: Type[MappingT],
    key: Callable[[ObjectT], HashableT],
) -> Callable[[Callable[P, pd.DataFrame]], Callable[P, MappingT[HashableT, ObjectT]]]:
    ...

# Ignore the types here as the intention is for the overloads to take precedence
def df_to_types(
    class_type: Type[ObjectT],
    container_type: Union[Type[SequenceT], Type[MappingT]],
    key: Optional[Callable[[ObjectT], Hashable]] = None
) -> Any:
    def inner_decorator(func: Callable[P, pd.DataFrame]) -> Any:
        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
            df = func(*args, **kwargs)
            non_null = df.where(pd.notnull, None).to_dict(orient="records")
            obj_iter = map(lambda kwargs: class_type(**kwargs), non_null)
            if isinstance(container_type, Sequence):
                return container_type(obj_iter)  # Allow any sequence, not just a list
            elif isinstance(container_type, Mapping):  # Allow OrderedDict or any mapping
                return container_type((key(obj), obj) for obj in obj_iter)
            else:
                raise ValueError(f"{container_type} is not supported")
        return wrapper
    return inner_decorator


@dataclass
class Car:
    make: str
    model: str
    engine: str
    weight: float

def get_cars_df(a: int, b: str) -> pd.DataFrame:
    # Ignore the parameters for the sake of example simplicity
    return pd.DataFrame([("Honda", "Accord", "I4", 500), ("Toyota", "RAV4", "V6", 600)], columns=["make", "model", "engine", "weight"])

get_cars_list = df_to_types(Car, list)(get_cars_df)
get_cars_dict = df_to_types(Car, dict, key=lambda x: x.make + x.model)(get_cars_df)

cars_dict = get_cars_dict(1, "a")
reveal_type(cars_dict)  # This should be a Dict[str, Car], but it is Dict[Any, Any]

cars_list = get_cars_list(1, "a")
reveal_type(cars_list)  # This should be a List[Car], but it is List[Any]

After doing some research it seems that I may need higher-kinded types to do this: #548

Is that right?