Trac #32990: pickling support for real balls and related objects

URL: https://trac.sagemath.org/32990
Reported by: mmezzarobba
Ticket author(s): Marc Mezzarobba
Reviewer(s): Matthias Koeppe

Author: Release Manager

src/sage/libs/arb/arb.pxd

@@ -42,6 +42,9 @@ cdef extern from "arb_wrap.h":
     void arb_printd(const arb_t x, long digits)
     void arb_printn(const arb_t x, long digits, unsigned long flags)
 
+    char *arb_dump_str(const arb_t x)
+    bint arb_load_str(arb_t x, const char *str)
+
     # void arb_randtest(arb_t x, flint_rand_t state, long prec, long mag_bits)
     # void arb_randtest_exact(arb_t x, flint_rand_t state, long prec, long mag_bits)
     # void arb_randtest_precise(arb_t x, flint_rand_t state, long prec, long mag_bits)

src/sage/matrix/matrix_complex_ball_dense.pyx

@@ -18,6 +18,10 @@ TESTS::
     sage: pol(mat)
     [8.000000000000000 11.00000000000000]
     [22.00000000000000 41.00000000000000]
+
+    sage: mat = matrix(ComplexBallField(20), 2, 2, list(range(4)))*i/3
+    sage: loads(dumps(mat)).identical(mat)
+    True
 """
 
 #*****************************************************************************

src/sage/rings/complex_arb.pyx

@@ -331,6 +331,9 @@ class ComplexBallField(UniqueRepresentation, sage.rings.abc.ComplexBallField):
 
         sage: ComplexBallField().is_finite()
         False
+
+        sage: loads(dumps(ComplexBallField(60))) is ComplexBallField(60)
+        True
     """
     Element = ComplexBall
 
@@ -1487,6 +1490,18 @@ cdef class ComplexBall(RingElement):
             return "{} + {}*I".format(self.real()._repr_(),
                                         self.imag()._repr_())
 
+    def __reduce__(self):
+        r"""
+        Serialize a ComplexBall.
+
+        TESTS::
+
+            sage: [loads(dumps(b)).identical(b) for b in
+            ....:     [ComplexBallField(60)(1/3 + i*pi), CBF(NaN)]]
+            [True, True]
+        """
+        return self.__class__, (self._parent, self.real(), self.imag())
+
     def _is_atomic(self):
         r"""
         Declare that complex balls print atomically in some cases.

src/sage/rings/polynomial/polynomial_complex_arb.pyx

@@ -184,6 +184,23 @@ cdef class Polynomial_complex_arb(Polynomial):
                 ball = Coeff(x)
                 acb_poly_set_coeff_acb(self.__poly, 0, ball.value)
 
+    def __reduce__(self):
+        r"""
+        Serialize a polynomial for pickling.
+
+        TESTS::
+
+            sage: Pol.<x> = ComplexBallField(42)[]
+            sage: pol = (x + i)/3
+            sage: pol2 = loads(dumps(pol))
+            sage: pol.degree() == pol2.degree()
+            True
+            sage: all(a.identical(b) for (a, b) in zip(pol, pol2))
+            True
+        """
+        return (self.__class__,
+               (self.parent(), self.list(), False, self.is_gen()))
+
     # Access
 
     def degree(self):

src/sage/rings/real_arb.pyx

@@ -369,6 +369,9 @@ class RealBallField(UniqueRepresentation, sage.rings.abc.RealBallField):
 
         sage: RealBallField().is_finite()
         False
+
+        sage: loads(dumps(RealBallField(60))) is RealBallField(60)
+        True
     """
     Element = RealBall
 
@@ -1148,6 +1151,30 @@ cdef inline bint _do_sig(long prec):
 cdef inline long prec(RealBall ball):
     return ball._parent._prec
 
+def create_RealBall(parent, serialized):
+    r"""
+    Create a RealBall from a serialized representation.
+
+    TESTS::
+
+        sage: from sage.rings.real_arb import create_RealBall
+        sage: create_RealBall(RBF, b'15555555555555 -36 1 -36')
+        [0.3333333333333333 +/- 7.04e-17]
+        sage: create_RealBall(RBF, b'foo')
+        Traceback (most recent call last):
+        ...
+        ValueError: incorrect format
+    """
+    cdef RealBall res = RealBall.__new__(RealBall)
+    res._parent = parent
+    sig_on()
+    cdef bint error = arb_load_str(res.value, serialized)
+    sig_off()
+    if error:
+        raise ValueError("incorrect format")
+    else:
+        return res
+
 cdef class RealBall(RingElement):
     """
     Hold one ``arb_t`` of the `Arb library
@@ -1494,6 +1521,26 @@ cdef class RealBall(RingElement):
 
         return py_string
 
+    def __reduce__(self):
+        r"""
+        Serialize a RealBall.
+
+        TESTS::
+
+            sage: [loads(dumps(b)).identical(b) for b in
+            ....:     [RealBallField(60).pi(), RBF(infinity), RBF(NaN)]]
+            [True, True, True]
+        """
+        cdef bytes py_val
+        sig_on()
+        cdef char* c_val = arb_dump_str(self.value)
+        sig_off()
+        try:
+            py_val = <bytes> c_val
+        finally:
+            flint_free(c_val)
+        return create_RealBall, (self._parent, py_val)
+
     # Conversions
 
     cpdef RealIntervalFieldElement _real_mpfi_(self, RealIntervalField_class parent):