Python Typehinting для функции

У меня есть следующий код Python:

foo.create("TireModelA")

Функция принимает на вход строку. Вариантов всего 3: TireModelA, TireModelB, TireModelC.

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

Спасибо.


person Joao Antunes    schedule 31.07.2020    source источник


Ответы (1)


У вас есть как минимум две возможности Literal и Enum:

from typing import Literal
from enum import Enum


class FooParamsEnum(str, Enum):
    tire_model_a = "TireModelA"
    tire_model_b = "TireModelB"
    tire_model_c = "TireModelC"


FooParamsLiterals = Literal["TireModelA", "TireModelB", "TireModelC"]


class Foo1:
    def create(self, param: FooParamsLiterals) -> FooParamsLiterals:
        return param


class Foo2:
    def create(self, param: FooParamsEnum) -> FooParamsEnum:
        return param


foo1 = Foo1()
foo1.create("TireModelA")
foo2 = Foo2()
foo2.create(FooParamsEnum.tire_model_a)
person alex_noname    schedule 31.07.2020