Merge pull request #280 from Point72/ac/add_postprocess_dict_hook

Add postprocess_to_dict hook for to_dict method in structs
Point72 · Jul 8, 2024 · a198332 · a198332
2 parents b3385f2 + 5deccad
commit a198332
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 0 deletions.
diff --git a/cpp/csp/python/PyStructToDict.cpp b/cpp/csp/python/PyStructToDict.cpp
@@ -108,6 +108,14 @@ PyObjectPtr parseStructToDictRecursive( const StructPtr& self, PyObject * callab
             } );
         PyDict_SetItemString( new_dict.get(), key.c_str(), py_obj.get() );
     }
+
+    // Optional postprocess hook in python to allow caller to customize to_dict behavior for struct
+    PyObject * py_type = ( PyObject * ) meta -> pyType();
+    if( PyObject_HasAttrString( py_type, "postprocess_to_dict" ) )
+    {
+        auto postprocess_dict_callable = PyObjectPtr::own( PyObject_GetAttrString( py_type, "postprocess_to_dict" ) );
+        new_dict = PyObjectPtr::check( PyObject_CallFunction( postprocess_dict_callable.get(), "(O)", new_dict.get() ) );
+    }
     return new_dict;
 }
 

diff --git a/csp/impl/struct.py b/csp/impl/struct.py
@@ -165,6 +165,16 @@ def to_dict_depr(self):
         res = self._obj_to_python(self)
         return res
 
+    @classmethod
+    def postprocess_to_dict(self, obj):
+        """Postprocess hook for to_dict method
+
+        This method is invoked by to_dict after converting a struct to a dict
+        as an additional hook for users to modify the dict before it is returned
+        by the to_dict method
+        """
+        return obj
+
     def to_dict(self, callback=None):
         """Create a dictionary representation of the struct
 

diff --git a/csp/tests/impl/test_struct.py b/csp/tests/impl/test_struct.py
@@ -1385,6 +1385,84 @@ class A(csp.Struct):
         r = repr(a)
         self.assertTrue(repr(raw) in r)
 
+    def test_to_dict_recursion(self):
+        class MyStruct(csp.Struct):
+            l1: list
+            l2: list
+            d1: dict
+            d2: dict
+            t1: tuple
+            t2: tuple
+
+        test_struct = MyStruct(l1=[1], l2=[2])
+        result_dict = {"l1": [1], "l2": [2]}
+        self.assertEqual(test_struct.to_dict(), result_dict)
+
+        test_struct = MyStruct(l1=[1], l2=[2])
+        test_struct.l1.append(test_struct.l2)
+        test_struct.l2.append(test_struct.l1)
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+        test_struct = MyStruct(l1=[1])
+        test_struct.l1.append(test_struct.l1)
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+        test_struct = MyStruct(l1=[1])
+        test_struct.l1.append(test_struct)
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+        test_struct = MyStruct(d1={1: 1}, d2={2: 2})
+        result_dict = {"d1": {1: 1}, "d2": {2: 2}}
+        self.assertEqual(test_struct.to_dict(), result_dict)
+
+        test_struct = MyStruct(d1={1: 1}, d2={2: 2})
+        test_struct.d1["d2"] = test_struct.d2
+        test_struct.d2["d1"] = test_struct.d1
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+        test_struct = MyStruct(d1={1: 1}, d2={2: 2})
+        test_struct.d1["d1"] = test_struct.d1
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+        test_struct = MyStruct(d1={1: 1}, d2={2: 2})
+        test_struct.d1["d1"] = test_struct
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+        test_struct = MyStruct(t1=(1, 1), t2=(2, 2))
+        result_dict = {"t1": (1, 1), "t2": (2, 2)}
+        self.assertEqual(test_struct.to_dict(), result_dict)
+
+        test_struct = MyStruct(t1=(1, 1))
+        test_struct.t1 = (1, 2, test_struct)
+        with self.assertRaises(RecursionError):
+            test_struct.to_dict()
+
+    def test_to_dict_postprocess(self):
+        class MySubStruct(csp.Struct):
+            i: int = 0
+
+            def postprocess_to_dict(obj):
+                obj["postprocess_called"] = True
+                return obj
+
+        class MyStruct(csp.Struct):
+            i: int = 1
+            mss: MySubStruct = MySubStruct()
+
+            def postprocess_to_dict(obj):
+                obj["postprocess_called"] = True
+                return obj
+
+        test_struct = MyStruct()
+        result_dict = {"i": 1, "postprocess_called": True, "mss": {"i": 0, "postprocess_called": True}}
+        self.assertEqual(test_struct.to_dict(), result_dict)
+
     def test_to_json_primitives(self):
         class MyStruct(csp.Struct):
             b: bool = True