Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 2 commits into from
Jul 17, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3332483
Add support for additional python types as dict keys in Struct.to_json
arhamchopra Jul 11, 2024
92be596
Fix memory leaks, add docs, and address comments
arhamchopra Jul 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 63 additions & 7 deletions cpp/csp/python/PyStructToJson.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" ) ) );
}
arhamchopra marked this conversation as resolved.
Show resolved Hide resolved
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 ) )
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

might want to check string first since that is by far the most common

{
Py_ssize_t len;
auto str = PyUnicode_AsUTF8AndSize( py_key, &len );
Expand All @@ -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
Expand All @@ -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 ) );
Expand Down Expand Up @@ -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() );
}
Expand Down
61 changes: 61 additions & 0 deletions csp/tests/impl/test_struct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import enum
import json
import numpy as np
arhamchopra marked this conversation as resolved.
Show resolved Hide resolved
import pickle
Expand Down Expand Up @@ -1731,6 +1732,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)
Expand All @@ -1749,6 +1751,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()}
arhamchopra marked this conversation as resolved.
Show resolved Hide resolved
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
Expand Down
1 change: 1 addition & 0 deletions docs/wiki/api-references/csp.Struct-API.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading