From 5deccad943fe196496303054078a268f79a128d1 Mon Sep 17 00:00:00 2001
From: Arham Chopra <arham.chopra@cubistsystematic.com>
Date: Thu, 13 Jun 2024 12:19:12 -0400
Subject: [PATCH] Add postprocess_to_dict hook for to_dict method in structs

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>
---
 cpp/csp/python/PyStructToDict.cpp |  8 ++++
 csp/impl/struct.py                | 10 ++++
 csp/tests/impl/test_struct.py     | 78 +++++++++++++++++++++++++++++++
 3 files changed, 96 insertions(+)

diff --git a/cpp/csp/python/PyStructToDict.cpp b/cpp/csp/python/PyStructToDict.cpp
index 83ea191a..9f521276 100644
--- 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
index bcdab845..88910edf 100644
--- 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
index db0fd852..fe5bd288 100644
--- a/csp/tests/impl/test_struct.py
+++ b/csp/tests/impl/test_struct.py
@@ -1376,6 +1376,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