typ: Fix mpoly_p typing protocol

Author: oscarbenjamin

src/flint/types/fmpq_mpoly.pyi

@@ -28,7 +28,7 @@ class fmpq_mpoly_ctx(flint_mpoly_context[fmpq_mpoly, fmpq, ifmpq]):
 
 class fmpq_mpoly(flint_mpoly[fmpq_mpoly_ctx, fmpq, ifmpq]):
     def __init__(self,
-        val: fmpq_mpoly | fmpz_mpoly | ifmpq | dict[tuple[int, ...], ifmpq] | _str = 0,
+        val: fmpq_mpoly | fmpz_mpoly | ifmpq | Mapping[tuple[int, ...], ifmpq] | _str = 0,
         ctx: fmpq_mpoly_ctx | None = None
     ) -> None: ...
 
@@ -53,7 +53,7 @@ class fmpq_mpoly(flint_mpoly[fmpq_mpoly_ctx, fmpq, ifmpq]):
     def __getitem__(self, index: tuple[int, ...]) -> fmpq: ...
     def __setitem__(self, index: tuple[int, ...], coeff: ifmpq) -> None: ...
 
-    def subs(self, mapping: dict[_str | int, ifmpq]) -> fmpq_mpoly: ...
+    def subs(self, mapping: Mapping[_str | int, ifmpq]) -> fmpq_mpoly: ...
     def compose(self, *args: fmpq_mpoly, ctx: fmpq_mpoly_ctx | None = None) -> fmpq_mpoly: ...
 
     def __call__(self, *args: ifmpq) -> fmpq: ...

src/flint/types/fmpz_mod_mpoly.pyi

@@ -31,7 +31,7 @@ class fmpz_mod_mpoly_ctx(flint_mod_mpoly_context[fmpz_mod_mpoly, fmpz_mod, ifmpz
 
 class fmpz_mod_mpoly(flint_mpoly[fmpz_mod_mpoly_ctx, fmpz_mod, ifmpz_mod]):
     def __init__(self,
-        val: fmpz_mod_mpoly | fmpz_mpoly | ifmpz_mod | dict[tuple[int, ...], ifmpz_mod] | _str = 0,
+        val: fmpz_mod_mpoly | fmpz_mpoly | ifmpz_mod | Mapping[tuple[int, ...], ifmpz_mod] | _str = 0,
         ctx: fmpz_mod_mpoly_ctx | None = None
     ) -> None: ...
 
@@ -56,7 +56,7 @@ class fmpz_mod_mpoly(flint_mpoly[fmpz_mod_mpoly_ctx, fmpz_mod, ifmpz_mod]):
     def __getitem__(self, index: tuple[int, ...]) -> fmpz_mod: ...
     def __setitem__(self, index: tuple[int, ...], coeff: ifmpz_mod) -> None: ...
 
-    def subs(self, mapping: dict[_str | int, ifmpz_mod]) -> fmpz_mod_mpoly: ...
+    def subs(self, mapping: Mapping[_str | int, ifmpz_mod]) -> fmpz_mod_mpoly: ...
     def compose(self, *args: fmpz_mod_mpoly, ctx: fmpz_mod_mpoly_ctx | None = None) -> fmpz_mod_mpoly: ...
 
     def __call__(self, *args: ifmpz_mod) -> fmpz_mod: ...

src/flint/types/fmpz_mpoly.pyi

@@ -27,7 +27,7 @@ class fmpz_mpoly_ctx(flint_mpoly_context[fmpz_mpoly, fmpz, ifmpz]):
 class fmpz_mpoly(flint_mpoly[fmpz_mpoly_ctx, fmpz, ifmpz]):
 
     def __init__(self,
-        val: fmpz_mpoly | ifmpz | dict[tuple[int, ...], ifmpz] | _str = 0,
+        val: fmpz_mpoly | ifmpz | Mapping[tuple[int, ...], ifmpz] | _str = 0,
         ctx: fmpz_mpoly_ctx | None = None
     ) -> None: ...
 
@@ -52,7 +52,7 @@ class fmpz_mpoly(flint_mpoly[fmpz_mpoly_ctx, fmpz, ifmpz]):
     def __getitem__(self, index: tuple[int, ...]) -> fmpz: ...
     def __setitem__(self, index: tuple[int, ...], coeff: ifmpz) -> None: ...
 
-    def subs(self, mapping: dict[_str | int, ifmpz]) -> fmpz_mpoly: ...
+    def subs(self, mapping: Mapping[_str | int, ifmpz]) -> fmpz_mpoly: ...
     def compose(self, *args: fmpz_mpoly, ctx: fmpz_mpoly_ctx | None = None) -> fmpz_mpoly: ...
 
     def __call__(self, *args: ifmpz) -> fmpz: ...

src/flint/types/nmod_mpoly.pyi

@@ -32,7 +32,7 @@ class nmod_mpoly_ctx(flint_mpoly_context[nmod_mpoly, nmod, inmod]):
 
 class nmod_mpoly(flint_mpoly[nmod_mpoly_ctx, nmod, inmod]):
     def __init__(self,
-        val: nmod_mpoly | fmpz_mpoly | inmod | dict[tuple[int, ...], inmod] | _str = 0,
+        val: nmod_mpoly | fmpz_mpoly | inmod | Mapping[tuple[int, ...], inmod] | _str = 0,
         ctx: nmod_mpoly_ctx | None = None
     ) -> None: ...
 
@@ -57,7 +57,7 @@ class nmod_mpoly(flint_mpoly[nmod_mpoly_ctx, nmod, inmod]):
     def __getitem__(self, index: tuple[int, ...]) -> nmod: ...
     def __setitem__(self, index: tuple[int, ...], coeff: inmod) -> None: ...
 
-    def subs(self, mapping: dict[_str | int, inmod]) -> nmod_mpoly: ...
+    def subs(self, mapping: Mapping[_str | int, inmod]) -> nmod_mpoly: ...
     def compose(self, *args: nmod_mpoly, ctx: nmod_mpoly_ctx | None = None) -> nmod_mpoly: ...
 
     def __call__(self, *args: inmod) -> nmod: ...

src/flint/typing.py

@@ -1,6 +1,8 @@
 #
 #
 #
+from __future__ import annotations
+
 from typing import (
     TYPE_CHECKING,
     Protocol,
@@ -39,9 +41,9 @@
     "elem_p",
     "scalar_p",
     "poly_p",
-    # "mpoly_p",
-    # "mpoly_context_p",
     # "series_p",
+    "mpoly_p",
+    "mpoly_context_p",
 
     # Types
     "Ordering",
@@ -50,16 +52,8 @@
 ]
 
 
-if TYPE_CHECKING:
-    import flint as _flint
-
-
 _Telem = TypeVar("_Telem", bound=flint_scalar)
-_Telem_co = TypeVar("_Telem_co", bound=flint_scalar, covariant=True)
-_Telem_coerce = TypeVar("_Telem_coerce")
-_Telem_coerce_contra = TypeVar("_Telem_coerce_contra", contravariant=True)
-_Tmpoly = TypeVar("_Tmpoly", bound=flint_mpoly, covariant=True)
-_Tctx = TypeVar("_Tctx", bound=flint_mpoly_context)
+_Telem_contra = TypeVar("_Telem_contra", bound=flint_scalar, contravariant=True)
 _Sctx = TypeVar("_Sctx", bound=flint_mpoly_context)
 
 _str = str
@@ -149,101 +143,106 @@ def factor_squarefree(self) -> tuple[_Tscalar, list[tuple[Self, int]]]: ...
     def deflation(self) -> tuple[Self, int]: ...
 
 
-class _mpoly_p(elem_p, Protocol[_Tctx, _Telem, _Telem_coerce]):
-    """FLINT multivariate polynomial Protocol."""
-    def __init__(
-        self,
-        val: Self
-        | _Telem
-        | _Telem_coerce
-        | int
-        | dict[tuple[int, ...], _Telem | _Telem_coerce | int]
-        | str = 0,
-        ctx: _Tctx | None = None,
-    ) -> None: ...
-    def str(self) -> _str: ...
-    def repr(self) -> _str: ...
-    def context(self) -> _Tctx: ...
-    def degrees(self) -> tuple[int, ...]: ...
-    def total_degree(self) -> int: ...
-    def leading_coefficient(self) -> _Telem: ...
-    def to_dict(self) -> dict[tuple[int, ...], _Telem]: ...
-    def is_one(self) -> bool: ...
-    def is_zero(self) -> bool: ...
-    def is_constant(self) -> bool: ...
-    def __len__(self) -> int: ...
-    def __getitem__(self, index: tuple[int, ...]) -> _Telem: ...
-    def __setitem__(
-        self, index: tuple[int, ...], coeff: _Telem | _Telem_coerce | int
-    ) -> None: ...
-    def __iter__(self) -> Iterable[tuple[int, ...]]: ...
-    def __contains__(self, index: tuple[int, ...]) -> bool: ...
-    def coefficient(self, i: int) -> _Telem: ...
-    def monomial(self, i: int) -> tuple[int, ...]: ...
-    def terms(self) -> Iterable[tuple[tuple[int, ...], _Telem]]: ...
-    def monoms(self) -> list[tuple[int, ...]]: ...
+class _series_p(elem_p, Protocol[_Telem]):
+    """FLINT univariate power series."""
+
+    def __iter__(self) -> Iterator[_Telem]: ...
     def coeffs(self) -> list[_Telem]: ...
-    def __pos__(self) -> Self: ...
-    def __neg__(self) -> Self: ...
-    def __add__(self, other: Self | _Telem | _Telem_coerce | int) -> Self: ...
-    def __radd__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __sub__(self, other: Self | _Telem | _Telem_coerce | int) -> Self: ...
-    def __rsub__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __mul__(self, other: Self | _Telem | _Telem_coerce | int) -> Self: ...
-    def __rmul__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __truediv__(self, other: Self | _Telem | _Telem_coerce | int) -> Self: ...
-    def __rtruediv__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __floordiv__(self, other: Self | _Telem | _Telem_coerce | int) -> Self: ...
-    def __rfloordiv__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __mod__(self, other: Self | _Telem | _Telem_coerce | int) -> Self: ...
-    def __rmod__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __divmod__(
-        self, other: Self | _Telem | _Telem_coerce | int
-    ) -> tuple[Self, Self]: ...
-    def __rdivmod__(self, other: _Telem | _Telem_coerce | int) -> tuple[Self, Self]: ...
-    def __pow__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def __rpow__(self, other: _Telem | _Telem_coerce | int) -> Self: ...
-    def iadd(self, other: _Telem | _Telem_coerce | int) -> None: ...
-    def isub(self, other: _Telem | _Telem_coerce | int) -> None: ...
-    def imul(self, other: _Telem | _Telem_coerce | int) -> None: ...
-    def gcd(self, other: Self) -> Self: ...
-    def term_content(self) -> Self: ...
-    def factor(self) -> tuple[_Telem, Sequence[tuple[Self, int]]]: ...
-    def factor_squarefree(self) -> tuple[_Telem, Sequence[tuple[Self, int]]]: ...
-    def sqrt(self) -> Self: ...
-    def resultant(self, other: Self, var: _str | int) -> Self: ...
-    def discriminant(self, var: _str | int) -> Self: ...
-    def deflation_index(self) -> tuple[list[int], list[int]]: ...
-    def deflation(self) -> tuple[Self, list[int]]: ...
-    def deflation_monom(self) -> tuple[Self, list[int], Self]: ...
-    def inflate(self, N: list[int]) -> Self: ...
-    def deflate(self, N: list[int]) -> Self: ...
-    def subs(self, mapping: dict[_str | int, _Telem | _Telem_coerce | int]) -> Self: ...
-    def compose(self, *args: Self, ctx: _Tctx | None = None) -> Self: ...
-    def __call__(self, *args: _Telem | _Telem_coerce) -> _Telem: ...
-    def derivative(self, var: _str | int) -> Self: ...
-    def unused_gens(self) -> tuple[_str, ...]: ...
-    def project_to_context(
-        self, other_ctx: _Tctx, mapping: dict[_str | int, _str | int] | None = None
-    ) -> Self: ...
-
-
-class _mpoly_context_p(
-    elem_p, Protocol[_Tmpoly, _Telem_co, _Telem_coerce_contra]
-):
-    """FLINT multivariate polynomial context protocol."""
+
+
+class mpoly_p(elem_p, Protocol[_Telem]):
+   def __init__(
+       self,
+       val: Self
+       | _Telem
+       | int
+       | Mapping[tuple[int, ...], _Telem | int]
+       | str = 0,
+       ctx: Any | None = None,
+   ) -> None: ...
+   def str(self) -> _str: ...
+   def repr(self) -> _str: ...
+   def context(self) -> mpoly_context_p[Self, _Telem]: ...
+   def degrees(self) -> tuple[int, ...]: ...
+   def total_degree(self) -> int: ...
+   def leading_coefficient(self) -> _Telem: ...
+   def to_dict(self) -> dict[tuple[int, ...], _Telem]: ...
+   def is_one(self) -> bool: ...
+   def is_zero(self) -> bool: ...
+   def is_constant(self) -> bool: ...
+   def __len__(self) -> int: ...
+   def __getitem__(self, index: tuple[int, ...]) -> _Telem: ...
+   def __setitem__(
+       self, index: tuple[int, ...], coeff: _Telem | int
+   ) -> None: ...
+   def __iter__(self) -> Iterable[tuple[int, ...]]: ...
+   def __contains__(self, index: tuple[int, ...]) -> bool: ...
+   def coefficient(self, i: int) -> _Telem: ...
+   def monomial(self, i: int) -> tuple[int, ...]: ...
+   def terms(self) -> Iterable[tuple[tuple[int, ...], _Telem]]: ...
+   def monoms(self) -> list[tuple[int, ...]]: ...
+   def coeffs(self) -> list[_Telem]: ...
+   def __pos__(self) -> Self: ...
+   def __neg__(self) -> Self: ...
+   def __add__(self, other: Self | _Telem | int) -> Self: ...
+   def __radd__(self, other: _Telem | int) -> Self: ...
+   def __sub__(self, other: Self | _Telem | int) -> Self: ...
+   def __rsub__(self, other: _Telem | int) -> Self: ...
+   def __mul__(self, other: Self | _Telem | int) -> Self: ...
+   def __rmul__(self, other: _Telem | int) -> Self: ...
+   def __truediv__(self, other: Self | _Telem | int) -> Self: ...
+   def __rtruediv__(self, other: _Telem | int) -> Self: ...
+   def __floordiv__(self, other: Self | _Telem | int) -> Self: ...
+   def __rfloordiv__(self, other: _Telem | int) -> Self: ...
+   def __mod__(self, other: Self | _Telem | int) -> Self: ...
+   def __rmod__(self, other: _Telem | int) -> Self: ...
+   def __divmod__(
+       self, other: Self | _Telem | int
+   ) -> tuple[Self, Self]: ...
+   def __rdivmod__(self, other: _Telem | int) -> tuple[Self, Self]: ...
+   def __pow__(self, other: _Telem | int) -> Self: ...
+   def __rpow__(self, other: _Telem | int) -> Self: ...
+   def iadd(self, other: _Telem | int) -> None: ...
+   def isub(self, other: _Telem | int) -> None: ...
+   def imul(self, other: _Telem | int) -> None: ...
+   def gcd(self, other: Self) -> Self: ...
+   def term_content(self) -> Self: ...
+   def factor(self) -> tuple[_Telem, Sequence[tuple[Self, int]]]: ...
+   def factor_squarefree(self) -> tuple[_Telem, Sequence[tuple[Self, int]]]: ...
+   def sqrt(self) -> Self: ...
+   def resultant(self, other: Self, var: _str | int) -> Self: ...
+   def discriminant(self, var: _str | int) -> Self: ...
+   def deflation_index(self) -> tuple[list[int], list[int]]: ...
+   def deflation(self) -> tuple[Self, list[int]]: ...
+   def deflation_monom(self) -> tuple[Self, list[int], Self]: ...
+   def inflate(self, N: list[int]) -> Self: ...
+   def deflate(self, N: list[int]) -> Self: ...
+   def subs(self, mapping: Mapping[_str | int, _Telem | int]) -> Self: ...
+   def compose(self, *args: Self, ctx: Any | None = None) -> Self: ...
+   def __call__(self, *args: _Telem) -> _Telem: ...
+   def derivative(self, var: _str | int) -> Self: ...
+   def unused_gens(self) -> tuple[_str, ...]: ...
+   def project_to_context(
+       self, other_ctx: Any, mapping: dict[_str | int, _str | int] | None = None
+   ) -> Self: ...
+
+
+_Tmpoly_p = TypeVar("_Tmpoly_p", bound=mpoly_p, covariant=True)
+
+
+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: ...
-    def from_dict(self, d: Mapping[tuple[int, ...], _Telem_coerce_contra], /) -> _Tmpoly: ...
-    def constant(self, z: _Telem_coerce_contra, /) -> _Tmpoly: ...
+    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 name(self, i: int, /) -> str: ...
     def names(self) -> tuple[str]: ...
-    def gens(self) -> tuple[_Tmpoly, ...]: ...
+    def gens(self) -> tuple[_Tmpoly_p, ...]: ...
     def variable_to_index(self, var: str, /) -> int: ...
     def term(
-        self, coeff: _Telem_coerce_contra | None = None, exp_vec: Iterable[int] | None = None
-    ) -> _Tmpoly: ...
+        self, coeff: _Telem_contra | None = None, exp_vec: Iterable[int] | None = None
+    ) -> _Tmpoly_p: ...
     def drop_gens(self, gens: Iterable[str | int], /) -> Self: ...
     def append_gens(self, gens: Iterable[str | int], /) -> Self: ...
     def infer_generator_mapping(
@@ -258,31 +257,60 @@ def from_context(
     ) -> _Sctx: ...
 
 
-class _series_p(elem_p, Protocol[_Telem]):
-    """FLINT univariate power series."""
+if TYPE_CHECKING:
 
-    def __iter__(self) -> Iterator[_Telem]: ...
-    def coeffs(self) -> list[_Telem]: ...
+    import flint as _f
 
+    _x1: scalar_p = _f.fmpz(1)
+    _x2: scalar_p = _f.fmpq(1, 2)
+    _x3: scalar_p = _f.nmod(1, 2)
+    _x4: scalar_p = _f.fmpz_mod(1, _f.fmpz_mod_ctx(2))
+    _x5: scalar_p = _f.fq_default(1, _f.fq_default_ctx(2))
+    # XXX: Add arf, acf, arb, acb, ...
 
-if TYPE_CHECKING:
+    _Tscalar0 = TypeVar("_Tscalar0", bound=scalar_p)
 
-    _x1: scalar_p = _flint.fmpz(1)
-    _x2: scalar_p = _flint.fmpq(1, 2)
-    _x3: scalar_p = _flint.nmod(1, 2)
-    _x4: scalar_p = _flint.fmpz_mod(1, _flint.fmpz_mod_ctx(2))
-    _x5: scalar_p = _flint.fq_default(1, _flint.fq_default_ctx(2))
-    # XXX: Add arf, acf, arb, acb, ...
+    def _scalar(x: _Tscalar0, y: _Tscalar0) -> _Tscalar0:
+        z = x + y + 1
+        z = z / z + x**2
+        z = 1 / z + 1**z
+        return z
 
-    _y1: poly_p[_flint.fmpz] = _flint.fmpz_poly([1, 2])
-    _y2: poly_p[_flint.fmpq] = _flint.fmpq_poly([1, 2])
-    _y3: poly_p[_flint.nmod] = _flint.nmod_poly(1, 2)
-    _y4: poly_p[_flint.fmpz_mod] = _flint.fmpz_mod_poly(1, _flint.fmpz_mod_poly_ctx(2))
-    _y5: poly_p[_flint.fq_default] = _flint.fq_default_poly(1, _flint.fq_default_poly_ctx(2))
+    _y1: poly_p[_f.fmpz] = _f.fmpz_poly([1, 2])
+    _y2: poly_p[_f.fmpq] = _f.fmpq_poly([1, 2])
+    _y3: poly_p[_f.nmod] = _f.nmod_poly(1, 2)
+    _y4: poly_p[_f.fmpz_mod] = _f.fmpz_mod_poly(1, _f.fmpz_mod_poly_ctx(2))
+    _y5: poly_p[_f.fq_default] = _f.fq_default_poly(1, _f.fq_default_poly_ctx(2))
     # XXX: Add arb_poly, acb_poly, ...
 
-    _z1: epoly_p[_flint.fmpz] = _flint.fmpz_poly([1, 2])
-    _z2: epoly_p[_flint.fmpq] = _flint.fmpq_poly([1, 2])
-    _z3: epoly_p[_flint.nmod] = _flint.nmod_poly(1, 2)
-    _z4: epoly_p[_flint.fmpz_mod] = _flint.fmpz_mod_poly(1, _flint.fmpz_mod_poly_ctx(2))
-    _z5: epoly_p[_flint.fq_default] = _flint.fq_default_poly(1, _flint.fq_default_poly_ctx(2))
+    def _poly(x: epoly_p[_Tscalar0], y: epoly_p[_Tscalar0]) -> epoly_p[_Tscalar0]:
+        z = x + y + 1
+        z = z / z + x**2
+        z = 1 / z + z.coeffs()[0]
+        return z
+
+    _z1: epoly_p[_f.fmpz] = _f.fmpz_poly([1, 2])
+    _z2: epoly_p[_f.fmpq] = _f.fmpq_poly([1, 2])
+    _z3: epoly_p[_f.nmod] = _f.nmod_poly(1, 2)
+    _z4: epoly_p[_f.fmpz_mod] = _f.fmpz_mod_poly(1, _f.fmpz_mod_poly_ctx(2))
+    _z5: epoly_p[_f.fq_default] = _f.fq_default_poly(1, _f.fq_default_poly_ctx(2))
+
+    def _epoly(x: epoly_p[_Tscalar0], y: epoly_p[_Tscalar0]) -> epoly_p[_Tscalar0]:
+        return x.gcd(y).factor()[1][0][0]
+
+    _ctx1: mpoly_context_p[_f.fmpz_mpoly, _f.fmpz] = _f.fmpz_mpoly_ctx.get(['x', 'y'])
+    _ctx2: mpoly_context_p[_f.fmpq_mpoly, _f.fmpq] = _f.fmpq_mpoly_ctx.get(['x', 'y'])
+    _ctx3: mpoly_context_p[_f.nmod_mpoly, _f.nmod] = _f.nmod_mpoly_ctx.get(['x', 'y'], modulus=7)
+    _ctx4: mpoly_context_p[_f.fmpz_mod_mpoly, _f.fmpz_mod] = _f.fmpz_mod_mpoly_ctx.get(['x', 'y'], modulus=7)
+
+    _p1: mpoly_p[_f.fmpz] = _ctx1.gens()[0]
+    _p2: mpoly_p[_f.fmpq] = _ctx2.gens()[0]
+    _p3: mpoly_p[_f.nmod] = _ctx3.gens()[0]
+    _p4: mpoly_p[_f.fmpz_mod] = _ctx4.gens()[0]
+
+    _Tmpoly0 = TypeVar("_Tmpoly0", bound=mpoly_p)
+
+    def _mpoly(x: mpoly_p[_Telem], y: mpoly_p[_Telem], ctx: mpoly_context_p[_Tmpoly0, _Telem]) -> mpoly_p[_Telem]:
+        z = x.gcd(y).factor()[1][0][0] + ctx.gens()[0]
+        t = z.compose(z, ctx=ctx) + z.context().gens()[0]
+        return t