Merge pull request #197 from pyiron/safe_save_dynamics

Fail early if h5io won't be able to re-instantiate a node

pyiron_workflow/composite.py

@@ -763,3 +763,15 @@ def _restore_signal_connections_from_strings(
             self._get_signals_input,
             self._get_signals_output,
         )
+
+    @property
+    def import_ready(self) -> bool:
+        return super().import_ready and all(node.import_ready for node in self)
+
+    def _report_import_readiness(self, tabs=0, report_so_far=""):
+        report = super()._report_import_readiness(
+            tabs=tabs, report_so_far=report_so_far
+        )
+        for node in self:
+            report = node._report_import_readiness(tabs=tabs + 1, report_so_far=report)
+        return report

pyiron_workflow/node.py

@@ -10,6 +10,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from concurrent.futures import Executor as StdLibExecutor, Future
+from importlib import import_module
 from typing import Any, Literal, Optional, TYPE_CHECKING
 
 from pyiron_workflow.channels import (
@@ -156,6 +157,13 @@ class Node(HasToDict, ABC, metaclass=AbstractHasPost):
         - On instantiation, nodes will load automatically if they find saved content.
           - Discovered content can instead be deleted with a kwarg.
           - You can't load saved content _and_ run after instantiation at once.
+        - The nodes must be somewhere importable, and the imported object must match
+            the type of the node being saved. This basically just rules out one edge
+            case where a node class is defined like
+            `SomeFunctionNode = Workflow.wrap_as.function_node()(some_function)`, since
+            then the new class gets the name `some_function`, which when imported is
+            the _function_ "some_function" and not the desired class "SomeFunctionNode".
+            This is checked for at save-time and will cause a nice early failure.
         - [ALPHA ISSUE] If the source code (cells, `.py` files...) for a saved graph is
             altered between saving and loading the graph, there are no guarantees about
             the loaded state; depending on the nature of the changes everything may
@@ -231,6 +239,8 @@ class Node(HasToDict, ABC, metaclass=AbstractHasPost):
             connected.
         future (concurrent.futures.Future | None): A futures object, if the node is
             currently running or has already run using an executor.
+        import_ready (bool): Whether importing the node's class from its class's module
+            returns the same thing as its type. (Recursive on sub-nodes for composites.)
         inputs (pyiron_workflow.io.Inputs): **Abstract.** Children must define
             a property returning an :class:`Inputs` object.
         label (str): A name for the node.
@@ -1224,3 +1234,34 @@ def tidy_working_directory(self):
             self._working_directory = None
             # Touching the working directory may have created it -- if it's there and
             # empty just clean it up
+
+    @property
+    def import_ready(self) -> bool:
+        """
+        Checks whether `importlib` can find this node's class, and if so whether the
+        imported object matches the node's type.
+
+        Returns:
+            (bool): Whether the imported module and name of this node's class match
+                its type.
+        """
+        try:
+            module = self.__class__.__module__
+            class_ = getattr(import_module(module), self.__class__.__name__)
+            if module == "__main__":
+                warnings.warn(f"{self.label} is only defined in __main__")
+            return type(self) is class_
+        except (ModuleNotFoundError, AttributeError):
+            return False
+
+    @property
+    def import_readiness_report(self):
+        print(self._report_import_readiness())
+
+    def _report_import_readiness(self, tabs=0, report_so_far=""):
+        newline = "\n" if len(report_so_far) > 0 else ""
+        tabspace = tabs * "\t"
+        return (
+            report_so_far + f"{newline}{tabspace}{self.label}: "
+            f"{'ok' if self.import_ready else 'NOT IMPORTABLE'}"
+        )

pyiron_workflow/storage.py

@@ -18,6 +18,13 @@
 ALLOWED_BACKENDS = ["h5io", "tinybase"]
 
 
+class TypeNotFoundError(ImportError):
+    """
+    Raised when you try to save a node, but importing its module and class give
+    something other than its type.
+    """
+
+
 class StorageInterface:
 
     _TINYBASE_STORAGE_FILE_NAME = "project.h5"
@@ -40,6 +47,15 @@ def save(self, backend: Literal["h5io", "tinybase"]):
             root.storage.save(backend=backend)
 
     def _save(self, backend: Literal["h5io", "tinybase"]):
+        if not self.node.import_ready:
+            raise TypeNotFoundError(
+                f"{self.node.label} cannot be saved because it (or one "
+                f"of its child nodes) has a type that cannot be imported. Did you "
+                f"dynamically define this node? Try using the node wrapper as a "
+                f"decorator instead. \n"
+                f"Import readiness report: \n"
+                f"{self.node._report_import_readiness()}"
+            )
         if backend == "h5io":
             h5io.write_hdf5(
                 fname=self._h5io_storage_file_path,

tests/static/demo_nodes.py

@@ -27,4 +27,10 @@ def AddPlusOne(obj, other):
     return obj + other + 1
 
 
-nodes = [OptionallyAdd, AddThree, AddPlusOne]
+def dynamic(x):
+    return x + 1
+
+
+Dynamic = Workflow.wrap_as.single_value_node()(dynamic)
+
+nodes = [OptionallyAdd, AddThree, AddPlusOne, Dynamic]

tests/unit/test_composite.py

@@ -652,6 +652,48 @@ def test_graph_info(self):
                     "from all depths."
             )
 
+    def test_import_ready(self):
+        self.comp.register("static.demo_nodes", "demo")
+
+        totally_findable = Composite.create.demo.OptionallyAdd()
+        self.assertTrue(
+            totally_findable.import_ready,
+            msg="The node class is well defined and in an importable module"
+        )
+        bad_class = Composite.create.demo.dynamic()
+        self.assertFalse(
+            bad_class.import_ready,
+            msg="The node is in an importable location, but the imported object is not "
+                "the node class (but rather the node function)"
+        )
+        with self.subTest(msg="Made up module"):
+            og_module = totally_findable.__class__.__module__
+            try:
+                totally_findable.__class__.__module__ = "something I totally made up"
+                self.assertFalse(
+                    totally_findable.import_ready,
+                    msg="The node class is well defined, but the module is not in the "
+                        "python path so import fails"
+                )
+            finally:
+                totally_findable.__class__.__module__ = og_module  # Fix what you broke
+
+        self.assertTrue(
+            self.comp.import_ready,
+            msg="Sanity check on initial condition -- tests are in the path, so this "
+                "is importable"
+        )
+        self.comp.totally_findable = totally_findable
+        self.assertTrue(
+            self.comp.import_ready,
+            msg="Adding importable children should leave the parent import-ready"
+        )
+        self.comp.bad_class = bad_class
+        self.assertFalse(
+            self.comp.import_ready,
+            msg="Adding un-importable children should make the parent not import ready"
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

tests/unit/test_node_package.py

@@ -36,7 +36,7 @@ def test_nodes(self):
 
     def test_length(self):
         package = NodePackage("static.demo_nodes")
-        self.assertEqual(3, len(package))
+        self.assertEqual(4, len(package))
 
 
 if __name__ == '__main__':

tests/unit/test_workflow.py

@@ -6,6 +6,7 @@
 from pyiron_workflow._tests import ensure_tests_in_python_path
 from pyiron_workflow.channels import NOT_DATA
 from pyiron_workflow.snippets.dotdict import DotDict
+from pyiron_workflow.storage import TypeNotFoundError
 from pyiron_workflow.workflow import Workflow
 
 
@@ -381,6 +382,20 @@ def test_storage_scopes(self):
                 finally:
                     wf.storage.delete()
 
+        with self.subTest("No unimportable nodes for either back-end"):
+            try:
+                wf.import_type_mismatch = wf.create.demo.dynamic()
+                for backend in ["h5io", "tinybase"]:
+                    with self.subTest(backend):
+                        with self.assertRaises(
+                            TypeNotFoundError,
+                            msg="Imported object is function but node type is node -- "
+                                "should fail early on save"
+                        ):
+                            wf.save(backend=backend)
+            finally:
+                wf.remove_node(wf.import_type_mismatch)
+
         wf.add_node(PlusOne(label="local_but_importable"))
         try:
             wf.save(backend="h5io")
@@ -415,13 +430,12 @@ def UnimportableScope(x):
             wf.unimportable_scope = UnimportableScope()
 
             try:
-                wf.save(backend="h5io")
                 with self.assertRaises(
-                    AttributeError,
+                    TypeNotFoundError,
                     msg="Nodes must live in an importable scope to save with the h5io "
                         "backend"
                 ):
-                    Workflow(wf.label, storage_backend="h5io")
+                    wf.save(backend="h5io")
             finally:
                 wf.remove_node(wf.unimportable_scope)
                 wf.storage.delete()