dev(feat): add pickle for constraints in TreeValue

test/tree/tree/base.py

@@ -701,4 +701,31 @@ def get_treevalue_test(treevalue_class: Type[TreeValue]):
             t3 = t1.with_constraints([GreaterThanConstraint(10), int], clear=True)
             assert t3.constraint.equiv([int, GreaterThanConstraint(10)])
 
+        def test_pickle_constraints(self):
+            t1 = get_demo_constraint_tree()
+            assert t1.a == 21
+            assert t1.b.x == 'f-49'
+            assert t1.b.y == pytest.approx(7.7)
+            assert t1.constraint.equiv([
+                object, {
+                    'a': [int, GreaterThanConstraint(3)],
+                    'b': {'x': [cleaf(), str], 'y': float}
+                }
+            ])
+
+            binary = pickle.dumps(t1)
+            newt1 = pickle.loads(binary)
+            assert newt1.a == 21
+            assert newt1.b.x == 'f-49'
+            assert newt1.b.y == pytest.approx(7.7)
+            assert newt1.constraint.equiv([
+                object, {
+                    'a': [int, GreaterThanConstraint(3)],
+                    'b': {'x': [cleaf(), str], 'y': float}
+                }
+            ])
+
+            assert newt1 == t1
+            assert newt1.constraint == t1.constraint
+
     return _TestClass

treevalue/tree/tree/tree.pyx

@@ -792,7 +792,7 @@ cdef class TreeValue:
             return False
 
     @cython.binding(True)
-    def __setstate__(self, TreeStorage state):
+    def __setstate__(self, tuple state):
         """
         Deserialize operation, can support `pickle.loads`.
 
@@ -807,7 +807,7 @@ cdef class TreeValue:
             >>> bin_ = pickle.dumps(t)  # dump it to binary
             >>> pickle.loads(bin_)      #  TreeValue({'a': 1, 'b': 2, 'x': {'c': 3}})
         """
-        self._st = state
+        self._st, self.constraint = state
 
     @cython.binding(True)
     def __getstate__(self):
@@ -824,7 +824,7 @@ cdef class TreeValue:
             >>> bin_ = pickle.dumps(t)  # dump it to binary
             >>> pickle.loads(bin_)      #  TreeValue({'a': 1, 'b': 2, 'x': {'c': 3}})
         """
-        return self._st
+        return self._st, self.constraint
 
     @cython.binding(True)
     cpdef treevalue_keys keys(self):