nmod_mat.pyx: add type annotations

Author: oscarbenjamin

src/flint/test/test_all.py

@@ -1594,6 +1594,7 @@ def test_nmod_mat():
     assert str(M(2,2,[1,2,3,4],17)) == '[1, 2]\n[3, 4]'
     assert repr(M(2,2,[1,2,3,4],17)) == '[1, 2]\n[3, 4]'
     assert M(1,2,[3,4],17) / 3 == M(1,2,[3,4],17) * (~G(3,17))
+    assert raises(lambda: pow(M([[1]],17), 2, 3), NotImplementedError)
     assert M(2,2,[1,2,3,4], 17).inv().det() == ~(M(2,2,[1,2,3,4], 17).det())
     assert M(2,2,[1,2,3,4], 17).inv().inv() == M(2,2,[1,2,3,4], 17)
     assert M(2,2,[0,1,2,3],17) * M(2, 2, [2,3,4,5], 17) == M(2,2,[4,5,16,4],17)
@@ -1611,6 +1612,11 @@ def test_nmod_mat():
     assert (M([[1]],17) != M([[1]],13)) is True
     assert (M([[1]],17) == None) is False
     assert (M([[1]],17) != None) is True
+    class _nmod_mat_subclass(M):
+        pass
+
+    assert raises(lambda: M([[1]],3) + _nmod_mat_subclass([[1]],5), ValueError)
+    assert raises(lambda: M([[1]],3) - _nmod_mat_subclass([[1]],5), ValueError)
     M2 = M.randtest(3,4,5)
     assert all(0 <= int(x) < 5 for x in M2.entries())
     assert (M2.nrows(), M2.ncols()) == (3, 4)

src/flint/types/meson.build

@@ -20,7 +20,7 @@ pyfiles = [
   'nmod.pyi',
   'nmod_poly.pyi',
   'nmod_mpoly.pyi',
-  # 'nmod_mat.pyi',
+  'nmod_mat.pyi',
   # 'nmod_series.pyi',
 
   'fmpz_mod.pyi',

src/flint/types/nmod_mat.pyi

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator, Sequence
+from typing import overload
+
+from flint.flint_base.flint_base import flint_mat
+from flint.types.fmpz import fmpz
+from flint.types.fmpz_mat import fmpz_mat
+from flint.types.nmod import inmod, nmod
+from flint.types.nmod_poly import nmod_poly
+
+_str = str
+
+inmod_mat = nmod_mat | fmpz_mat | Sequence[Sequence[inmod]]
+
+
+class nmod_mat(flint_mat[nmod]):
+    @overload
+    def __init__(self, val: nmod_mat, /) -> None: ...
+    @overload
+    def __init__(self, val: fmpz_mat | Sequence[Sequence[inmod]], mod: int, /) -> None: ...
+    @overload
+    def __init__(self, m: int, n: int, mod: int, /) -> None: ...
+    @overload
+    def __init__(self, m: int, n: int, entries: Iterable[inmod], mod: int, /) -> None: ...
+
+    def __bool__(self) -> bool: ...
+    def __eq__(self, other: object, /) -> bool: ...
+    def __ne__(self, other: object, /) -> bool: ...
+
+    def nrows(self) -> int: ...
+    def ncols(self) -> int: ...
+    def modulus(self) -> int: ...
+    @classmethod
+    def randtest(cls, m: int, n: int, mod: int, /) -> nmod_mat: ...
+
+    def __getitem__(self, index: tuple[int, int], /) -> nmod: ...
+    def __setitem__(self, index: tuple[int, int], value: inmod, /) -> None: ...
+    def __iter__(self) -> Iterator[nmod]: ...
+    def entries(self) -> list[nmod]: ...
+    def table(self) -> list[list[nmod]]: ...
+    def tolist(self) -> list[list[nmod]]: ...
+
+    def repr(self) -> _str: ...
+    def __str__(self) -> _str: ...
+    def __repr__(self) -> _str: ...
+
+    def __pos__(self) -> nmod_mat: ...
+    def __neg__(self) -> nmod_mat: ...
+    def __add__(self, other: inmod_mat, /) -> nmod_mat: ...
+    def __radd__(self, other: inmod_mat, /) -> nmod_mat: ...
+    def __sub__(self, other: inmod_mat, /) -> nmod_mat: ...
+    def __rsub__(self, other: inmod_mat, /) -> nmod_mat: ...
+    def __mul__(self, other: inmod_mat | inmod, /) -> nmod_mat: ...
+    def __rmul__(self, other: inmod_mat | inmod, /) -> nmod_mat: ...
+    def __truediv__(self, other: inmod, /) -> nmod_mat: ...
+    def __pow__(self, exponent: int, modulo: None = None, /) -> nmod_mat: ...
+
+    def det(self) -> nmod: ...
+    def inv(self) -> nmod_mat: ...
+    def transpose(self) -> nmod_mat: ...
+    def solve(self, other: inmod_mat, /) -> nmod_mat: ...
+    def rref(self, inplace: bool = False, /) -> tuple[nmod_mat, int]: ...
+    def rank(self) -> int: ...
+    def nullspace(self) -> tuple[nmod_mat, int]: ...
+    def charpoly(self) -> nmod_poly: ...
+    def minpoly(self) -> nmod_poly: ...

src/flint/types/nmod_mat.pyx

@@ -358,9 +358,6 @@ cdef class nmod_mat(flint_mat):
     def __truediv__(s, t):
         return nmod_mat._div_(s, t)
 
-    def __div__(s, t):
-        return nmod_mat._div_(s, t)
-
     def det(self):
         """
         Returns the determinant of self as an nmod.