IMO it would be nice if there was some official way to type-hint numpy arrays to describe shape. Has that happened yet?

I went down a rabbit hole trying to find the perfect way to type hint a matrix. Here's what I learned. First, the naive approach:

matrix3x3: list[list[int]] = [[1,2,3],[4,5,6],[7,8,9]]

There are two problems with this. The first is that list[list[int]] is a concrete type, and we'd like it to be abstract. As is, mypy would raise an error if we tried to do this:

import numpy as np
matrix3x3 = np.ndarray(shape=(3, 3), buffer=np.array([[1,2,3],[4,5,6],[7,8,9]])) # error

We would like to be able to do this though, because an NDArray shares all the relevant qualities of a matrix for our application.

The second problem is more subtle. matrix3x3 is meant to always be 3x3, but Python's lists are dynamically resizable, which means the shape can be tampered with. Ideally, we'd like mypy to be able to raise an error before runtime if someone else later tries to write matrix3x3.pop() or matrix3x3[0].append(something). This is not a problem in a language like Java, where Arrays are fixed-size.

There are three ways around these issues:

1. Switch to a statically-typed language.

This is the least preferable option, but something everyone should consider if they keep resisting duck typing. I still prefer duck typing at least for prototyping.

2. Modify the implementation.

This is certainly better, but not the best option. It's worth demonstrating how you could do this. For example, we can start with this:

class FixedShapeMatrix:
  def __init__(rows: int, cols: int) -> None:
    _matrix = [[0 for c in cols] for r in rows]

and continue defining the functionality of the FixedShapeMatrix object so that it has an immutable shape with mutable entries.

Another example is to just use numpy instead:

import numpy as np
from numpy import typing as npt

matrix3x3: npt.NDArray[np.int64] = np.ndarray((3,3), buffer=np.array([[1,2,3],[4,5,6],[7,8,9]])

Both of these solutions suffer from the same problem: they require significant refactoring of the existing project. And even if you had the time, you will lose generality when you pick the NDArray or FixedShapeMatrix implementations. Ideally, you want matrix3x3 to be structurally typed such that any of these implementations can be assigned to it. When you pigeonhole your matrix3x3 type, you lose the Abstraction of OOP. Thankfully, with Protocol, there's another way.

3. Structural subtyping.

Note: I'm going to be using Python 3.12 typing notation. As a quick reference, this is code in 3.11:

from typing import TypeVar, Generic

T = TypeVar('T', bound=int|float)


class MyClass(Generic[T]):

  def Duplicates(self, val: T) -> list[T]:
    return [val] * 2

And this is the same code in 3.12 (no imports needed):

class MyClass[T: int|float]:

  def Duplicates(self, val: T) -> list[T]:
    return [val] * 2

So, let's finally try to make an abstract matrix type directly. I'm going to show you how I iteratively figured it out. If you're already a little familiar with Protocol, you might have guessed this:

type Matrix[T] = Sequence[Sequence[T]]

But the problem is that Sequence is read-only. We're going to have to create our own type from scratch. The best way to start is to realize which methods we really need from the matrix:

1. indexing (read + write)

2. iterable

3. sized

The first attempt might be this:

from typing import Protocol


class Matrix(Protocol):

  def __getitem__(): ...

  def __setitem__(): ...

  def __len__(): ...

  def __iter__(): ...

But there are multiple problems with this. The first is that we need to explicitly annotate the types of each of these functions, or else our matrix won't be statically hinted.

from typing import Protocol, Iterator


class Matrix(Protocol):

  def __getitem__(self, index: int) -> int | Matrix: ...

  def __setitem__(self, index: int, val: int | Matrix) -> None: ...

  def __len__(self) -> int: ...

  def __iter__(self) -> Iterator[int | Matrix]: ...

The idea here is that matrix3x3[0][0] is an int, while the type of matrix3x3[0] is recursively a matrix that contains ints. But this doesn't protect against matrix3x3: Matrix = [1,2,3,[4,5,6],7,8,9] , which is not a matrix.

Here we realize that we should handle the internal rows as their own type.

from typing import Protocol, Iterator


class MatrixRow(Protocol):

  def __getitem__(self, index: int) -> int: ...

  def __setitem__(self, index: int, value: int) -> None: ...

  def __len__(self) -> int: ...

  def __iter__(self) -> Iterator[int]: ...


class Matrix(Protocol):

  def __getitem__(self, index: int) -> MatrixRow: ...

  def __setitem__(self, index: int, value: MatrixRow) -> None: ...

  def __len__(self) -> int: ...

  def __iter__(self) -> Iterator[MatrixRow]: ...

Now both the matrix and its rows are iterable, sized, and have accessible and mutable indexes.

matrix3x3: Matrix = [[1,2,3],[4,5,6],[7,8,9]] # good
matrix3x3.append([10,11,12]) # error - good!
matrix3x3[0][2] = 10 # good
matrix3x3[0][0] += 1 # good
matrix3x3[1].append(7) # error - good!

There's just one bug though. See if you can find it first:

matrix3x3[1] = [4,5,6,7] # no error - bad!

The solution is we need to remove __setitem__ from Matrix. We will still be able to modify the elements of any MatrixRow without it. Bonus points if you understand why (hint: references).

So let's go ahead and do that, and as a final touch, let's make it so that the matrix values all must have the same type. To do this, we enforce a generic type that supports integer operations (int, float, np.int32, np.float64, etc). Here's how I did that:

from typing import Protocol, Iterator, SupportsInt


class MatrixRow[T: SupportsInt](Protocol):

  def __getitem__(self, index: int) -> T: ...

  def __setitem__(self, index: int, value: T) -> None: ...

  def __len__(self) -> int: ...

  def __iter__(self) -> Iterator[T]: ...


class Matrix[S: SupportsInt](Protocol):

  def __getitem__(self, index: int) -> MatrixRow[S]: ...

  def __len__(self) -> int: ...

  def __iter__(self) -> Iterator[MatrixRow[S]]: ...

Now all of these work!

matrix3x3: Matrix[int]
matrix3x3 = [[1,2,3],[4,5,6],[7,8,9]]
matrix3x3 = np.array([[1,2,3],[4,5,6],[7,8,9]])
matrix3x3 = np.ndarray(shape=(3, 3), buffer=np.array([[1,2,3],[4,5,6],[7,8,9]]))
for row in matrix3x3:
  for val in row:
    print(val)
print(len(matrix3x3), len(matrix3x3[0]))

And all of these raise errors!

matrix3x3.append([10,11,12])
matrix3x3[2].append(10)
matrix3x3.pop()
matrix3x3[0].pop()
matrix3x3[0][0] = "one"

And even if some of those implementations are intrinsically mutable in size, our type system lets mypy catch any bug where matrix3x3 is reshaped. Unfortunately, there's no way to prevent someone from assigning a 4x7 matrix to matrix3x3, but at least it's named clearly. Maybe someday there will be Python support for fixed-size lists as types.