fix: adding support for other data types in orjson_dumps

libs/libcommon/src/libcommon/utils.py

@@ -11,6 +11,7 @@
 from typing import Any, Optional, TypedDict
 
 import orjson
+import pandas as pd
 
 from libcommon.exceptions import DatasetInBlockListError
 
@@ -113,11 +114,15 @@ def orjson_default(obj: Any) -> Any:
         # the bytes are encoded with base64, and then decoded as utf-8
         # (ascii only, by the way) to get a string
         return base64.b64encode(obj).decode("utf-8")
-    raise TypeError
+    if isinstance(obj, pd.Timestamp):
+        return obj.to_pydatetime()
+    return str(obj)
 
 
 def orjson_dumps(content: Any) -> bytes:
-    return orjson.dumps(content, option=orjson.OPT_UTC_Z, default=orjson_default)
+    return orjson.dumps(
+        content, option=orjson.OPT_UTC_Z | orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS, default=orjson_default
+    )
 
 
 def get_datetime(days: Optional[float] = None) -> datetime:

libs/libcommon/tests/test_utils.py

@@ -4,10 +4,17 @@
 from contextlib import nullcontext as does_not_raise
 from typing import Any
 
+import numpy as np
+import pandas as pd
 import pytest
 
 from libcommon.exceptions import DatasetInBlockListError
-from libcommon.utils import inputs_to_string, is_image_url, raise_if_blocked
+from libcommon.utils import (
+    inputs_to_string,
+    is_image_url,
+    orjson_dumps,
+    raise_if_blocked,
+)
 
 
 @pytest.mark.parametrize(
@@ -62,3 +69,19 @@ def test_is_image_url(text: str, expected: bool) -> None:
 def test_raise_if_blocked(dataset: str, blocked: list[str], expectation: Any) -> None:
     with expectation:
         raise_if_blocked(dataset=dataset, blocked_datasets=blocked)
+
+
+def test_orjson_dumps() -> None:
+    obj = {
+        "int": 1,
+        "bool": True,
+        "str": "text",
+        "np_array": np.array([1, 2, 3, 4]),
+        "np_int64": np.int64(12),
+        "pd_timedelta": pd.Timedelta(1, "d"),
+        "object": {"a": 1, "b": 10.2},
+        "non_string_key": {1: 20},
+        "pd_timestmap": pd.Timestamp("2023-10-06"),
+    }
+    serialized_obj = orjson_dumps(obj)
+    assert serialized_obj is not None