Make fmpq_mpoly and fmpz_mpoly more dict-like

Author: Jake-Moss

src/flint/flint_base/flint_base.pyx

@@ -248,7 +248,7 @@ cdef class flint_mpoly(flint_elem):
         """
         Returns True if `self` contains a term with exponent vector `x` and a non-zero coefficient.
 
-            >>> from flint import Ordering
+            >>> from flint import fmpq_mpoly_ctx, Ordering
             >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
             >>> p = ctx.from_dict({(0, 1): 2, (1, 1): 3})
             >>> (1, 1) in p
@@ -259,6 +259,25 @@ cdef class flint_mpoly(flint_elem):
         """
         return bool(self[x])
 
+    def __iter__(self):
+        return iter(self.keys())
+
+    def __pos__(self):
+        return self
+
+    def items(self):
+        """
+        Return the exponent vectors and coefficient of each term.
+
+            >>> from flint import fmpq_mpoly_ctx, Ordering
+            >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
+            >>> f = ctx.from_dict({(0, 0): 1, (1, 0): 2, (0, 1): 3, (1, 1): 4})
+            >>> list(f.items())
+            [((1, 1), 4), ((1, 0), 2), ((0, 1), 3), ((0, 0), 1)]
+
+        """
+        return zip(self.keys(), self.values())
+
 
 cdef class flint_series(flint_elem):
     """

src/flint/test/test_all.py

@@ -2790,6 +2790,11 @@ def quick_poly():
         assert p(0, 0) == p(0, S(0)) == p(S(0), S(0)) == S(1) == 1
         assert p(1, 1) == S(10) == 10
 
+        p = quick_poly()
+        assert p.keys() == [(2, 2), (1, 0), (0, 1), (0, 0)]
+        assert p.values() == [4, 3, 2, 1]
+        assert list(p.items()) == list(zip([(2, 2), (1, 0), (0, 1), (0, 0)], [4, 3, 2, 1]))
+
         assert p.subs({"x1": S(0), "x0": S(0)}) == ctx.from_dict({(0, 0): 1})
         assert p.compose(p.subs({"x1": 0}), ctx.from_dict({(0, 1): 1})) == mpoly({
             (2, 2): 36,

src/flint/types/fmpq_mpoly.pyx

@@ -317,9 +317,6 @@ cdef class fmpq_mpoly(flint_mpoly):
         exp_vec = fmpz_vec(x, double_indirect=True)
         fmpq_mpoly_set_coeff_fmpq_fmpz(self.val, (<fmpq>coeff).val, exp_vec.double_indirect, self.ctx.val)
 
-    def __pos__(self):
-        return self
-
     def __neg__(self):
         cdef fmpq_mpoly res
         res = create_fmpq_mpoly(self.ctx)
@@ -587,6 +584,51 @@ cdef class fmpq_mpoly(flint_mpoly):
             raise ValueError("Unreasonably large polynomial")  # pragma: no cover
         return vres
 
+    def keys(self):
+        """
+        Return the exponent vectors of each term as a tuple of fmpz.
+
+            >>> from flint import Ordering
+            >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
+            >>> f = ctx.from_dict({(0, 0): 1, (1, 0): 2, (0, 1): 3, (1, 1): 4})
+            >>> f.keys()
+            [(1, 1), (1, 0), (0, 1), (0, 0)]
+
+        """
+        cdef:
+            slong i, nvars = self.ctx.nvars()
+            fmpz_vec vec = fmpz_vec(nvars, double_indirect=True)
+
+        res = []
+        for i in range(len(self)):
+            fmpq_mpoly_get_term_exp_fmpz(vec.double_indirect, self.val, i, self.ctx.val)
+            res.append(vec.to_tuple())
+
+        return res
+
+    def values(self):
+        """
+        Return the coefficients of each term as a fmpq.
+
+            >>> from flint import Ordering
+            >>> ctx = fmpq_mpoly_ctx.get_context(2, Ordering.lex, 'x')
+            >>> f = ctx.from_dict({(0, 0): 1, (1, 0): 2, (0, 1): 3, (1, 1): 4})
+            >>> f.values()
+            [4, 2, 3, 1]
+
+        """
+        cdef:
+            fmpq coeff
+            slong i
+
+        res = []
+        for i in range(len(self)):
+            coeff = fmpq.__new__(fmpq)
+            fmpq_mpoly_get_term_coeff_fmpq(coeff.val, self.val, i, self.ctx.val)
+            res.append(coeff)
+
+        return res
+
     def subs(self, dict_args) -> fmpq_mpoly:
         """
         Partial evaluate this polynomial.

src/flint/types/fmpz_mpoly.pyx

@@ -297,9 +297,6 @@ cdef class fmpz_mpoly(flint_mpoly):
         exp_vec = fmpz_vec(x, double_indirect=True)
         fmpz_mpoly_set_coeff_fmpz_fmpz(self.val, (<fmpz>coeff).val, exp_vec.double_indirect, self.ctx.val)
 
-    def __pos__(self):
-        return self
-
     def __neg__(self):
         cdef fmpz_mpoly res
         res = create_fmpz_mpoly(self.ctx)
@@ -569,6 +566,51 @@ cdef class fmpz_mpoly(flint_mpoly):
             raise ValueError("Unreasonably large polynomial")  # pragma: no cover
         return vres
 
+    def keys(self):
+        """
+        Return the exponent vectors of each term as a tuple of fmpz.
+
+            >>> from flint import Ordering
+            >>> ctx = fmpz_mpoly_ctx.get_context(2, Ordering.lex, 'x')
+            >>> f = ctx.from_dict({(0, 0): 1, (1, 0): 2, (0, 1): 3, (1, 1): 4})
+            >>> f.keys()
+            [(1, 1), (1, 0), (0, 1), (0, 0)]
+
+        """
+        cdef:
+            slong i, nvars = self.ctx.nvars()
+            fmpz_vec vec = fmpz_vec(nvars, double_indirect=True)
+
+        res = []
+        for i in range(len(self)):
+            fmpz_mpoly_get_term_exp_fmpz(vec.double_indirect, self.val, i, self.ctx.val)
+            res.append(vec.to_tuple())
+
+        return res
+
+    def values(self):
+        """
+        Return the coefficients of each term as a fmpz
+
+            >>> from flint import Ordering
+            >>> ctx = fmpz_mpoly_ctx.get_context(2, Ordering.lex, 'x')
+            >>> f = ctx.from_dict({(0, 0): 1, (1, 0): 2, (0, 1): 3, (1, 1): 4})
+            >>> f.values()
+            [4, 2, 3, 1]
+
+        """
+        cdef:
+            fmpz coeff
+            slong i
+
+        res = []
+        for i in range(len(self)):
+            coeff = fmpz.__new__(fmpz)
+            fmpz_mpoly_get_term_coeff_fmpz(coeff.val, self.val, i, self.ctx.val)
+            res.append(coeff)
+
+        return res
+
     def subs(self, dict_args) -> fmpz_mpoly:
         """
         Partial evaluate this polynomial with select constants. All arguments must be fmpz.