Как я могу напечатать подсказку функции, где тип возвращаемого значения зависит от типа ввода аргумента? - PullRequest
Новая
       6

Как я могу напечатать подсказку функции, где тип возвращаемого значения зависит от типа ввода аргумента?

0 голосов
/ 21 сентября 2018

Предположим, у меня есть функция, которая преобразует типы данных Python в типы данных Postgres, например: 

def map_type(input):
    if isinstance(input, int):
        return MyEnum(input)
    elif isinstance(input, str):
        return MyCustomClass(str)

Я мог бы напечатать подсказку так: 

def map_type(input: Union[int, str]) -> Union[MyEnum, MyCustomClass]: ...

Нотогда код, подобный следующему, не сможет проверить тип, даже если он верен: 

myvar = map_type('foobar')
print(myvar.property_of_my_custom_class)

Полный пример (рабочий код, но ошибки в подсказке типа): 

from typing import Union
from enum import Enum


class MyEnum(Enum):
    VALUE_1 = 1
    VALUE_2 = 2


class MyCustomClass:

    def __init__(self, value: str) -> None:
        self.value = value

    @property
    def myproperty(self) -> str:
        return 2 * self.value


def map_type(value: Union[int, str]) -> Union[MyEnum, MyCustomClass]:

    if isinstance(value, int):
        return MyEnum(value)
    elif isinstance(value, str):
        return MyCustomClass(value)
    raise TypeError('Invalid input type')


myvar1 = map_type(1)
print(myvar1.value, myvar1.name)

myvar2 = map_type('foobar')
print(myvar2.myproperty)

IЯ знаю, что я мог бы разделить отображение на две функции, но цель в том, чтобы иметь общую функцию отображения типов.

Я также думал о работе с классами и полиморфизмом, но тогда как бы я набралПодсказка самые верхние методы класса?Потому что их тип вывода будет зависеть от конкретного типа экземпляра.

1 Ответ

0 голосов
/ 21 сентября 2018

Это именно то, для чего предназначены перегрузки функции .

Короче говоря, вы делаете следующее: 

from typing import overload

# ...snip...

@overload
def map_type(value: int) -> MyEnum: ...

@overload
def map_type(value: str) -> MyCustomClass: ...

def map_type(value: Union[int, str]) -> Union[MyEnum, MyCustomClass]:
    if isinstance(value, int):
        return MyEnum(value)
    elif isinstance(value, str):
        return MyCustomClass(value)
    raise TypeError('Invalid input type')

Теперь, когда вы делаете map_type(3),mypy поймет, что тип возвращаемого значения MyEnum.

И во время выполнения единственная функция, которая должна фактически выполняться, является последней - первые две полностью переопределены и игнорируются.

...