Use typing.Literal and TypeAlias to define exact axis-like parameter values in Python APIs

Learn how to replace broad Union hints with typing.Literal and TypeAlias to specify exact axis options (0/1/index/columns) and improve hints and runtime checks

Conveying a tight, explicit set of valid values for a function parameter is harder than it looks if you only reach for a broad Union. A typical case is an axis-like argument, where both integers and strings are accepted, but only in specific pairs. Relying on a generic Union[int, str] leaves too much ambiguity for users and tooling. Here’s a precise way to communicate and enforce that intent at the type level.

Problem

Suppose you mirror the behavior of pandas.DataFrame.dropna and accept both int and str for the axis parameter. A straightforward type hint might look like this:

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

This signature tells callers that both int and str are allowed, but it does not convey which specific values make sense, nor the fact that these values come in pairs.

Why Union is not enough

Union[int, str] is too permissive. Type checkers and IDEs cannot infer the intended set of valid options, so developers won’t get targeted completions or early static feedback. A docstring helps humans, but it does not guide tooling and cannot restrict the accepted values at analysis time.

Solution with precise value sets

Use typing.Literal to enumerate exact allowed values and connect them into a union via a TypeAlias. This makes the contract explicit and discoverable for both users and type checkers.

from typing import Literal, TypeAlias
RowAxisSpec: TypeAlias = Literal[0, "index"]
ColAxisSpec: TypeAlias = Literal[1, "columns"]
AxisSpec: TypeAlias    = RowAxisSpec | ColAxisSpec  # 0/"index" or 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
    """
    ...

This approach communicates both domains and their mapping. Developers see exactly which literals are expected, IDEs can suggest them, and static analysis can flag unsupported values early.

If you plan to encourage string usage going forward, narrow the alias to strings while keeping integers as legacy with a deprecation path.

AxisSpec = Literal["index", "columns"]  # keep 0/1 as legacy with a deprecation warning

Static hints vs. runtime behavior

Literal improves IDE support and type checking, but it does not enforce values at runtime. If you also want a runtime guard, add a small check up front.

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

Why this matters

Being explicit about value sets reduces ambiguity in APIs and improves developer experience. Tooling can surface the exact options at the call site, and static analyzers can catch mistakes before code runs. Documentation remains useful for narrative guidance, but the type system carries the authoritative contract.

Takeaways

When a parameter accepts a closed set of values, especially when they come in well-defined pairs, use typing.Literal with TypeAlias to encode that contract. Keep the docstring aligned with the literals for human readers and, if needed, add a lightweight runtime guard. This combination balances clarity, discoverability, and safety without complicating the implementation.