Typing: Inferred type too general during class inheritance

Question

I'm facing the following problem. Suppose I want to define generic class for Dataset and its Samples. Then, I want to create specific subclasses of these that will in the end be instantiated (E.g. DatasetOfSpecificSamples that contains MoreSpecificSample samples). To save some LoC, I want to implement common functionality in Dataset class.

Yet, I can't get the typehints in the inherited sublcasses right. They are too general. See the following example, where I would expect that __iter__ return type of DatasetOfSpecificSamples would be MoreSpecificSample, yet MyPy infers just Sample.

Is there a way to enforce more specific hint without redifining the method itself?

from typing import Mapping, Iterator

class Sample:
    pass


class MoreSpecificSample(Sample):
    pass


class Dataset:
    def __init__(self, samples: Mapping[str, Sample]):
        self.samples = samples

    def __iter__(self) -> Iterator[Sample]:
        yield from self.samples.values()


class DatasetOfSpecificSamples(Dataset):
    samples: Mapping[str, MoreSpecificSample]

    # def __iter__(self) -> Iterator[CommonCriteriaCert]: <-- This fixes the problem, as it specifies the type
    #     yield from self.certs.values()

dset = DatasetOfSpecificSamples()
for sample in dset: # <-- the inferred type of "sample" is Sample, not MoreSpecificSample
    pass

Answer 1

Make Dataset class Generic:

T = TypeVar("T")

class Dataset(Generic[T]):

    def __init__(self, samples: Mapping[str, T]):
        self.samples = samples

    def __iter__(self) -> Iterator[T]:
        yield from self.samples.values()


class DatasetOfSpecificSamples(Dataset[MoreSpecificSample]):
    samples: Mapping[str, MoreSpecificSample]