[Question] Integer `Literal` arithmetic for, *e.g.*, physical units?

[ferdezgarcia]

Hello! Apologies beforehand for polluting the repo's issues with questions tangential to PEP 827.

For the past few days I've been banging my head against Python's type system trying to replace a dependency on the otherwise excellent pint library with my own bare-bones implementation. I'm by no means a computer scientist nor a professional programmer, so I've struggled to come up with a way to express physical units' dimensionality through the type system.

While racking my brain I came upon PEP 827 and I found the array shape example particularly fitting. Does your current proposal allow for performing arithmetic operations on integer literals? E.g., Literal[1] + Literal[2] resolving to Literal[3]. Is there even a way to represent any integer literal similarly to LiteralString?

For further context, here is a toy example of what I'm trying to achieve:

from decimal import Decimal
from fractions import Fraction
from typing import Any, Literal, Self, final, overload, reveal_type

from typemap_extensions import GetArg, IsAssignable, RaiseError

type Range[T] = tuple[Literal[i] for i in ...] # How would one generate a tuple of integer `Literal`s?
type AddDims[T, S] = (
    tuple[GetArg[T, tuple, i] + GetArg[S, tuple, i] for i in Range[Length[T]]] # Can integer `Literal`s be added?
    if IsAssignable[Length[T], Length[S]]  # Units belong to same system
    else RaiseError[Literal["Incompatible unit system"], T, S]
)
type EqDims[T, S] = (
    T
    if IsAssignable[T, S]
    else RaiseError[Literal["Incompatible dimensions"], T, S]
)


@final
class Unit[*T]:
    """Physical unit abstraction."""

    def __init__(self: Self, dimensions: tuple[*T], /) -> None:
        """Infer dimensionality from exponents."""

        self.dimensions: tuple[*T] = dimensions

    @overload
    def __mul__[*S](
        self: Self, other: Unit[*S], /
    ) -> Unit[AddDims[tuple[*T], tuple[*S]]]: ...

    @overload
    def __mul__[S: (Decimal, Fraction)](
        self: Self, other: S, /
    ) -> Quantity[S, *T]: ...

    def __mul__(self: Self, other: Any, /) -> Any:
        """Return `Unit` with computed dimensions or `Quantity`."""

        raise NotImplementedError()

    def __rmul__[S: (Decimal, Fraction)](
        self: Self, other: S, /
    ) -> Quantity[S, *T]:
        """Return `Quantity` from scalar and `Unit`."""

        raise NotImplementedError()

@final
class Quantity[T: (Decimal, Fraction), *S]:
    """Physical quantity abstraction."""

    def __init__(self: Self, magnitude: T, unit: Unit[*S], /) -> None:
        """Infer dimensionality from `Unit`."""

        self.magnitude: T = magnitude
        self.unit: Unit[*S] = unit

    def __add__[*U](
        self: Self, other: Quantity[T, *U], /
    ) -> Quantity[T, EqDims[tuple[*S], tuple[*U]]]:
        """Disallow addition when dimensions do not match."""

        raise NotImplementedError()


# Expected behavior:

s = Unit((0, 1))
mps = Unit((1, -1))

reveal_type(s)  # Unit[Literal[0], Literal[1]]
reveal_type(mps)  # Unit[Literal[1], Literal[-1]]

m = s * mps

reveal_type(m)  # Unit[Literal[1], Literal[0]]

dist = Decimal(10) * m
speed = Decimal(1) * mps

reveal_type(dist)  # Quantity[Decimal, Unit[Literal[1], Literal[0]]
reveal_type(speed)  # Quantity[Decimal, Unit[Literal[1], Literal[-1]]]

reveal_type(dist + dist)  # Quantity[Decimal, Unit[Literal[1], Literal[0]]]
reveal_type(dist + speed)  # Error: Incompatible dimensions

Range and AddDims are pseudo-code; I'm not exactly sure whether the operation would be allowed and whether that's how you'd retrieve the exponent from the tuple.

Does this kind of behavior fit within the scope of the PEP? Or would it require additional modifications to the type system such as allowing subclassing or changing how Literals are inferred?

Looking forward to hearing what becomes of the PEP regardless; we're all sure to find some interesting applications for the proposed additions!