add pandas timestamps and test

imatiach-msft · Sep 6, 2023 · e60e1f3 · e60e1f3
1 parent ff48382
commit e60e1f3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/raiutils/raiutils/data_processing/data_processing_utils.py b/raiutils/raiutils/data_processing/data_processing_utils.py
@@ -143,7 +143,7 @@ def serialize_json_safe(o: Any):
         if isinstance(o, str):
             return json.dumps(o)[1:-1]
         return o
-    elif isinstance(o, datetime.datetime):
+    elif isinstance(o, datetime.datetime) or isinstance(o, pd.Timestamp):
         return o.__str__()
     elif isinstance(o, dict):
         return {k: serialize_json_safe(v, ) for k, v in o.items()}

diff --git a/raiutils/tests/test_data_processing_utils.py b/raiutils/tests/test_data_processing_utils.py
@@ -253,6 +253,12 @@ def test_serialize_timestamp(self):
         datetime_object = datetime.datetime.strptime(datetime_str, "%Y-%m-%d")
         result = serialize_json_safe(datetime_object)
         assert datetime_str in result
+
+    def test_serialize_pandas_timestamp(self):
+        datetime_str = "2020-10-10"
+        datetime_object = pd.Timestamp(datetime_str)
+        result = serialize_json_safe(datetime_object)
+        assert datetime_str in result
 
     def test_serialize_via_json_timestamp(self):
         timestamp_obj = pd.Timestamp(2020, 1, 1)