Search code examples
pythonpython-typing

Using TypeVarTuple with inner TypeVar (variadic generics transformation)


How would I use TypeVarTuple for this example?

T = TypeVar("T")
Ts = TypeVarTuple("Ts")

@dataclass
class S(Generic[T]):
    data: T

def data_from_s(*structs: ??) -> ??: 
    return tuple(x.data for x in structs)

a = data_from_s(S(1), S("3"))  # is type tuple[int, str]

Solution

  • I don't see any way to do this with the current spec. The main issue I see is that TypeVarTuple does not support bounds. You can't constrain the types referred to by Ts to be bounded to S.

    You need to translate somehow tuple[S[T1], S[T2], ...] -> tuple[T1, T2, ...], but you have no way to know that the types contained by Ts are specializations of S or that types themselves are generic with further parameterization.


    Without using TypeVarTuple, your goal can be accomplished to some extent with a pattern like the following, using overload to handle subsets of the signature for differing amounts of arguments. I also use an ending / in the overloads to prevent usage of named arguments (forcing positional args to be used), which allows the overloads to match the real method definition.

    Obviously, this pattern becomes awkward as you add more ranges of arguments, but in some cases it can be a nice escape hatch.

    from dataclasses import dataclass
    from typing import Any, Generic, TypeVar, assert_type, overload
    
    T = TypeVar("T")
    
    @dataclass
    class S(Generic[T]):
        data: T
    ...
    
    T1 = TypeVar("T1")
    T2 = TypeVar("T2")
    T3 = TypeVar("T3")
    
    @overload
    def data_from_s(s1: S[T1], /) -> tuple[T1]:
        ...
    
    @overload
    def data_from_s(s1: S[T1], s2: S[T2], /) -> tuple[T1, T2]:
        ...
    
    @overload
    def data_from_s(s1: S[T1], s2: S[T2], s3: S[T3], /) -> tuple[T1, T2, T3]:
        ...
    
    def data_from_s(*structs: S[Any]) -> tuple[Any, ...]:
        return tuple(x.data for x in structs)
    

    Which will pass this test:

    assert_type(
        data_from_s(S(1)),
        tuple[int]
    )
    assert_type(
        data_from_s(S(1), S("3")), 
        tuple[int, str]
    )
    assert_type(
        data_from_s(S(1), S("3"), S(3.9)), 
        tuple[int, str, float]
    )