fmpz_mat: add type stubs

Author: oscarbenjamin

src/flint/flint_base/flint_base.pyi

@@ -87,6 +87,16 @@ class flint_poly(flint_elem, Generic[Telem]):
     def complex_roots(self) -> list[Any]: ...
     def derivative(self) -> Self: ...
 
+class flint_mat(flint_elem, Generic[Telem]):
+    def nrows(self) -> int: ...
+    def ncols(self) -> int: ...
+    def __getitem__(self, index: tuple[int, int], /) -> Telem: ...
+    def __setitem__(self, index: tuple[int, int], value: Telem | int, /) -> None: ...
+    def entries(self) -> list[Telem]: ...
+    def table(self) -> list[list[Telem]]: ...
+    def tolist(self) -> list[list[Telem]]: ...
+    def __iter__(self) -> Iterator[Telem]: ...
+
 class Ordering(enum.Enum):
     lex = "lex"
     deglex = "deglex"

src/flint/types/fmpz_mat.pyi

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator, Sequence
+from typing import Literal, overload
+
+from flint.flint_base.flint_base import flint_mat
+from flint.types.fmpq import fmpq
+from flint.types.fmpz import fmpz, ifmpz
+from flint.types.fmpz_poly import fmpz_poly
+
+_str = str
+
+ifmpz_mat = fmpz_mat | Sequence[Sequence[ifmpz]]
+
+
+class fmpz_mat(flint_mat[fmpz]):
+    @overload
+    def __init__(self, val: ifmpz_mat, /) -> None: ...
+    @overload
+    def __init__(self, m: int, n: int, /) -> None: ...
+    @overload
+    def __init__(self, m: int, n: int, entries: Iterable[ifmpz], /) -> None: ...
+
+    def __bool__(self) -> bool: ...
+    def __eq__(self, other: object, /) -> bool: ...
+    def __ne__(self, other: object, /) -> bool: ...
+
+    def __getitem__(self, index: tuple[int, int], /) -> fmpz: ...
+    def __setitem__(self, index: tuple[int, int], value: ifmpz, /) -> None: ...
+    def __iter__(self) -> Iterator[fmpz]: ...
+    def entries(self) -> list[fmpz]: ...
+    def table(self) -> list[list[fmpz]]: ...
+    def tolist(self) -> list[list[fmpz]]: ...
+
+    def det(self) -> fmpz: ...
+    def __pos__(self) -> fmpz_mat: ...
+    def __neg__(self) -> fmpz_mat: ...
+    def __add__(self, other: fmpz_mat, /) -> fmpz_mat: ...
+    def __sub__(self, other: fmpz_mat, /) -> fmpz_mat: ...
+    @overload
+    def __mul__(self, other: fmpz_mat | ifmpz, /) -> fmpz_mat: ...
+    @overload
+    def __mul__(self, other: fmpq, /) -> flint_mat: ...
+    @overload
+    def __rmul__(self, other: ifmpz, /) -> fmpz_mat: ...
+    @overload
+    def __rmul__(self, other: fmpq, /) -> flint_mat: ...
+    def __truediv__(self, other: ifmpz, /) -> fmpz_mat: ...
+    def __pow__(self, exponent: int, modulo: None = None, /) -> fmpz_mat: ...
+
+    def is_square(self) -> bool: ...
+    def is_empty(self) -> bool: ...
+    def is_zero(self) -> bool: ...
+    def is_one(self) -> bool: ...
+    def is_neg_one(self) -> bool: ...
+    def is_upper_triangular(self) -> bool: ...
+    def is_lower_triangular(self) -> bool: ...
+    def is_diagonal(self) -> bool: ...
+    def is_scalar(self) -> bool: ...
+    def is_hadamard(self) -> bool: ...
+    def is_hnf(self) -> bool: ...
+    def is_snf(self) -> bool: ...
+
+    @classmethod
+    def hadamard(cls, n: int, /) -> fmpz_mat: ...
+    @classmethod
+    def randtest(cls, m: int, n: int, bits: int, /) -> fmpz_mat: ...
+    @classmethod
+    def randbits(cls, m: int, n: int, bits: int, /) -> fmpz_mat: ...
+    @classmethod
+    def randrank(cls, m: int, n: int, rank: int, bits: int, /) -> fmpz_mat: ...
+
+    def rank(self) -> int: ...
+    @overload
+    def inv(self, integer: Literal[False] = False, /) -> flint_mat: ...
+    @overload
+    def inv(self, integer: Literal[True], /) -> fmpz_mat: ...
+    def transpose(self) -> fmpz_mat: ...
+    @overload
+    def solve(self, other: fmpz_mat, integer: Literal[False] = False, /) -> flint_mat: ...
+    @overload
+    def solve(self, other: fmpz_mat, integer: Literal[True], /) -> fmpz_mat: ...
+    def _fflu(self) -> tuple[fmpz_mat, fmpz, list[int], int]: ...
+    def fflu(self) -> tuple[fmpz_mat, fmpz_mat, fmpz_mat, fmpz_mat]: ...
+    def rref(self, inplace: bool = False, /) -> tuple[fmpz_mat, fmpz, int]: ...
+    def nullspace(self) -> tuple[fmpz_mat, int]: ...
+    @overload
+    def lll(
+        self,
+        transform: Literal[False] = False,
+        delta: float = 0.99,
+        eta: float = 0.51,
+        rep: str = "zbasis",
+        gram: str = "approx",
+        /,
+    ) -> fmpz_mat: ...
+    @overload
+    def lll(
+        self,
+        transform: Literal[True],
+        delta: float = 0.99,
+        eta: float = 0.51,
+        rep: str = "zbasis",
+        gram: str = "approx",
+        /,
+    ) -> tuple[fmpz_mat, fmpz_mat]: ...
+    @overload
+    def hnf(self, transform: Literal[False] = False, /) -> fmpz_mat: ...
+    @overload
+    def hnf(self, transform: Literal[True], /) -> tuple[fmpz_mat, fmpz_mat]: ...
+    def snf(self) -> fmpz_mat: ...
+    def charpoly(self) -> fmpz_poly: ...
+    def minpoly(self) -> fmpz_poly: ...
+
+    def str(self, *args: object, **kwargs: object) -> _str: ...
+    def repr(self) -> _str: ...
+    def __str__(self) -> _str: ...
+    def __repr__(self) -> _str: ...

src/flint/types/meson.build

@@ -7,7 +7,7 @@ pyfiles = [
   'fmpz_poly.pyi',
   'fmpz_series.pyi',
   'fmpz_mpoly.pyi',
-  # 'fmpz_mat.pyi',
+  'fmpz_mat.pyi',
   'fmpz_vec.pyi',
 
   'fmpq.pyi',