base_model, examples: Add write/load methods, add sample JSON

examples/development/entrypoint_plotting.py

@@ -17,12 +17,15 @@ def main():
     results = results_raw.get_run_results()
 
     # Roundtrip save/load to/from JSON, just for demonstration
-    with open(results_path / "results.json", "w", encoding="utf-8") as file:
-        print("Dumping results to file")
-        file.write(results.model_dump_json())
-    with open(results_path / "results.json", "r", encoding="utf-8") as file:
-        print("Loading results from file")
-        results = mesmo.data_models.RunResults.model_validate_json(file.read())
+    print("Dumping results to file")
+    results.to_json_file(results_path / "results.json")
+    print("Loading results from file")
+    results = mesmo.data_models.RunResults.from_json_file(results_path / "results.json")
+
+    # Load sample results
+    results = mesmo.data_models.RunResults.from_json_file(
+        mesmo.config.base_path / "examples" / "development" / "results.json.bz2", decompress=True
+    )
 
     # Sample plotting to file, just for demonstration
     plots.plot_to_file(plots.der_active_power_time_series, results=results, results_path=results_path)

mesmo/data_models/base_model.py

@@ -1,13 +1,17 @@
 """Custom pydantic base model."""
 
-from typing import Annotated, Optional
+import bz2
+import pathlib
+from typing import Annotated, Optional, TypeVar
 
 import numpy as np
 import pandas as pd
 import pandera as pa
 import pydantic as pyd
 from pydantic.json import timedelta_isoformat
 
+Model = TypeVar("Model", bound="BaseModel")
+
 
 class BaseModel(pyd.BaseModel):
     """Custom base model which overrides pydantic default definitions."""
@@ -23,6 +27,36 @@ class BaseModel(pyd.BaseModel):
         },
     )
 
+    def to_json_file(self, file_path: pathlib.Path, compress=False):
+        """Write data model to given file path as JSON file.
+
+        Args:
+            file_path (pathlib.Path): File output path
+        """
+        if compress:
+            with bz2.open(file_path, "wt", encoding="utf-8") as file:
+                file.write(self.model_dump_json())
+        else:
+            with open(file_path, "w", encoding="utf-8") as file:
+                file.write(self.model_dump_json())
+
+    @classmethod
+    def from_json_file(cls: type[Model], file_path: pathlib.Path, decompress=False) -> Model:
+        """Load data model from given JSON file path.
+
+        Args:
+            file_path (pathlib.Path): File input path
+
+        Returns:
+            Instantiated data model based on given JSON file
+        """
+        if decompress:
+            with bz2.open(file_path, "rt", encoding="utf-8") as file:
+                return cls.model_validate_json(file.read())
+        else:
+            with open(file_path, "r", encoding="utf-8") as file:
+                return cls.model_validate_json(file.read())
+
 
 def get_index_annotation(dtype: type[pd.Timestamp | int | str], optional: bool = False) -> type[pd.Index]:
     """Get pydantic-compatible type annotation for pandas.Index."""