Merge pull request #303 from MC-kit/devel

Devel

Author: dvp2015

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "xpypact"
-version = "0.9.0"
+version = "0.9.1"
 description = "\"Python workflow framework for FISPACT.\""
 authors = ["dvp <[email protected]>"]
 license = "MIT"

src/xpypact/collector.py

@@ -356,6 +356,27 @@ class Result(ms.Struct):
         gbins: pl.DataFrame | None
         timestep_gamma: pl.DataFrame | None
 
+        def save_to_parquets(self, out: Path, *, override: bool = False) -> None:
+            """Save collectd data as parquet files.
+
+            Args:
+                out: directory where to save
+                override: override existing files, default - raise exception
+
+            Raises:
+                FileExistError: if destination file exists and override is not specified.
+            """
+            collected = ms.structs.asdict(self)
+            out.mkdir(parents=True, exist_ok=True)
+            for name, df in collected.items():
+                if df is None:  # pragma: no cover
+                    continue
+                dst = out / f"{name}.parquet"
+                if dst.exists() and not override:
+                    msg = f"File {dst} already exists and override is not specified."
+                    raise FileExistsError(msg)
+                df.write_parquet(dst)
+
     def get_result(self) -> FullDataCollector.Result:
         """Finish and present collected data."""
         return FullDataCollector.Result(
@@ -388,24 +409,3 @@ def get_result(self) -> FullDataCollector.Result:
             gbins=self.get_gbins(),
             timestep_gamma=self.get_timestep_gamma_as_spectrum(),
         )
-
-    def save_to_parquets(self, out: Path, *, override: bool = False) -> None:
-        """Save collectd data as parquet files.
-
-        Args:
-            out: directory where to save
-            override: override existing files, default - raise exception
-
-        Raises:
-            FileExistError: if destination file exists and override is not specified.
-        """
-        collected = ms.structs.asdict(self.get_result())
-        out.mkdir(parents=True, exist_ok=True)
-        for name, df in collected.items():
-            if df is None:  # pragma: no cover
-                continue
-            dst = out / f"{name}.parquet"
-            if dst.exists() and not override:
-                msg = f"File {dst} already exists and override is not specified."
-                raise FileExistsError(msg)
-            df.write_parquet(dst)

src/xpypact/dao/duckdb/implementation.py

@@ -139,15 +139,15 @@ def load_gamma(self, time_step_number: int | None = None) -> db.DuckDBPyRelation
 
 def save(
     cursor: db.DuckDBPyConnection,
-    collector: FullDataCollector,
+    collector_result: FullDataCollector.Result,
 ) -> None:
     """Save collected inventories to a DuckDB database.
 
     Args:
         cursor: separate multi-threaded cursor to access DuckDB, use con.cursor() in caller
-        collector: collected inventories as Polars frames
+        collector_result: collected inventories as Polars frames
     """
-    collected = ms.structs.asdict(collector.get_result())
+    collected = ms.structs.asdict(collector_result)
     for name, df in collected.items():
         if df is not None:  # pragma: no cover
             cursor.execute(f"create or replace table {name} as select * from df")  # noqa: S608

tests/test_collector.py

@@ -89,7 +89,7 @@ def test_one_cell_json(one_cell: Inventory, one_cell_time_step7_gamma_spectrum,
         )
 
     con = db.connect()
-    save(con, collector)
+    save(con, collector.get_result())
     gamma_from_db = con.sql(
         """
         select
@@ -109,10 +109,10 @@ def test_one_cell_json(one_cell: Inventory, one_cell_time_step7_gamma_spectrum,
     )
     collected = collector.get_result()
     assert collected.timestep_times.height == 7
-    collector.save_to_parquets(tmp_path)
-    collector.save_to_parquets(tmp_path, override=True)
+    collected.save_to_parquets(tmp_path)
+    collected.save_to_parquets(tmp_path, override=True)
     with pytest.raises(FileExistsError):
-        collector.save_to_parquets(tmp_path, override=False)
+        collected.save_to_parquets(tmp_path, override=False)
 
 
 if __name__ == "__main__":

tests/test_duckdb_dao.py

@@ -43,7 +43,7 @@ def test_save(inventory_with_gamma) -> None:
         dao = DataAccessObject(con)
         dc = FullDataCollector()
         dc.append(inventory_with_gamma, material_id=1, case_id=1)
-        save(con, dc)
+        save(con, dc.get_result())
         run_data = dao.load_rundata().df().loc[0]
         assert run_data["timestamp"] == pd.Timestamp("2022-02-21 01:52:45")
         assert run_data["run_name"] == "* Material Cu, fluxes 104_2_1_1"