typ: fix mpoly_p typing protocol

Author: oscarbenjamin

src/flint/flint_base/flint_base.pyi

@@ -36,13 +36,13 @@ class flint_scalar(flint_elem):
     def is_zero(self) -> bool: ...
     def __pos__(self) -> Self: ...
     def __neg__(self) -> Self: ...
-    def __add__(self, other: int, /) -> Self: ...
+    def __add__(self, other: Self | int, /) -> Self: ...
     def __radd__(self, other: int, /) -> Self: ...
-    def __sub__(self, other: int, /) -> Self: ...
+    def __sub__(self, other: Self | int, /) -> Self: ...
     def __rsub__(self, other: int, /) -> Self: ...
-    def __mul__(self, other: int, /) -> Self: ...
+    def __mul__(self, other: Self | int, /) -> Self: ...
     def __rmul__(self, other: int, /) -> Self: ...
-    def __truediv__(self, other: int, /) -> Self: ...
+    def __truediv__(self, other: Self | int, /) -> Self: ...
     def __rtruediv__(self, other: int, /) -> Self: ...
     def __pow__(self, other: int, /) -> Self: ...
     def __rpow__(self, other: int, /) -> Self: ...
@@ -193,14 +193,14 @@ class flint_mpoly_context(flint_elem, Generic[Tmpoly, Telem, Telem_coerce]):
     @classmethod
     def from_context(
         cls,
-        ctx: Sctx,
+        ctx: flint_mpoly_context,
         names: str | Iterable[str | tuple[str, int]] | tuple[str, int] | None = None,
         ordering: Ordering | str = Ordering.lex,
-    ) -> Sctx: ...
+    ) -> Self: ...
 
 class flint_mod_mpoly_context(flint_mpoly_context[Tmpoly, Telem, Telem_coerce]):
     @abstractmethod
-    def modulus(self) -> int: ...
+    def modulus(self): ...
 
 class flint_series(flint_elem, Generic[Telem]):
     """Base class for power series."""

src/flint/test/test_all.py

@@ -32,7 +32,9 @@ def raises(f, exception) -> bool:
 Tscalar = TypeVar('Tscalar', bound=flint_base.flint_scalar)
 Tscalar_co = TypeVar('Tscalar_co', bound=flint_base.flint_scalar, covariant=True)
 Tmpoly = TypeVar('Tmpoly', bound=flint_base.flint_mpoly)
+Tmpoly_p = TypeVar('Tmpoly_p', bound=typ.mpoly_p)
 Tmpolyctx_co = TypeVar('Tmpolyctx_co', bound=flint_base.flint_mpoly_context, covariant=True)
+Tmpolyctx_p_co = TypeVar('Tmpolyctx_p_co', bound=typ.mpoly_context_p, covariant=True)
 
 
 _default_ctx_string = """\
@@ -3093,17 +3095,17 @@ def _all_mpolys(): # -> _all_mpolys_type:
     )
 
 
-class _GetMPolyCtx(Protocol[Tmpolyctx_co]):
+class _GetMPolyCtx(Protocol[Tmpolyctx_p_co]):
     def __call__(self,
         names: Iterable[str | tuple[str, int]] | tuple[str, int],
         ordering: str | flint.Ordering = "lex"
-    ) -> Tmpolyctx_co:
+    ) -> Tmpolyctx_p_co:
         ...
 
 
 _MPolyTestCase = tuple[
-    type[Tmpoly],
-    _GetMPolyCtx['flint_base.flint_mpoly_context[Tmpoly, Tscalar, int]'],
+    type[Tmpoly_p],
+    _GetMPolyCtx[typ.mpoly_context_p[Tmpoly_p, Tscalar]],
     Callable[[int], Tscalar],
     bool,
     flint.fmpz
@@ -3176,7 +3178,7 @@ def wrapper():
 
 
 @all_mpolys
-def test_mpolys_constructor(args: _MPolyTestCase[Tmpoly, Tscalar]) -> None:
+def test_mpolys_constructor(args: _MPolyTestCase[Tmpoly_p, Tscalar]) -> None:
     P, get_context, S, _, _ = args
 
     ctx = get_context(("x", 2))
@@ -3301,7 +3303,7 @@ def quick_poly():
 
 
 @all_mpolys
-def test_mpolys_properties(args: _MPolyTestCase[Tmpoly, Tscalar]) -> None:
+def test_mpolys_properties(args: _MPolyTestCase[Tmpoly_p, Tscalar]) -> None:
 
     P, get_context, S, is_field, characteristic = args
 

src/flint/types/fmpq_mpoly.pyi

@@ -1,4 +1,4 @@
-from typing import Iterable, Mapping
+from typing import Iterable, Mapping, Any
 from ..flint_base.flint_base import flint_mpoly, flint_mpoly_context, Ordering
 from .fmpz import fmpz
 from .fmpq import fmpq
@@ -20,11 +20,28 @@ class fmpq_mpoly_ctx(flint_mpoly_context[fmpq_mpoly, fmpq, ifmpq]):
 
     def nvars(self) -> int: ...
     def ordering(self) -> Ordering: ...
-
-    def gen(self, i: int) -> fmpq_mpoly: ...
-    def from_dict(self, d: Mapping[tuple[int, ...], ifmpq]) -> fmpq_mpoly: ...
-    def constant(self, z: ifmpq) -> fmpq_mpoly: ...
-
+    def gen(self, i: int, /) -> fmpq_mpoly: ...
+    def from_dict(self, d: Mapping[tuple[int, ...], ifmpq], /) -> fmpq_mpoly: ...
+    def constant(self, q: ifmpq, /) -> fmpq_mpoly: ...
+    def name(self, i: int, /) -> str: ...
+    def names(self) -> tuple[str]: ...
+    def gens(self) -> tuple[fmpq_mpoly, ...]: ...
+    def variable_to_index(self, var: str, /) -> int: ...
+    def term(
+        self, coeff: ifmpq | None = None, exp_vec: Iterable[int] | None = None
+    ) -> fmpq_mpoly: ...
+    def drop_gens(self, gens: Iterable[str | int], /) -> fmpq_mpoly_ctx: ...
+    def append_gens(self, gens: Iterable[str | int], /) -> fmpq_mpoly_ctx: ...
+    def infer_generator_mapping(
+        self, ctx: flint_mpoly_context, /
+    ) -> dict[int, int]: ...
+    @classmethod
+    def from_context(
+        cls,
+        ctx: flint_mpoly_context[Any, Any, Any],
+        names: str | Iterable[str | tuple[str, int]] | tuple[str, int] | None = None,
+        ordering: Ordering | str = Ordering.lex,
+    ) -> fmpq_mpoly_ctx: ...
 
 class fmpq_mpoly(flint_mpoly[fmpq_mpoly_ctx, fmpq, ifmpq]):
     def __init__(self,

src/flint/types/fmpz_mod_mpoly.pyi

@@ -1,5 +1,10 @@
-from typing import Iterable, Mapping
-from ..flint_base.flint_base import flint_mpoly, flint_mod_mpoly_context, Ordering
+from typing import Iterable, Mapping, Any
+from ..flint_base.flint_base import (
+    flint_mpoly,
+    flint_mpoly_context,
+    flint_mod_mpoly_context,
+    Ordering,
+)
 from .fmpz import fmpz
 from .fmpz_mod import fmpz_mod
 from .fmpz_mpoly import fmpz_mpoly
@@ -17,16 +22,38 @@ class fmpz_mod_mpoly_ctx(flint_mod_mpoly_context[fmpz_mod_mpoly, fmpz_mod, ifmpz
         names: _str | Iterable[_str | tuple[_str, int]] | tuple[_str, int],
         ordering: Ordering | _str = Ordering.lex,
         *,
-        modulus: ifmpz,
+        modulus: int,
     ) -> fmpz_mod_mpoly_ctx:
         ...
 
+    def modulus(self) -> fmpz: ...
+
     def nvars(self) -> int: ...
     def ordering(self) -> Ordering: ...
-    def modulus(self) -> int: ...
-    def gen(self, i: int, /) -> fmpz_mod_mpoly: ...
-    def constant(self, z: ifmpz_mod) -> fmpz_mod_mpoly: ...
+
+    def gen(self, i: int) -> fmpz_mod_mpoly: ...
     def from_dict(self, d: Mapping[tuple[int, ...], ifmpz_mod]) -> fmpz_mod_mpoly: ...
+    def constant(self, z: ifmpz_mod) -> fmpz_mod_mpoly: ...
+
+    def name(self, i: int, /) -> str: ...
+    def names(self) -> tuple[str]: ...
+    def gens(self) -> tuple[fmpz_mod_mpoly, ...]: ...
+    def variable_to_index(self, var: str, /) -> int: ...
+    def term(
+        self, coeff: ifmpz_mod | None = None, exp_vec: Iterable[int] | None = None
+    ) -> fmpz_mod_mpoly: ...
+    def drop_gens(self, gens: Iterable[str | int], /) -> fmpz_mod_mpoly_ctx: ...
+    def append_gens(self, gens: Iterable[str | int], /) -> fmpz_mod_mpoly_ctx: ...
+    def infer_generator_mapping(
+        self, ctx: flint_mpoly_context, /
+    ) -> dict[int, int]: ...
+    @classmethod
+    def from_context(
+        cls,
+        ctx: flint_mpoly_context[Any, Any, Any],
+        names: str | Iterable[str | tuple[str, int]] | tuple[str, int] | None = None,
+        ordering: Ordering | str = Ordering.lex,
+    ) -> fmpz_mod_mpoly_ctx: ...
 
 
 class fmpz_mod_mpoly(flint_mpoly[fmpz_mod_mpoly_ctx, fmpz_mod, ifmpz_mod]):

src/flint/types/fmpz_mpoly.pyi

@@ -1,4 +1,4 @@
-from typing import Iterable, Iterator, Mapping
+from typing import Iterable, Iterator, Mapping, Any
 
 from ..flint_base.flint_base import flint_mpoly, flint_mpoly_context, Ordering
 from .fmpz import fmpz
@@ -18,11 +18,28 @@ class fmpz_mpoly_ctx(flint_mpoly_context[fmpz_mpoly, fmpz, ifmpz]):
 
     def nvars(self) -> int: ...
     def ordering(self) -> Ordering: ...
-
-    def gen(self, i: int) -> fmpz_mpoly: ...
-    def from_dict(self, d: Mapping[tuple[int, ...], ifmpz]) -> fmpz_mpoly: ...
-    def constant(self, z: ifmpz) -> fmpz_mpoly: ...
-
+    def gen(self, i: int, /) -> fmpz_mpoly: ...
+    def from_dict(self, d: Mapping[tuple[int, ...], ifmpz], /) -> fmpz_mpoly: ...
+    def constant(self, z: ifmpz, /) -> fmpz_mpoly: ...
+    def name(self, i: int, /) -> str: ...
+    def names(self) -> tuple[str]: ...
+    def gens(self) -> tuple[fmpz_mpoly, ...]: ...
+    def variable_to_index(self, var: str, /) -> int: ...
+    def term(
+        self, coeff: ifmpz | None = None, exp_vec: Iterable[int] | None = None
+    ) -> fmpz_mpoly: ...
+    def drop_gens(self, gens: Iterable[str | int], /) -> fmpz_mpoly_ctx: ...
+    def append_gens(self, gens: Iterable[str | int], /) -> fmpz_mpoly_ctx: ...
+    def infer_generator_mapping(
+        self, ctx: flint_mpoly_context, /
+    ) -> dict[int, int]: ...
+    @classmethod
+    def from_context(
+        cls,
+        ctx: flint_mpoly_context[Any, Any, Any],
+        names: str | Iterable[str | tuple[str, int]] | tuple[str, int] | None = None,
+        ordering: Ordering | str = Ordering.lex,
+    ) -> fmpz_mpoly_ctx: ...
 
 class fmpz_mpoly(flint_mpoly[fmpz_mpoly_ctx, fmpz, ifmpz]):
 

src/flint/types/nmod_mpoly.pyi

@@ -1,5 +1,10 @@
-from typing import Iterable, Mapping, Self
-from ..flint_base.flint_base import flint_mpoly, flint_mpoly_context, Ordering
+from typing import Iterable, Mapping, Self, Any
+from ..flint_base.flint_base import (
+    flint_mpoly,
+    flint_mpoly_context,
+    flint_mod_mpoly_context,
+    Ordering,
+)
 from .fmpz import fmpz
 from .fmpq import fmpq
 from .nmod import nmod
@@ -11,7 +16,7 @@ ifmpz = int | fmpz
 inmod = int | fmpz | fmpq | nmod
 
 
-class nmod_mpoly_ctx(flint_mpoly_context[nmod_mpoly, nmod, inmod]):
+class nmod_mpoly_ctx(flint_mod_mpoly_context[nmod_mpoly, nmod, inmod]):
 
     @classmethod
     def get(cls,
@@ -22,13 +27,35 @@ class nmod_mpoly_ctx(flint_mpoly_context[nmod_mpoly, nmod, inmod]):
     ) -> nmod_mpoly_ctx:
         ...
 
+    def modulus(self) -> int: ...
+
     def nvars(self) -> int: ...
     def ordering(self) -> Ordering: ...
 
     def gen(self, i: int) -> nmod_mpoly: ...
     def from_dict(self, d: Mapping[tuple[int, ...], inmod]) -> nmod_mpoly: ...
     def constant(self, z: inmod) -> nmod_mpoly: ...
 
+    def name(self, i: int, /) -> str: ...
+    def names(self) -> tuple[str]: ...
+    def gens(self) -> tuple[nmod_mpoly, ...]: ...
+    def variable_to_index(self, var: str, /) -> int: ...
+    def term(
+        self, coeff: inmod | None = None, exp_vec: Iterable[int] | None = None
+    ) -> nmod_mpoly: ...
+    def drop_gens(self, gens: Iterable[str | int], /) -> nmod_mpoly_ctx: ...
+    def append_gens(self, gens: Iterable[str | int], /) -> nmod_mpoly_ctx: ...
+    def infer_generator_mapping(
+        self, ctx: flint_mpoly_context, /
+    ) -> dict[int, int]: ...
+    @classmethod
+    def from_context(
+        cls,
+        ctx: flint_mpoly_context[Any, Any, Any],
+        names: str | Iterable[str | tuple[str, int]] | tuple[str, int] | None = None,
+        ordering: Ordering | str = Ordering.lex,
+    ) -> nmod_mpoly_ctx: ...
+
 
 class nmod_mpoly(flint_mpoly[nmod_mpoly_ctx, nmod, inmod]):
     def __init__(self,

src/flint/typing.py

@@ -234,8 +234,8 @@ class mpoly_context_p(elem_p, Protocol[_Tmpoly_p, _Telem_contra]):
     def nvars(self) -> int: ...
     def ordering(self) -> Ordering: ...
     def gen(self, i: int, /) -> _Tmpoly_p: ...
-    def from_dict(self, d: Mapping[tuple[int, ...], _Telem_contra], /) -> _Tmpoly_p: ...
-    def constant(self, z: _Telem_contra, /) -> _Tmpoly_p: ...
+    def from_dict(self, d: Mapping[tuple[int, ...], _Telem_contra | int], /) -> _Tmpoly_p: ...
+    def constant(self, z: _Telem_contra | int, /) -> _Tmpoly_p: ...
     def name(self, i: int, /) -> str: ...
     def names(self) -> tuple[str]: ...
     def gens(self) -> tuple[_Tmpoly_p, ...]: ...
@@ -246,15 +246,15 @@ def term(
     def drop_gens(self, gens: Iterable[str | int], /) -> Self: ...
     def append_gens(self, gens: Iterable[str | int], /) -> Self: ...
     def infer_generator_mapping(
-        self, ctx: flint_mpoly_context, /
+        self, ctx: Self, /
     ) -> dict[int, int]: ...
     @classmethod
     def from_context(
         cls,
-        ctx: _Sctx,
+        ctx: flint_mpoly_context,
         names: str | Iterable[str | tuple[str, int]] | tuple[str, int] | None = None,
         ordering: Ordering | str = Ordering.lex,
-    ) -> _Sctx: ...
+    ) -> Self: ...
 
 
 if TYPE_CHECKING: