[Python] PyTorch type for dataset with length

I am creating a meta-dataset that combines data from multiple input datasets.

from torch.utils.data import Dataset, IterableDataset

class MetaDataset(Dataset):
def __init__(self, regular_dataset: Dataset, iterable_dataset: IterableDataset):
self.regular_dataset = regular_dataset
self.iterable_dataset = iterable_dataset
pass # Do other stuff...


I got a type warning when I tried to access len(self.regular_dataset) from within MetaDataset

It turns out that the PyTorch type definition for Dataset intentionally doesn't include __len__. Therefore, I have to build my own type:

from torch.utils.data import Dataset, IterableDataset

class DatasetWithLength(Dataset):
def __len__(self) -> int:
pass

class MetaDataset(Dataset):
def __init__(self, regular_dataset: DatasetWithLength, iterable_dataset: IterableDataset):
self.regular_dataset = regular_dataset
self.iterable_dataset = iterable_dataset
pass # Do other stuff...


But now, I get a Expected type 'DatasetWithLength', got 'FirstDataset' instead warning when I try to do this:

foo = MetaDataset(
FirstDataset(),
FirstIterableDataset(),
)


How do I correctly define the type of a PyTorch Dataset that has the length property?

Continue reading...