Annotate Callable input params from a generic tuple

[maciejzj]

Hi, I am trying to annotate a rather simple PyTorch Dataset wrapper for data transformations.

The snippet below should be the best explanation of what I want to achieve. Right now I get:

p.py:32: error: Argument 2 to "TransformableDatasetWrapper" has incompatible type "Callable[[str, str], tuple[int, int]]"; expected "Callable[[Any], tuple[int, int]]"  [arg-type]

from mypy which works as expected. How do I solve this and indicate that the t function should take expanded InT as args? I have already tried a lot with ParamSpec, Unpack, TypeVarTuple, etc. but failed every time (I don't include my messy non-working solutions to reduce the noise).

from typing import Sequence, Callable, TypeVar, Generic
from torch.utils.data import Dataset, IterableDataset, TensorDataset


class StrDs(Dataset[tuple[str, str]]):
    def __init__(self):
        self.d = [("1", "10"), ("2", "20"), ("3", "30")]

    def __getitem__(self, idx: int) -> tuple[str, str]:
        return self.d[idx]


InT = TypeVar("InT", bound=tuple)
OutT = TypeVar("OutT", bound=tuple)


class TransformableDatasetWrapper(Generic[InT, OutT], Dataset[InT]):
    def __init__(
        self, d: Dataset[InT], t: Callable[[*InT], OutT]   # <-- * is pseudocode to show my intentions
    ):
        self.d = d
        self.t = t

    def get(self, i: int) -> OutT:
        return self.t(*self.d[i])


def str_to_int(i: str, ii: str) -> tuple[int, int]:
    return int(i), int(ii)


TransformableDatasetWrapper[tuple[str, str], tuple[int, int]](StrDs(), str_to_int)

Thanks for any help. I have found some fairly similar questions to this one, but still couldn't solve it.

[sobolevn]

You are looking for https://peps.python.org/pep-0646/ which is not fully supported yet :)

[maciejzj]

Thx, seems that I was looking in the right direction when I was trying to use Unpack and TypeVarTuple.

Out of curiosity, when I try the following using the new features:

from typing import Sequence, Callable, TypeVar, Generic, Unpack, TypeVarTuple
from torch.utils.data import Dataset, IterableDataset, TensorDataset


class StrDs(Dataset[tuple[str, str]]):
    def __init__(self):
        self.d = [("1", "10"), ("2", "20"), ("3", "30")]

    def __getitem__(self, idx: int) -> tuple[str, str]:
        return self.d[idx]


InT = TypeVarTuple("InT")
OutT = TypeVar("OutT", bound=tuple)


class TransformableDatasetWrapper(Generic[Unpack[InT], OutT], Dataset[InT]):
    def __init__(
        self,
        d: Dataset[InT],
        t: Callable[[Unpack[InT]], OutT],
    ):
        self.d = d
        self.t = t

    def get(self, i: int) -> OutT:
        return self.t(*self.d[i])


def str_to_int(i: str, ii: str) -> tuple[int, int]:
    return int(i), int(ii)


TransformableDatasetWrapper[tuple[str, str], tuple[int, int]](StrDs(), str_to_int)

and run mypy with enabled support for incomplete I get:

mypy --enable-incomplete-feature=Unpack --enable-incomplete-feature=TypeVarTuple p.py 
p.py:27: error: INTERNAL ERROR -- Please try using mypy master on GitHub:
https://mypy.readthedocs.io/en/stable/common_issues.html#using-a-development-mypy-build
If this issue continues with mypy master, please report a bug at https://github.com/python/mypy/issues
version: 1.5.1
p.py:27: : note: please use --show-traceback to print a traceback when reporting a bug

Have I used the Unpack and TypeVarTuple in a wrong way, or have I reached the limit of the incomplete implementation of these features?

[hmc-cs-mdrissi]

Internal error always means bug in mypy and you can open an issue there.

There are multiple type checkers available. Pyright and pyre both have support for typevartuple and you can test your code with them.

Your code does have a few errors, but is close. Here's a fixed version that type checks with pyright,

from typing import Sequence, Callable, TypeVar, Generic, Unpack, TypeVarTuple

InT = TypeVarTuple("InT")
class Dataset(Generic[Unpack[InT]]):
    def __getitem__(self, idx: int) -> tuple[Unpack[InT]]:
        ...

class StrDs(Dataset[Unpack[tuple[str, str]]]):
    def __init__(self):
        self.d = [("1", "10"), ("2", "20"), ("3", "30")]

    def __getitem__(self, idx: int) -> tuple[str, str]:
        return self.d[idx]


OutT = TypeVar("OutT", bound=tuple)


class TransformableDatasetWrapper(Generic[Unpack[InT], OutT], Dataset[Unpack[InT]]):
    def __init__(
        self,
        d: Dataset[Unpack[InT]],
        t: Callable[[Unpack[InT]], OutT],
    ):
        self.d = d
        self.t = t

    def get(self, i: int) -> OutT:
        return self.t(*self.d[i])


def str_to_int(i: str, ii: str) -> tuple[int, int]:
    return int(i), int(ii)


TransformableDatasetWrapper[Unpack[tuple[str, str]], tuple[int, int]](StrDs(), str_to_int)

Maybe this version will work with mypy as well. You almost always need to have Unpack next to type var tuple and your original code was missing several unpacks.