From cc927e7234def2b487a4222913b3096f0c223157 Mon Sep 17 00:00:00 2001
From: HansBug <hansbug@buaa.edu.cn>
Date: Mon, 5 Dec 2022 00:49:41 +0800
Subject: [PATCH] dev(feat): add pickle for constraints in TreeValue

---
 test/tree/tree/base.py       | 27 +++++++++++++++++++++++++++
 treevalue/tree/tree/tree.pyx |  6 +++---
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/test/tree/tree/base.py b/test/tree/tree/base.py
index 12fbad1239..0b37f15dfc 100644
--- a/test/tree/tree/base.py
+++ b/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
diff --git a/treevalue/tree/tree/tree.pyx b/treevalue/tree/tree/tree.pyx
index cd18178ea2..f4f5e0b194 100644
--- a/treevalue/tree/tree/tree.pyx
+++ b/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):
-- 
GitLab