Search code examples
pythonobjectabstract-classtype-hinting

Specify return type of a wrapper function that calls an abstract method in Python


For this example, consider the simplified scenario where a Solver will return a Solution.

We have Solutions:

class Solution(ABC):
    pass


class AnalyticalSolution(Solution):
    pass


class NumericalSolution(Solution):
    def get_mesh_size(self) -> float:
        return 0.12345

And Solvers:

class Solver(ABC):
    def solve(self, task: int) -> Solution:
        # Do some pre-processing with task
        # ...
        return self._solve(task)

    @abstractmethod
    def _solve(self, task: int) -> Solution:
        pass


class NumericalSolver(Solver):
    def _solve(self, task: int) -> NumericalSolution:
        return NumericalSolution()


class AnalyticalSolver(Solver):
    def _solve(self, task: int) -> AnalyticalSolution:
        return AnalyticalSolution()

The problem I encounter results from the implementation of the wrapper method solve that then calls the abstract method _solve. I often encounter a situation like this where I want to do some preprocessing in the solve method that is the same for all solver, but then the actual implementation of _solve might differ.

If I now call the numerical solver and call the get_mesh_size() method, Pylance (correctly) tells me that a Solution object has no get_mesh_sizemember.

if __name__ == "__main__":
    solver = NumericalSolver()
    solution = solver.solve(1)
    print(solution.get_mesh_size())

I understand that Pylance only sees the interface of solve which indicates that the return type is a Solution object that does not need to have a get_mesh_size method. I am also aware that this example works at runtime.

I tried to use TypeVar like this (actually, because ChatGPT suggested it):

class Solution(ABC):
    pass
T = TypeVar("T", bound=Solution)

and then rewrite the Solver class:

class Solver(ABC):
    def solve(self, task: int) -> T:
        # Do some pre-processing with task
        # ...
        return self._solve(task)

    @abstractmethod
    def _solve(self, task: int) -> T:
        pass

But Pylance now tells me TypeVar "T" appears only once in generic function signature. So this can't be the solution.

How do I get typing to work with this example?


Solution

  • You may use Generic[T] as a base for Solver and then extend it as follows

    from abc import ABC, abstractmethod
    from typing import TypeVar, Generic
    
    
    class Solution(ABC):
        pass
    
    
    class AnalyticalSolution(Solution):
        pass
    
    
    class NumericalSolution(Solution):
        def get_mesh_size(self) -> float:
            return 0.12345
    
    
    SolutionGeneric = TypeVar("SolutionGeneric", bound=Solution)
    
    
    class Solver(ABC, Generic[SolutionGeneric]):
        def solve(self, task: int) -> SolutionGeneric:
            # Do some pre-processing with task
            # ...
            return self._solve(task)
    
        @abstractmethod
        def _solve(self, task: int) -> SolutionGeneric:
            pass
    
    
    class NumericalSolver(Solver[NumericalSolution]):
        def _solve(self, task: int) -> NumericalSolution:
            return NumericalSolution()
    
    
    class AnalyticalSolver(Solver):
        def _solve(self, task: int) -> AnalyticalSolution:
            return AnalyticalSolution()
    
    
    if __name__ == "__main__":
        solver = NumericalSolver()
        solution = solver.solve(1)
        print(solution.get_mesh_size())