Add support for additional python types as dict keys in Struct.to_json (#321)

* Add support for additional python types as dict keys in Struct.to_json

Signed-off-by: Arham Chopra <[email protected]>

* Fix memory leaks, add docs, and address comments

Signed-off-by: Arham Chopra <[email protected]>

---------

Signed-off-by: Arham Chopra <[email protected]>

Author: arhamchopra

cpp/csp/python/PyStructToJson.cpp

@@ -198,12 +198,40 @@ rapidjson::Value toJsonRecursive( const StructPtr& self, rapidjson::Document& do
     return new_dict;
 }
 
-rapidjson::Value pyDictKeyToName( PyObject * py_key, rapidjson::Document& doc )
+rapidjson::Value pyDictKeyToName( PyObject * py_key, rapidjson::Document& doc, PyObject * callable )
 {
-    // Only support strings, ints, and floats as keys
+    // NOTE: Only support None, bool, strings, ints, and floats, date, time, datetime, enums, csp.Enums as keys
     // JSON encoding requires all names to be strings so convert them to strings
+
+    static thread_local PyTypeObjectPtr s_tl_enum_type;
+    // Get the enum type on the first call and save it for future use
+    if( s_tl_enum_type.get() == nullptr ) [[unlikely]]
+    {
+        // Import enum module to extract the Enum type
+        auto py_enum_module = PyObjectPtr::own( PyImport_ImportModule( "enum" ) );
+        if( py_enum_module.get() )
+        {
+            s_tl_enum_type = PyTypeObjectPtr::own( reinterpret_cast<PyTypeObject*>( PyObject_GetAttrString( py_enum_module.get(), "Enum" ) ) );
+        }
+        else
+        {
+            CSP_THROW( RuntimeException, "Unable to import enum module from the python standard library" );
+        }
+    }
+
     rapidjson::Value val;
-    if( PyUnicode_Check( py_key ) )
+    if( py_key == Py_None )
+    {
+        val.SetString( "null" );
+    }
+    else if( PyBool_Check( py_key ) )
+    {
+        auto str_obj = PyObjectPtr::own( PyObject_Str( py_key ) );
+        Py_ssize_t len = 0;
+        const char * str = PyUnicode_AsUTF8AndSize( str_obj.get(), &len );
+        val.SetString( str, len, doc.GetAllocator() );
+    }
+    else if( PyUnicode_Check( py_key ) )
     {
         Py_ssize_t len;
         auto str = PyUnicode_AsUTF8AndSize( py_key, &len );
@@ -220,9 +248,8 @@ rapidjson::Value pyDictKeyToName( PyObject * py_key, rapidjson::Document& doc )
         auto json_obj = doubleToJson( key, doc );
         if ( json_obj.IsNull() )
         {
-            auto * str_obj = PyObject_Str( py_key );
-            Py_ssize_t len = 0;
-            const char * str = PyUnicode_AsUTF8AndSize( str_obj, &len );
+            auto str_obj = PyObjectPtr::own( PyObject_Str( py_key ) );
+            const char * str = PyUnicode_AsUTF8( str_obj.get() );
             CSP_THROW( ValueError, "Cannot serialize " + std::string( str ) + " to key in JSON" );
         }
         else
@@ -233,6 +260,35 @@ rapidjson::Value pyDictKeyToName( PyObject * py_key, rapidjson::Document& doc )
             val.SetString( s.str(), doc.GetAllocator() );
         }
     }
+    else if( PyTime_CheckExact( py_key ) )
+    {
+        auto v = fromPython<Time>( py_key );
+        val = toJson( v, CspType( CspType::Type::TIME ), doc, callable );
+    }
+    else if( PyDate_CheckExact( py_key ) )
+    {
+        auto v = fromPython<Date>( py_key );
+        val = toJson( v, CspType( CspType::Type::DATE ), doc, callable );
+    }
+    else if( PyDateTime_CheckExact( py_key ) )
+    {
+        auto v = fromPython<DateTime>( py_key );
+        val = toJson( v, CspType( CspType::Type::DATETIME ), doc, callable );
+    }
+    else if( PyType_IsSubtype( Py_TYPE( py_key ), &PyCspEnum::PyType ) )
+    {
+        auto enum_ptr = static_cast<PyCspEnum *>( py_key ) -> enum_;
+        val = toJson( enum_ptr, CspType( CspType::Type::ENUM ), doc, callable );
+    }
+    else if( PyType_IsSubtype( Py_TYPE( py_key ), s_tl_enum_type.get() ) )
+    {
+        // Use the `name` attribute of the enum for the string representation
+        auto py_enum_name = PyObjectPtr::own( PyObject_GetAttrString( py_key, "name" ) );
+        auto str_obj = PyObjectPtr::own( PyObject_Str( py_enum_name.get() ) );
+        Py_ssize_t len = 0;
+        const char * str = PyUnicode_AsUTF8AndSize( str_obj.get(), &len );
+        val.SetString( str, len, doc.GetAllocator() );
+    }
     else
     {
         CSP_THROW( ValueError, "Cannot serialize key of type: " + std::string( Py_TYPE( py_key ) -> tp_name ) );
@@ -284,7 +340,7 @@ rapidjson::Value pyDictToJson( PyObject * py_dict, rapidjson::Document& doc, PyO
 
     while( PyDict_Next( py_dict, &pos, &py_key, &py_value ) )
     {
-        auto key = pyDictKeyToName( py_key, doc );
+        auto key = pyDictKeyToName( py_key, doc, callable );
         auto res = pyObjectToJson( py_value, doc, callable, false );
         new_dict.AddMember( key, res, doc.GetAllocator() );
     }

csp/tests/impl/test_struct.py

@@ -1,3 +1,4 @@
+import enum
 import json
 import numpy as np
 import pickle
@@ -1742,6 +1743,7 @@ class MyStruct(csp.Struct):
         result_dict = {"i": 456, "d_any": d_any_res}
         self.assertEqual(json.loads(test_struct.to_json()), result_dict)
 
+        # Special floats not supported as keys
         d_f = {float("nan"): 2, 2.3: 4, 3.4: 6, 4.5: 7}
         d_f_res = {str(k): v for k, v in d_f.items()}
         test_struct = MyStruct(i=456, d_any=d_f)
@@ -1760,6 +1762,65 @@ class MyStruct(csp.Struct):
         with self.assertRaises(ValueError):
             test_struct.to_json()
 
+        # None as key
+        d_none = {
+            None: 2,
+        }
+        d_none_res = {"null": 2}
+        test_struct = MyStruct(i=456, d_any=d_none)
+        result_dict = {"i": 456, "d_any": d_none_res}
+        self.assertEqual(json.loads(test_struct.to_json()), result_dict)
+
+        # Bool as key
+        d_bool = {True: 2, False: "abc"}
+        d_bool_res = {str(k): v for k, v in d_bool.items()}
+        test_struct = MyStruct(i=456, d_any=d_bool)
+        result_dict = {"i": 456, "d_any": d_bool_res}
+        self.assertEqual(json.loads(test_struct.to_json()), result_dict)
+
+        # Datetime as key
+        dt = datetime.now(tz=pytz.utc)
+        d_datetime = {dt: "datetime"}
+        test_struct = MyStruct(i=456, d_any=d_datetime)
+        result_dict = json.loads(test_struct.to_json())
+        self.assertEqual({datetime.fromisoformat(k): v for k, v in result_dict["d_any"].items()}, d_datetime)
+
+        dt = datetime.now(tz=pytz.utc)
+        d_datetime = {dt.date(): "date"}
+        test_struct = MyStruct(i=456, d_any=d_datetime)
+        result_dict = json.loads(test_struct.to_json())
+        self.assertEqual({date.fromisoformat(k): v for k, v in result_dict["d_any"].items()}, d_datetime)
+
+        dt = datetime.now(tz=pytz.utc)
+        d_datetime = {dt.time(): "time"}
+        test_struct = MyStruct(i=456, d_any=d_datetime)
+        result_dict = json.loads(test_struct.to_json())
+        self.assertEqual({time.fromisoformat(k): v for k, v in result_dict["d_any"].items()}, d_datetime)
+
+        # csp.Enum as key
+        class MyCspEnum(csp.Enum):
+            KEY1 = csp.Enum.auto()
+            KEY2 = csp.Enum.auto()
+            KEY3 = csp.Enum.auto()
+
+        d_csp_enum = {MyCspEnum.KEY1: "key1", MyCspEnum.KEY2: "key2", MyCspEnum.KEY3: "key3"}
+        d_csp_enum_res = {k.name: v for k, v in d_csp_enum.items()}
+        test_struct = MyStruct(i=456, d_any=d_csp_enum)
+        result_dict = {"i": 456, "d_any": d_csp_enum_res}
+        self.assertEqual(json.loads(test_struct.to_json()), result_dict)
+
+        # enum as key
+        class MyPyEnum(enum.Enum):
+            KEY1 = enum.auto()
+            KEY2 = enum.auto()
+            KEY3 = enum.auto()
+
+        d_py_enum = {MyPyEnum.KEY1: "key1", MyPyEnum.KEY2: "key2", MyPyEnum.KEY3: "key3"}
+        d_py_enum_res = {k.name: v for k, v in d_csp_enum.items()}
+        test_struct = MyStruct(i=456, d_any=d_py_enum)
+        result_dict = {"i": 456, "d_any": d_py_enum_res}
+        self.assertEqual(json.loads(test_struct.to_json()), result_dict)
+
     def test_to_json_struct(self):
         class MySubSubStruct(csp.Struct):
             b: bool = True

docs/wiki/api-references/csp.Struct-API.md

@@ -142,6 +142,7 @@ print(f"Using FastList field: value {s.a}, type {type(s.a)}, is Python list: {is
 - **`from_dict(self, dict)`**: convert a regular python dict to an instance of the struct
 - **`metadata(self)`**: returns the struct's metadata as a dictionary of key : type pairs
 - **`to_dict(self)`**: convert struct instance to a python dictionary
+- **`to_json(self, callback=lambda x: x)`**: convert struct instance to a json string, callback is invoked for any values encountered when processing the struct that are not basic Python types, datetime types, tuples, lists, dicts, csp.Structs, or csp.Enums. The callback should convert the unhandled type to a combination of the known types.
 - **`all_fields_set(self)`**: returns `True` if all the fields on the struct are set. Note that this will not recursively check sub-struct fields
 
 # Note on inheritance