Python/MyPy: How to annotate a method that can return one of several different types of objects?

user_5478942 0
Viewed 1088

How should I annotate the return type of a method that can return multiple different types of objects?

Specifically this is the method I'm having trouble with:

def _bin_factory(self) -> Any:
    """
    Returns a bin with the specificed algorithm,
    heuristic, and dimensions
    """
    if self.algorithm == 'guillotine':
        return guillotine.Guillotine(self.bin_width, self.bin_height, self.rotation,
                                     self.rectangle_merge, self.split_heuristic)
    elif self.algorithm == 'shelf':
        return shelf.Sheet(self.bin_width, self.bin_height, self.rotation, self.wastemap)
    elif self.algorithm == 'maximal_rectangle':
        return maximal_rectangles.MaximalRectangle(self.bin_width, self.bin_height, self.rotation)
    raise ValueError('Error: No such Algorithm')

I tried Union[shelf.Sheet, guillotine.Guillotine, maximal_rectangles.MaximalRectangle] but MyPy gives me a ton of errors where I use the _bin_factory method later on in my code. The errors seem to center around the fact that all three object types in the Union have different attributes from one another.

mypy python-3.x oop python types
1 Answers

Here's a solution using typing.Generic

from typing import Generic, TypeVar

T = TypeVar('T', 'Guillotine', 'Sheet', 'MaximalRectangle')

class Guillotine:
    pass

class Sheet:
    pass

class MaximalRectangle:
    pass

class Algo(Generic[T]):
    def __init__(self, algorithm: str) -> None:
        self.algorithm = algorithm

    def _bin_factory(self) -> T:
        """
        Returns a bin with the specificed algorithm,
        heuristic, and dimensions
        """
        if self.algorithm == 'guillotine':
            return Guillotine()  # type: ignore
        elif self.algorithm == 'shelf':
            return Sheet()  # type: ignore
        elif self.algorithm == 'maximal_rectangle':
            return MaximalRectangle()  # type: ignore
        raise ValueError('Error: No such Algorithm')


algo: Algo[Guillotine] = Algo('guillotine')
reveal_type(algo._bin_factory())

Alternately, if you're willing to modify your approach a bit more, you can provide a cleaner API:

from typing import Generic, TypeVar, Type

T = TypeVar('T', 'Guillotine', 'Sheet', 'MaximalRectangle')

class Guillotine:
    pass

class Sheet:
    pass

class MaximalRectangle:
    pass

class Algo(Generic[T]):
    def __init__(self, algorithm: Type[T]) -> None:
        self.algorithm = algorithm  # type: Type[T]

    def _bin_factory(self) -> T:
        """
        Returns a bin with the specificed algorithm,
        heuristic, and dimensions
        """
        if self.algorithm is Guillotine:
            # handle custom arguments:
            return self.algorithm()
        elif self.algorithm is Sheet:
            # handle custom arguments:
            return self.algorithm()
        elif self.algorithm is MaximalRectangle:
            # handle custom arguments:
            return self.algorithm()
        raise ValueError('Error: No such Algorithm')

algo = Algo(Guillotine)
reveal_type(algo._bin_factory())
user_9675260
Related