From 4a1493d91dadee61179f680f24c37a03561fcfc0 Mon Sep 17 00:00:00 2001
From: finlayclark <s1629668@ed.ac.uk>
Date: Thu, 4 Jul 2024 11:07:00 +0100
Subject: [PATCH] Avoid pickling `simulation_runners` before complete
 initialisation

This should close https://github.com/michellab/a3fe/issues/1.
---
 a3fe/run/_simulation_runner.py | 6 +++++-
 a3fe/run/calculation.py        | 1 +
 a3fe/run/lambda_window.py      | 1 +
 a3fe/run/leg.py                | 1 +
 a3fe/run/simulation.py         | 1 +
 a3fe/run/stage.py              | 1 +
 docs/CHANGELOG.rst             | 4 ++++
 7 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/a3fe/run/_simulation_runner.py b/a3fe/run/_simulation_runner.py
index 0220ba7..78d5ddc 100644
--- a/a3fe/run/_simulation_runner.py
+++ b/a3fe/run/_simulation_runner.py
@@ -55,6 +55,7 @@ def __init__(
         dg_multiplier: int = 1,
         ensemble_size: int = 5,
         update_paths: bool = True,
+        dump: bool = True,
     ) -> None:
         """
         base_dir : str, Optional, default: None
@@ -80,6 +81,8 @@ def __init__(
         update_paths: bool, Optional, default: True
             If True, if the simulation runner is loaded by unpickling, then
             update_paths() is called.
+        dump: bool, Optional, default: True
+            If True, the state of the simulation runner is saved to a pickle file.
         """
         # Set up the directories (which may be overwritten if the
         # simulation runner is subsequently loaded from a pickle file)
@@ -140,7 +143,8 @@ def __init__(
             self._set_up_logging()
 
             # Save state
-            self._dump()
+            if dump:
+                self._dump()
 
     def _set_up_logging(self, null: bool = False) -> None:
         """
diff --git a/a3fe/run/calculation.py b/a3fe/run/calculation.py
index afa63bc..5755e3d 100644
--- a/a3fe/run/calculation.py
+++ b/a3fe/run/calculation.py
@@ -93,6 +93,7 @@ def __init__(
             stream_log_level=stream_log_level,
             ensemble_size=ensemble_size,
             update_paths=update_paths,
+            dump=False,
         )
 
         if not self.loaded_from_pickle:
diff --git a/a3fe/run/lambda_window.py b/a3fe/run/lambda_window.py
index 8bc5d87..73edfe1 100644
--- a/a3fe/run/lambda_window.py
+++ b/a3fe/run/lambda_window.py
@@ -122,6 +122,7 @@ def __init__(
             stream_log_level=stream_log_level,
             ensemble_size=ensemble_size,
             update_paths=update_paths,
+            dump=False,
         )
 
         if not self.loaded_from_pickle:
diff --git a/a3fe/run/leg.py b/a3fe/run/leg.py
index 2be768c..5f01a6b 100644
--- a/a3fe/run/leg.py
+++ b/a3fe/run/leg.py
@@ -119,6 +119,7 @@ def __init__(
             stream_log_level=stream_log_level,
             ensemble_size=ensemble_size,
             update_paths=update_paths,
+            dump=False,
         )
 
         if not self.loaded_from_pickle:
diff --git a/a3fe/run/simulation.py b/a3fe/run/simulation.py
index 076eacb..363ac94 100644
--- a/a3fe/run/simulation.py
+++ b/a3fe/run/simulation.py
@@ -103,6 +103,7 @@ def __init__(
             output_dir=output_dir,
             stream_log_level=stream_log_level,
             update_paths=update_paths,
+            dump=False,
         )
 
         if not self.loaded_from_pickle:
diff --git a/a3fe/run/stage.py b/a3fe/run/stage.py
index 22a1491..fa27a45 100644
--- a/a3fe/run/stage.py
+++ b/a3fe/run/stage.py
@@ -146,6 +146,7 @@ def __init__(
             stream_log_level=stream_log_level,
             ensemble_size=ensemble_size,
             update_paths=update_paths,
+            dump=False,
         )
 
         if not self.loaded_from_pickle:
diff --git a/docs/CHANGELOG.rst b/docs/CHANGELOG.rst
index e71920e..3deede0 100644
--- a/docs/CHANGELOG.rst
+++ b/docs/CHANGELOG.rst
@@ -2,6 +2,10 @@
 Change Log
 ===============
 
+v0.1.1
+
+- Ensured that ``simulation_runner`` objects only get pickled once they've been fully initialised. This avoids issues where an error occurs during initialisation and the object is pickled in an incomplete state. Addresses https://github.com/michellab/a3fe/issues/1.
+
 v0.1.0
 ====================