Literal values not bound to TypeVar as Literal values?

[ringohoffman]

I'm trying to add matrix shape type checking to taichi. They have a MatrixType class which can be used to create a 2D Matrix. The 2 shape parameters to MatrixType, N and M, can be provided as literal integers, which taichi.math actually does. But unless I explicitly declare a MatrixType using Literal, N and M are always inferred as int, even when using literals... implicitly opting you out of static shape checking.

I feel that it is unexpected and burdensome to have to manually declare the shape when you are already using a literal values to define the shape. Could the default instead be to use the Literal values?

from __future__ import annotations

from typing import Any, Generic, TypeVar, overload
from typing import Literal as L

M = TypeVar("M", bound=int)
N = TypeVar("N", bound=int)
P = TypeVar("P", bound=int)
DtypeT = TypeVar("DtypeT")


class Matrix(Generic[DtypeT, N, M]):
    def __init__(self, n: N, m: M, dt: DtypeT | None = None) -> None:
        ...

    @overload
    def __matmul__(self: Matrix[DtypeT, N, L[1]], other: Matrix[DtypeT, N, L[1]], /) -> DtypeT:
        ...

    @overload
    def __matmul__(self: Matrix[DtypeT, N, L[1]], other: Matrix[DtypeT, N, P], /) -> Matrix[DtypeT, P, L[1]]:
        ...

    @overload
    def __matmul__(self: Matrix[DtypeT, N, M], other: Matrix[DtypeT, M, L[1]], /) -> Matrix[DtypeT, N, L[1]]:
        ...

    @overload
    def __matmul__(self: Matrix[DtypeT, N, M], other: Matrix[DtypeT, M, P], /) -> Matrix[DtypeT, N, P]:
        ...

    def __matmul__(self, other: Any, /) -> Any:
        ...
    

class MatrixType(Generic[DtypeT, N, M]):
    def __init__(self, n: N, m: M, ndim: L[0, 1, 2], dtype: DtypeT) -> None:
        self.n = n
        self.m = m
        self.ndim = ndim

    def __call__(self, *args: *tuple[float | int, *tuple[float | int, ...]]) -> Matrix[DtypeT, N, M]:
        ...


m4x4 = MatrixType(4, 4, 2, float)
m4x2 = MatrixType(3, 3, 2, float)

m4x2(0) @ m4x4(0)  # no error?

reveal_type(m4x4)  # Type of "m4x4" is "MatrixType[type[float], int, int]"

m4x4_explicit: MatrixType[type[float], L[4], L[4]] = MatrixType(4, 4, 2, float)
m4x2_explicit: MatrixType[type[float], L[4], L[2]] = MatrixType(4, 2, 2, float)

m4x2_explicit(0) @ m4x4_explicit(0)  # error as expected
m4x4_explicit(0) @ m4x2_explicit(0)  # OK

[juk-r]

Possible solution is to add __new___:

from __future__ import annotations

from typing import Any, Generic, TypeVar, overload
from typing import Literal as L, Literal

IntLiteral = Literal[1,2,3,4,5,6,7,8,9,10]
M_literal = TypeVar("M_literal", bound=IntLiteral)
N_literal = TypeVar("N_literal", bound=IntLiteral)
M = TypeVar("M", bound=int, covariant=True)
N = TypeVar("N", bound=int, covariant=True)
P = TypeVar("P", bound=int)
DtypeT = TypeVar("DtypeT")


class Matrix(Generic[DtypeT, N, M]):
    def __init__(self, n: int, m: int, dt: DtypeT | None = None):
        ...
    
    @overload
    def __new__(cls, n: N_literal, m: M_literal, dt: DtypeT | None = None) -> Matrix[DtypeT, N_literal, M_literal]: ...
    @overload
    def __new__(cls, n: int, m: int, dt: DtypeT | None = None) -> Matrix[DtypeT, int, int]: ...
    def __new__(cls, n: int, m: int, dt: DtypeT | None = None) -> Matrix[DtypeT, int, int]:
        return super().__new__(cls)

    @overload
    def __matmul__(self: Matrix[DtypeT, N, L[1]], other: Matrix[DtypeT, N, P], /) -> Matrix[DtypeT, P, L[1]]:
        ...

    @overload
    def __matmul__(self: Matrix[DtypeT, N, M], other: Matrix[DtypeT, M, L[1]], /) -> Matrix[DtypeT, N, L[1]]:
        ...

    @overload
    def __matmul__(self: Matrix[DtypeT, N, M], other: Matrix[DtypeT, M, P], /) -> Matrix[DtypeT, N, P]:
        ...

    def __matmul__(self, other: Any, /) -> Any:
        ...
    

m4x4 = Matrix(4, 4, float)
reveal_type(m4x4)  # Type of "m4x4" is "Matrix[type[float], Literal[4], Literal[4]]"

[erictraut]

Pyright's constraint solver (the portion of the type checker that "solves" type variables) retains literals only when they are required to satisfy a constraint. If no such constraint exists, the constraint solver widens them to their non-literal form. This is the expected behavior in the vast majority of cases, and it would be extremely cumbersome if it were otherwise. For example, [1, 2, 3] would then be inferred as list[Literal[1, 2, 3] which is almost definitely not what the user intends.

This behavior isn't likely to change in any Python type checker, so if your proposed approach depends on constraint solvers retaining literals, you should rethink your approach.

If you want the constraint solver to retain literals, you'd need to add some constraint (e.g. an explicit type annotation or an upper bound, as @juk-r suggests above) that forces it to do so. For example:

m4x4: MatrixType[type[float], L[4], L[4]] = MatrixType(4, 4, 2, float)
m4x2: MatrixType[type[float], L[3], L[3]] = MatrixType(3, 3, 2, float)

or

IntLiteral = L[1,2,3,4,5,6,7,8,9,10]
M = TypeVar("M", bound=IntLiteral)
N = TypeVar("N", bound=IntLiteral)
P = TypeVar("P", bound=IntLiteral)

[ringohoffman]

But there is nothing that can be used to say any Literal int? Like Literal[int]? Or just Literal?

[erictraut]

No, there's nothing in the Python type system today that would mean "the union of all possible int literals". PEP 675 added something similar for string literals called StringLiteral.