Как задать строгий параметр axis в Python с Literal и TypeAlias

Как описать параметр axis в стиле pandas через typing.Literal и TypeAlias: точные значения, подсказки IDE и статическая проверка типов. Без неоднозначностей.

Передать узкий, явно заданный набор допустимых значений для параметра функции сложнее, чем кажется, если ограничиться широким Union. Типичный пример — аргумент в духе axis, где принимаются и целые числа, и строки, но только в определённых парах. Полагаться на общий Union[int, str] — значит оставить слишком много неоднозначностей для пользователей и инструментов. Ниже — точный способ донести и закрепить это намерение на уровне типов.

Проблема

Предположим, вы повторяете поведение pandas.DataFrame.dropna и принимаете для параметра axis и int, и str. Наивная подсказка типов могла бы выглядеть так:

from typing import Union

def purge_na(*, axis: Union[int, str] = 0) -> None:
    ...

Эта сигнатура лишь сообщает, что допустимы int и str, но не уточняет, какие именно значения валидны и что они образуют пары.

Почему Union недостаточно

Union[int, str] слишком широк. Проверяющие типы и IDE не видят задуманный набор вариантов, поэтому не предложат точных автодополнений и не предупредят на этапе статической проверки. Докстринг помогает читателям, но не направляет инструменты и не ограничивает значения во время анализа.

Решение с точными наборами значений

Используйте typing.Literal, чтобы перечислить точные допустимые значения, и объедините их через TypeAlias. Так контракт становится явным и понятным и разработчикам, и анализаторам типов.

from typing import Literal, TypeAlias

RowAxisSpec: TypeAlias = Literal[0, "index"]
ColAxisSpec: TypeAlias = Literal[1, "columns"]
AxisSpec: TypeAlias    = RowAxisSpec | ColAxisSpec  # 0/"index" или 1/"columns"

def remove_na_entries(*, axis: AxisSpec = 0) -> None:
    """
    Parameters
    ----------
    axis : {0, 1, "index", "columns"}, default 0
        0 or "index"    -> operate over rows
        1 or "columns"  -> operate over columns
    """
    ...

Этот подход описывает оба домена и их соответствие. Разработчики видят, какие литералы ожидаются, IDE подсказывает их, а статический анализ заранее отметит неподдерживаемые значения.

Если дальше вы хотите поощрять использование строк, сузьте псевдоним до строк, оставив целые числа как наследие с путём к деприкации.

AxisSpec = Literal["index", "columns"]  # оставить 0/1 как наследие с предупреждением об устаревании

Статические подсказки и поведение во время выполнения

Literal улучшает поддержку в IDE и статическую проверку, но не навязывает значения во время выполнения. Если нужен и рантайм-контроль, добавьте короткую проверку в начале.

def remove_na_entries(*, axis: AxisSpec | int = 0) -> None:
    if axis not in (0, 1, "index", "columns"):
        raise ValueError("axis must be 0/1/'index'/'columns'")
    ...

Почему это важно

Явное описание наборов значений снимает неоднозначность API и улучшает работу разработчика. Инструменты показывают точные варианты прямо в месте вызова, а статические анализаторы ловят ошибки до запуска. Документация остаётся полезной для повествовательной части, но именно система типов несёт официальный контракт.

Выводы

Когда параметр принимает закрытый набор значений, особенно если это чёткие пары, используйте typing.Literal вместе с TypeAlias, чтобы зафиксировать контракт. Держите докстринг согласованным с литералами для читателей и, при необходимости, добавьте лёгкую проверку на этапе выполнения. Такое сочетание даёт ясность, удобство и безопасность без усложнения реализации.