CODEWITHSUNDEEP
pythonpython-3.xooptypesmypy
Solomon Bothwell
Solomon Bothwell

Reputation: 1134

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

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.

Upvotes: 3

Views: 1427

Answers (1)

chadrik
chadrik

Reputation: 3462

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())

Upvotes: 1

Related Questions