Python / MyPy: как аннотировать метод, который может возвращать один из нескольких разных типов объектов?

Как мне аннотировать возвращаемый тип метода, который может возвращать несколько различных типов объектов?

В частности, это метод, с которым у меня проблемы:

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

Я пробовал Union[shelf.Sheet, guillotine.Guillotine, maximal_rectangles.MaximalRectangle], но MyPy дает мне массу ошибок, когда я позже использую метод _bin_factory в своем коде. Ошибки, похоже, связаны с тем фактом, что все три типа объектов в Союзе имеют разные атрибуты друг от друга.


person Solomon Bothwell    schedule 05.11.2017    source источник
comment
Одна из идей состоит в том, чтобы изменить каждый из этих классов, чтобы они создавали подкласс некоторого класса Bin, который содержит метод bin(). Затем измените свой _bin_factory() метод, чтобы он возвращал объекты типа Bin. Таким образом, независимо от вашей стратегии биннинга, вызывающий может просто вызвать метод bin() и рассматривать возвращаемый тип как один и тот же тип.   -  person Michael0x2a    schedule 06.11.2017


Ответы (1)


Вот решение с использованием 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())

В качестве альтернативы, если вы хотите немного изменить свой подход, вы можете предоставить более чистый 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())
person chadrik    schedule 01.02.2018