How to implement associated types in Python/Mypy? Or, what to do when wanting sub-classes to have subclass arguments?

Question

Consider the (simplified) code, where I want Lean* and Isabelle* extend the Base*.

class BaseProblem: ...
class BaseStep: ...

class LeanProblem(BaseProblem): ...
class LeanStep(BaseStep): ...

class IsabelleProblem(BaseProblem): ...
class IsabelleStep(BaseStep): ...

class BaseProver:
    def f(self, problem: BaseProblem, step: BaseStep): ...
class LeanProver(BaseProver):
    def f(self, problem: LeanProblem, step: LeanStep): ...
class IsabelleProver(BaseProver):
    def f(self, problem: IsabelleProblem, step: IsabelleStep): ...

However, the f function will have problem in mypy:

Argument 1 of "f" is incompatible with supertype "LeanProblem";
supertype defines the argument type as "BaseProblem"  [override]

I know it can be solved by generics, such as:

TProblem = TypeVar('TProblem', bound=BaseProblem)
TStep = TypeVar('TStep', bound=BaseStep)

class BaseProver(Generic[TProblem, TStep]):
    def f(self, problem: TProblem, step: TStep): ...
class LeanProver(BaseProver[LeanProblem, LeanStep]):
    def f(self, problem: LeanProblem, step: LeanStep): ...
...

However, instead of only "Problem" and "Step", indeed I have more of such types (say, 10). Thus, the approach to use generic will be quite ugly IMHO.

When using Rust, I know we can have associated types and it can solve the problem; but I have not found an equivalent in Python.