Packing TypeVarTuples into bound, covariant TypeVars

I made a PR to numpy to improve shape typing. The primary purpose was to make the type variable for shape (a) bound to a tuple of ints and (b) covariant.

As an additional feature, someone requested an alias that used TypeVarTuple, as in the PEP 646 example. Although TypeVarTuple cannot be bound or made covariant, unpacking one inside a bound typevar seemed to work, and still does, for the most part. What is the intended intended behavior? As a related question, is there any activity to add bounds/variance to TypeVarTuple?

MWE (as debug.pyi, to avoid “___ is unbound” errors):

Define Array alias to ndarray (no numpy)

from typing import Generic, NewType, TypeVar
import sys
if sys.version_info < (3, 11):
    from typing_extensions import assert_type, reveal_type, TypeVarTuple, Unpack
else:
    from typing import assert_type, reveal_type, TypeVarTuple, Unpack

_ShapeType = TypeVar("_ShapeType", covariant=True, bound=tuple[int, ...])
_ShapeTypeTuple = TypeVarTuple("_ShapeTypeTuple")

class ndarray(Generic[_ShapeType]):
    shape: _ShapeType

Array = ndarray[tuple[Unpack[_ShapeTypeTuple]]]

The definition of Array causes a type error in both pyright and mypy: Type argument "tuple[*_ShapeTypeTuple]" of "ndarray" must be a subtype of "tuple[int, ...]". The crux here is whether type checkers still safely handle functions that use Array, and whether that safety is guaranteed or an accident of implementation.

Test 1: Array parameters unpacked correctly, bounds apply

Length = NewType("Length", int)
Width = NewType("Width", int)
arr: Array[Length, Width]
assert_type(arr.shape, tuple[Length, Width])
bad: Array[str]

This passes pyright. OTOH, mypy allows the definition of arr, but correctly flags the definiton of bad with "tuple[str]" of "ndarray" must be a subtype of "tuple[int, ...]". I assume this means that pyright has decided that since the Array definition is flawed, it can’t provide any safety to Array objects. Mypy, however, seems to pass the bounds of _ShapeType through _ShapeTypeTuple.

Test 2: I can use Array in function signatures

Ts = TypeVarTuple("Ts")
def stack(a: Array[Unpack[Ts]], b: Array[Unpack[Ts]]) -> Array[int, Unpack[Ts]]: ...

doubled = stack(arr, arr)
assert_type(doubled.shape, tuple[int, Length, Width])
doublebad = stack(bad, bad)

Mypy raises an error about all the types in the stack signature ("tuple[*Ts]"/"tuple[int, *Ts]" must be a subtype of subtype of "tuple[int, ...]"), but pyright does not. I assume that’s because of when each resolves type aliases. Both pass the assert_type, but both unfortunately also accept doubelbad. Nevertheless, if mypy prevents us from ever creating an Array[str], it would seem like a safe to ignore the function signature errors, no?

Corollary problem: Cannot unpack tuple-bound TypeVar

T = TypeVar("T", covariant=True, bound=tuple[int, ...])
def stack2(a: ndarray[T], b: ndarray[T]) -> ndarray[tuple[int, Unpack[T]]]: ...

doubled2 = stack2(arr, arr)
assert_type(doubled2.shape, tuple[int, Length, Width])
doublebad2 = stack2(bad, bad)

Here, both mypy and pyright complain that they can only unpack a tuple or a TypeVarTuple, but not a tuple-bound TypeVar. As a result, the assert_type fails, as since the shape of doubled2 is either tuple[int, *tuple[Any, ...]] (mypy) or tuple[int, Unknown] (pyright). This is unfortunate, because this definition (correctly) causes type checkers to complain about stack2(bad, bad)

I’m having difficulty following your example. It’s incomplete, and when I try to make it complete, I see behaviors that don’t match your description. In particular, I’m seeing that assert_type does not pass.

Could you take another crack at producing a minimal, self-contained example, preferably one that doesn’t reference symbols from np? Here’s my attempt to fix up the example.

I presume that you’re looking for a way to specify both the dtype and the shape of an array, right? The shape would presumably be described by a variadic type var and the dtype described by a non-variadic type var. That’s not what I’m seeing in your example though.

To answer your question, I’m not familiar with any active efforts to add bounds/variance to TypeVarTuple. The PEP and the typing spec allude to the fact that these could be added in the future.

Hey Eric, thanks for the reply! I’ve rewrote the MWE and checked it against mypy 1.11 and pyright 1.1.372 on python 3.9 and 3.11. I did my tests in a pyi file, so didn’t bind any variables (which the pyright playground flags), but here’s a working pyright playground variant.

I’ve eschewed dtype for brevity of the example. That part of numpy typing seems to be working fine. My question is “what safety can/should type checkers provide for TypeVarTuples that are used inside of a bound TypeVar”? In particular, some functions can’t be typed without TypeVarTuples (see the corollary in the MWE), and I’m not sure the best way to handle them.

If there isn’t a good solution, then I’d like to start the alluded-to work on adding variance/bounds to TypeVarTuples and/or enable unpacking of tuple-bound TypeVars.

Thanks for the updated sample. I think I understand now what you’re trying to do.

I don’t think there’s a good solution using TypeVarTuple currently — at least not one that guarantees that the tuple elements are all subtypes of int.

You’re welcome to start work on spec’ing support for upper bounds and variance for TypeVarTuple. This will require a grammar change (for PEP 695 generics syntax), so it will require a formal PEP and approval from the steering council.

It probably makes sense to consider upper bounds and variance for ParamSpecs at the same time.

1 Like