diff --git a/pyiron_workflow/storage.py b/pyiron_workflow/storage.py index 40cee293c..51b172247 100644 --- a/pyiron_workflow/storage.py +++ b/pyiron_workflow/storage.py @@ -231,14 +231,16 @@ def _save( if self._fallback(cloudpickle_fallback): attacks += [(self._CLOUDPICKLE, cloudpickle.dump)] - e = None + e: Exception | None = None for suffix, save_method in attacks: + e = None p = filename.with_suffix(suffix) try: with open(p, "wb") as filehandle: save_method(node, filehandle) return - except Exception: + except Exception as ee: + e = ee p.unlink(missing_ok=True) if e is not None: raise e diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py index 0883bd3e0..8adb54447 100644 --- a/tests/unit/test_storage.py +++ b/tests/unit/test_storage.py @@ -2,6 +2,9 @@ from pathlib import Path from tempfile import TemporaryDirectory +import cloudpickle +from pint import UnitRegistry + from pyiron_workflow.nodes.function import as_function_node from pyiron_workflow.nodes.standard import UserInput from pyiron_workflow.storage import PickleStorage, TypeNotFoundError, available_backends @@ -134,6 +137,21 @@ def Unimportable(x): finally: interface.delete(node=u, cloudpickle_fallback=True) + def test_uncloudpickleable(self): + ureg = UnitRegistry() + with self.assertRaises( + TypeError, + msg="Sanity check that this can't even be cloudpickled" + ): + cloudpickle.dumps(ureg) + + interface = PickleStorage(cloudpickle_fallback=True) + n = UserInput(ureg, label="uncloudpicklable_node") + with self.assertRaises( + TypeError, + msg="Exception should be caught and saving should fail" + ): + interface.save(n) if __name__ == "__main__": unittest.main()