Как задать функцию в Python с необязательными аргументами согласно литеральному типу?

Содержание

Вопрос или проблема
Ответ или решение
Объяснение кода:
Альтернативный подход:
Заключение:

Вопрос или проблема

Допустим, я хочу реализовать функцию маршрутизатора, которая будет принимать аргумент name и, при желании, аргумент data.

Например, функция будет вызываться так:

route("main")

для перехода на главный экран или

route("edit", item.id)

для перехода на экран редактирования элемента.

Этот код работает, но у меня нет никакой типизации, и поэтому типизаторы не поднимают ошибок:

from typing import Literal

def route(name: Literal["main", "edit"], *args) -> None:
    match name:
        case "main":
            print("Главный экран")
        case "edit":
            if not args:
                raise ValueError("Не предоставлен id элемента для редактирования.")
            print(f"Редактирование элемента с id: {args[0]}")
        case _:
            raise ValueError(f"Неверное имя маршрута: {name}")

route("main") # Главный экран
route("edit", 12) # Редактирование элемента с id: 12
route("edit")  # Поднимает ValueError

Fiddle

Я пробовал использовать TypeGuard или TypeIs, но, если я правильно понимаю PEP 647, они будут работать в противоположном направлении — они сужают тип внутри функции, но не влияют на сигнатуру функции.

Решение, предложенное @deceze, может заключаться в использовании @overload:

from typing import overload, Union, Literal

class Router:
    @overload
    def route(self, name: Literal["main"]) -> None:
        ...

    @overload
    def route(self, name: Literal["edit"], id: int) -> None:
        ...

    def route(self, name: Literal["main"] | Literal["edit"], id: int | None = None) -> None:
        if name == "main":
            print("главный")
        if name == "edit":
            if not id:
                print("ошибка")
            print(id)

# Использование
router = Router()
router.route("main")
router.route("edit", 5)
router.route("edit") # Ошибка линтера
router.route("main", 6) # Ошибка линтера

Лично я думаю, что это выглядит несколько перегруженно и не будет хорошо масштабироваться.

Fiddle

Вероятно, это должны быть две отдельные функции, которые вообще не требуют первого аргумента-литерала.

def route_main() -> None:
    print("Главный экран")

def route_edit(id: int) -> None:
    print(f"Редактирование элемента с id: {id}")

route_main()  # Главный экран
route_edit(12)  # Редактирование элемента с id: 12
route_edit()  # Поднимает TypeError, но также не проходит статическую типизацию

Любая логика, которая решает, какой литерал передать в route, может также быть использована для определения, какую из двух функций выше вызвать.

Ответ или решение

Чтобы реализовать функцию маршрутизации с использованием литеральных типов и опциональных аргументов в Python, вы можете воспользоваться декоратором @overload из модуля typing. Это позволит вам задать разные сигнатуры для функции в зависимости от ее аргументов.

Вот пример реализации функции route, которая принимает аргумент name и опционально аргумент data:

from typing import overload, Literal, Optional

class Router:
    @overload
    def route(self, name: Literal["main"]) -> None:
        ...

    @overload
    def route(self, name: Literal["edit"], data: int) -> None:
        ...

    def route(self, name: Literal["main", "edit"], data: Optional[int] = None) -> None:
        if name == "main":
            print("Main screen")
        elif name == "edit":
            if data is None:
                raise ValueError("No item id provided for editing.")
            print(f"Editing item id: {data}")
        else:
            raise ValueError(f"Invalid route name: {name}")

# Пример использования
router = Router()
router.route("main")       # Main screen
router.route("edit", 12)   # Editing item id: 12
router.route("edit")       # Raises ValueError

Объяснение кода:

Использование @overload:
- Декоратор @overload позволяет определить несколько версий функции route. В первой версии функция принимает только имя "main", а во второй — имя "edit" с обязательным параметром data типа int.
Основная реализация функции:
- В основной функции route реализована логика, которая проверяет значение name. Если это "main", функция выводит соответствующее сообщение. Если это "edit" и значение data отсутствует, генерируется ошибка ValueError.
Обработка ошибок:
- Если вызвана функция с недопустимым значением name или с отсутствующим data для "edit", вызывается соответствующее исключение.

Альтернативный подход:

Хотя описанный метод работает, можно рассмотреть более простой подход, если логика маршрутизатора не предполагает сложных манипуляций с аргументами. Например, определить отдельные функции для каждой маршрутизации:

def route_main() -> None:
    print("Main screen")

def route_edit(data: int) -> None:
    print(f"Editing item id: {data}")

# Пример использования
route_main()      # Main screen
route_edit(12)    # Editing item id: 12
# route_edit()   # Raises TypeError

Заключение:

Оба подхода имеют свои преимущества и недостатки. Первое решение с использованием перегрузок позволяет хранить всю логику в одном месте и может быть более читаемым для коротких маршрутизаторов. Второе решение с отдельными функциями проще и легче масштабируется при добавлении новых маршрутов. Выбор подхода зависит от ваших требований и предпочтений в структуре кода.