fix: custom date format in json (#359)

equinor · Jan 29, 2024 · d62e650 · d62e650
1 parent f586181
commit d62e650
Show file tree

Hide file tree

Showing 10 changed files with 20,380 additions and 20,349 deletions.
diff --git a/src/libecalc/infrastructure/file_utils.py b/src/libecalc/infrastructure/file_utils.py
@@ -1,8 +1,12 @@
 import enum
 import sys
-from typing import Callable, Optional, Union
+from datetime import datetime
+from typing import Any, Callable, Optional, Union
 
+import numpy as np
 import pandas as pd
+from orjson import orjson
+from pydantic import BaseModel
 
 from libecalc.common.datetime.utils import DateTimeFormats
 from libecalc.common.logger import logger
@@ -52,6 +56,20 @@ def dataframe_to_csv(
     )
 
 
+def find_float_64(data):
+    if isinstance(data, np.float64):
+        pass
+    elif isinstance(data, list):
+        for i in data:
+            find_float_64(i)
+    elif isinstance(data, dict):
+        for v in data.values():
+            find_float_64(v)
+    elif isinstance(data, BaseModel):
+        for v in data.model_dump().values():
+            find_float_64(v)
+
+
 def to_json(result: Union[ComponentResult, EcalcModelResult], simple_output: bool, date_format_option: int) -> str:
     """Dump result classes to json file
 
@@ -64,18 +82,36 @@ def to_json(result: Union[ComponentResult, EcalcModelResult], simple_output: boo
         String dump of json output
 
     """
-    DateTimeFormats.get_format(date_format_option)
-    return (
-        SimpleResultData.from_dto(result).model_dump_json(
-            indent=True,
+
+    data = (
+        SimpleResultData.from_dto(result).model_dump(
             exclude_none=True,
         )
         if simple_output
-        else result.model_dump_json(
-            indent=True,
+        else result.model_dump(
             exclude_none=True,
         )
     )
+    date_format = DateTimeFormats.get_format(date_format_option)
+
+    def default_serializer(x: Any):
+        if isinstance(x, datetime):
+            return x.strftime(date_format)
+        if isinstance(x, np.float64):
+            return float(x)
+
+        raise ValueError(f"Unable to serialize '{type(x)}'")
+
+    find_float_64(data)
+
+    # Using orjson to both allow custom date format and convert nan to null.
+    # NaN to null is not supported by json module.
+    # Custom date format is not supported by pydantic -> https://github.com/pydantic/pydantic/issues/7143
+    return orjson.dumps(
+        data,
+        default=default_serializer,
+        option=orjson.OPT_PASSTHROUGH_DATETIME | orjson.OPT_INDENT_2 | orjson.OPT_NON_STR_KEYS,
+    ).decode()
 
 
 def get_result_output(

diff --git a/src/libecalc/presentation/simple_result/simple.py b/src/libecalc/presentation/simple_result/simple.py
@@ -16,9 +16,6 @@
 class SimpleBase(EcalcResultBaseModel):
     model_config = ConfigDict(extra="ignore")
 
-    def model_dump(self, exclude_none=True, **kwargs):
-        return super().model_dump(exclude_none=exclude_none, **kwargs)
-
 
 class SimpleEmissionResult(SimpleBase):
     name: str

diff --git a/src/tests/ecalc_cli/snapshots/test_app/test_component_name_json/waterinj.json b/src/tests/ecalc_cli/snapshots/test_app/test_component_name_json/waterinj.json
@@ -19,18 +19,18 @@
             1.0
         ],
         "timesteps": [
-            "2020-01-01T00:00:00",
-            "2021-01-01T00:00:00",
-            "2022-01-01T00:00:00",
-            "2023-01-01T00:00:00",
-            "2024-01-01T00:00:00",
-            "2024-12-01T00:00:00",
-            "2026-01-01T00:00:00",
-            "2027-01-01T00:00:00",
-            "2028-01-01T00:00:00",
-            "2029-01-01T00:00:00",
-            "2030-01-01T00:00:00",
-            "2031-01-01T00:00:00"
+            "2020-01-01 00:00:00",
+            "2021-01-01 00:00:00",
+            "2022-01-01 00:00:00",
+            "2023-01-01 00:00:00",
+            "2024-01-01 00:00:00",
+            "2024-12-01 00:00:00",
+            "2026-01-01 00:00:00",
+            "2027-01-01 00:00:00",
+            "2028-01-01 00:00:00",
+            "2029-01-01 00:00:00",
+            "2030-01-01 00:00:00",
+            "2031-01-01 00:00:00"
         ],
         "unit": "MW",
         "values": [
@@ -50,18 +50,18 @@
     },
     "energy_usage_cumulative": {
         "timesteps": [
-            "2020-01-01T00:00:00",
-            "2021-01-01T00:00:00",
-            "2022-01-01T00:00:00",
-            "2023-01-01T00:00:00",
-            "2024-01-01T00:00:00",
-            "2024-12-01T00:00:00",
-            "2026-01-01T00:00:00",
-            "2027-01-01T00:00:00",
-            "2028-01-01T00:00:00",
-            "2029-01-01T00:00:00",
-            "2030-01-01T00:00:00",
-            "2031-01-01T00:00:00"
+            "2020-01-01 00:00:00",
+            "2021-01-01 00:00:00",
+            "2022-01-01 00:00:00",
+            "2023-01-01 00:00:00",
+            "2024-01-01 00:00:00",
+            "2024-12-01 00:00:00",
+            "2026-01-01 00:00:00",
+            "2027-01-01 00:00:00",
+            "2028-01-01 00:00:00",
+            "2029-01-01 00:00:00",
+            "2030-01-01 00:00:00",
+            "2031-01-01 00:00:00"
         ],
         "unit": "MWd",
         "values": [
@@ -82,18 +82,18 @@
     "id": "Sea water injection pump",
     "is_valid": {
         "timesteps": [
-            "2020-01-01T00:00:00",
-            "2021-01-01T00:00:00",
-            "2022-01-01T00:00:00",
-            "2023-01-01T00:00:00",
-            "2024-01-01T00:00:00",
-            "2024-12-01T00:00:00",
-            "2026-01-01T00:00:00",
-            "2027-01-01T00:00:00",
-            "2028-01-01T00:00:00",
-            "2029-01-01T00:00:00",
-            "2030-01-01T00:00:00",
-            "2031-01-01T00:00:00"
+            "2020-01-01 00:00:00",
+            "2021-01-01 00:00:00",
+            "2022-01-01 00:00:00",
+            "2023-01-01 00:00:00",
+            "2024-01-01 00:00:00",
+            "2024-12-01 00:00:00",
+            "2026-01-01 00:00:00",
+            "2027-01-01 00:00:00",
+            "2028-01-01 00:00:00",
+            "2029-01-01 00:00:00",
+            "2030-01-01 00:00:00",
+            "2031-01-01 00:00:00"
         ],
         "unit": "N/A",
         "values": [
@@ -130,18 +130,18 @@
             1.0
         ],
         "timesteps": [
-            "2020-01-01T00:00:00",
-            "2021-01-01T00:00:00",
-            "2022-01-01T00:00:00",
-            "2023-01-01T00:00:00",
-            "2024-01-01T00:00:00",
-            "2024-12-01T00:00:00",
-            "2026-01-01T00:00:00",
-            "2027-01-01T00:00:00",
-            "2028-01-01T00:00:00",
-            "2029-01-01T00:00:00",
-            "2030-01-01T00:00:00",
-            "2031-01-01T00:00:00"
+            "2020-01-01 00:00:00",
+            "2021-01-01 00:00:00",
+            "2022-01-01 00:00:00",
+            "2023-01-01 00:00:00",
+            "2024-01-01 00:00:00",
+            "2024-12-01 00:00:00",
+            "2026-01-01 00:00:00",
+            "2027-01-01 00:00:00",
+            "2028-01-01 00:00:00",
+            "2029-01-01 00:00:00",
+            "2030-01-01 00:00:00",
+            "2031-01-01 00:00:00"
         ],
         "unit": "MW",
         "values": [
@@ -161,18 +161,18 @@
     },
     "power_cumulative": {
         "timesteps": [
-            "2020-01-01T00:00:00",
-            "2021-01-01T00:00:00",
-            "2022-01-01T00:00:00",
-            "2023-01-01T00:00:00",
-            "2024-01-01T00:00:00",
-            "2024-12-01T00:00:00",
-            "2026-01-01T00:00:00",
-            "2027-01-01T00:00:00",
-            "2028-01-01T00:00:00",
-            "2029-01-01T00:00:00",
-            "2030-01-01T00:00:00",
-            "2031-01-01T00:00:00"
+            "2020-01-01 00:00:00",
+            "2021-01-01 00:00:00",
+            "2022-01-01 00:00:00",
+            "2023-01-01 00:00:00",
+            "2024-01-01 00:00:00",
+            "2024-12-01 00:00:00",
+            "2026-01-01 00:00:00",
+            "2027-01-01 00:00:00",
+            "2028-01-01 00:00:00",
+            "2029-01-01 00:00:00",
+            "2030-01-01 00:00:00",
+            "2031-01-01 00:00:00"
         ],
         "unit": "GWh",
         "values": [
@@ -191,17 +191,17 @@
         ]
     },
     "timesteps": [
-        "2020-01-01T00:00:00",
-        "2021-01-01T00:00:00",
-        "2022-01-01T00:00:00",
-        "2023-01-01T00:00:00",
-        "2024-01-01T00:00:00",
-        "2024-12-01T00:00:00",
-        "2026-01-01T00:00:00",
-        "2027-01-01T00:00:00",
-        "2028-01-01T00:00:00",
-        "2029-01-01T00:00:00",
-        "2030-01-01T00:00:00",
-        "2031-01-01T00:00:00"
+        "2020-01-01 00:00:00",
+        "2021-01-01 00:00:00",
+        "2022-01-01 00:00:00",
+        "2023-01-01 00:00:00",
+        "2024-01-01 00:00:00",
+        "2024-12-01 00:00:00",
+        "2026-01-01 00:00:00",
+        "2027-01-01 00:00:00",
+        "2028-01-01 00:00:00",
+        "2029-01-01 00:00:00",
+        "2030-01-01 00:00:00",
+        "2031-01-01 00:00:00"
     ]
 }