Merge branch 'main' into release/v1.0.0

Author: pesap

src/infrasys/arrow_storage.py

@@ -170,8 +170,38 @@ def _get_time_series_dispatch(
         msg = f"Bug: need to implement get_time_series for {type(metadata)}"
         raise NotImplementedError(msg)
 
-    @_get_time_series_dispatch.register(SingleTimeSeriesMetadata)
-    def _(
+    def remove_time_series(self, metadata: TimeSeriesMetadata, connection: Any = None) -> None:
+        fpath = self._ts_directory.joinpath(f"{metadata.time_series_uuid}{EXTENSION}")
+        if not fpath.exists():
+            msg = f"No time series with {metadata.time_series_uuid} is stored"
+            raise ISNotStored(msg)
+        fpath.unlink()
+
+    def serialize(
+        self, data: dict[str, Any], dst: Path | str, src: Path | str | None = None
+    ) -> None:
+        # From the shutil documentation: the copying operation will continue if
+        # it encounters existing directories, and files within the dst tree
+        # will be overwritten by corresponding files from the src tree.
+        if src is None:
+            src = self._ts_directory
+        src_path = Path(src)
+        dst_path = Path(dst)
+        # Note that src could be read-only. Don't copy it's permissions.
+        for path in src_path.iterdir():
+            if path.is_file():
+                if path.suffix == ".arrow":
+                    shutil.copyfile(path, dst_path / path.name)
+            else:
+                shutil.copytree(src, dst_path / path.name, dirs_exist_ok=True)
+        self.add_serialized_data(data)
+        logger.info("Copied time series data to {}", dst)
+
+    @staticmethod
+    def add_serialized_data(data: dict[str, Any]) -> None:
+        data["time_series_storage_type"] = TimeSeriesStorageType.ARROW.value
+
+    def _get_single_time_series(
         self,
         metadata: SingleTimeSeriesMetadata,
         start_time: datetime | None = None,

tests/test_serialization.py

@@ -13,7 +13,7 @@
 from pydantic import WithJsonSchema
 from typing_extensions import Annotated
 
-from infrasys import Location, NonSequentialTimeSeries, SingleTimeSeries
+from infrasys import Location, SingleTimeSeries, NonSequentialTimeSeries, System
 from infrasys.component import Component
 from infrasys.exceptions import ISInvalidParameter, ISOperationNotAllowed
 from infrasys.normalization import NormalizationMax
@@ -107,14 +107,42 @@ def test_serialize_single_time_series(tmp_path, time_series_storage_type):
 
     variable_name = "active_power"
     length = 8784
-    df = range(length)
+    data = range(length)
     start = datetime(year=2020, month=1, day=1)
     resolution = timedelta(hours=1)
-    ts = SingleTimeSeries.from_array(df, variable_name, start, resolution)
+    ts = SingleTimeSeries.from_array(data, variable_name, start, resolution)
     system.add_time_series(ts, gen1, gen2, scenario="high", model_year="2030")
     filename = tmp_path / "system.json"
     system.to_json(filename)
-    check_deserialize_with_read_write_time_series(filename)
+    system2 = check_deserialize_with_read_write_time_series(filename)
+    gen1b = system2.get_component(SimpleGenerator, gen1.name)
+    gen2b = system2.get_component(SimpleGenerator, gen2.name)
+    data2 = range(1, length + 1)
+    ts2 = SingleTimeSeries.from_array(data2, variable_name, start, resolution)
+    system2.add_time_series(ts2, gen1b, gen2b, scenario="low", model_year="2030")
+    filename2 = tmp_path / "system2.json"
+    system2.to_json(filename2)
+    system3 = SimpleSystem.from_json(filename2)
+    assert np.array_equal(
+        system3.get_time_series(
+            gen1b,
+            time_series_type=SingleTimeSeries,
+            variable_name=variable_name,
+            scenario="low",
+            model_year="2030",
+        ).data,
+        data2,
+    )
+    assert np.array_equal(
+        system3.get_time_series(
+            gen2b,
+            time_series_type=SingleTimeSeries,
+            variable_name=variable_name,
+            scenario="low",
+            model_year="2030",
+        ).data,
+        data2,
+    )
     check_deserialize_with_read_only_time_series(
         filename, gen1.name, gen2.name, variable_name, ts.data
     )
@@ -175,13 +203,14 @@ def test_serialize_nonsequential_time_series(tmp_path, time_series_storage_type)
     )
 
 
-def check_deserialize_with_read_write_time_series(filename):
+def check_deserialize_with_read_write_time_series(filename) -> System:
     system3 = SimpleSystem.from_json(filename, time_series_read_only=False)
     assert system3.get_time_series_directory() != SimpleSystem._make_time_series_directory(
         filename
     )
     system3_ts_dir = system3.get_time_series_directory()
     assert system3_ts_dir is not None
+    return system3
 
 
 @pytest.mark.parametrize(