Merge pull request #759 from pyiron/loading

Better loading

Author: liamhuber

pyiron_workflow/node.py

@@ -18,6 +18,7 @@
 from pyiron_snippets.colors import SeabornColors
 from pyiron_snippets.dotdict import DotDict
 
+from pyiron_workflow import overloading
 from pyiron_workflow.channels import (
     AccumulatingInputSignal,
     Channel,
@@ -978,15 +979,56 @@ def save_checkpoint(self, backend: BackendIdentifier | StorageInterface = "pickl
         """
         self.graph_root.save(backend=backend)
 
-    def load(
-        self,
+    @classmethod
+    def _new_instance_from_storage(
+        cls,
         backend: BackendIdentifier | StorageInterface = "pickle",
         only_requested=False,
         filename: str | Path | None = None,
+        _node: Node | None = None,
         **kwargs,
     ):
         """
+        Loads a node from file returns its instance.
 
+        Args:
+            backend (str | StorageInterface): The interface to use for serializing the
+                node. (Default is "pickle", which loads the standard pickling back end.)
+            only_requested (bool): Whether to _only_ try loading from the specified
+                backend, or to loop through all available backends. (Default is False,
+                try to load whatever you can find.)
+            filename (str | Path | None): The name of the file (without extensions)
+                from which to load the node. (Default is None, which uses the node's
+                lexical path.)
+            **kwargs: back end-specific arguments (only likely to work in combination
+                with :param:`only_requested`, otherwise there's nothing to be specific
+                _to_.)
+
+        Raises:
+            FileNotFoundError: when nothing got loaded.
+        """
+        inst = None
+        for selected_backend in available_backends(
+            backend=backend, only_requested=only_requested
+        ):
+            inst = selected_backend.load(node=_node, filename=filename, **kwargs)
+            if inst is not None:
+                break
+        if inst is None:
+            raise FileNotFoundError(
+                f"Could not find saved content at {filename} using backend={backend} "
+                f"using only_request={only_requested}."
+            )
+        return inst
+
+    def _update_instance_from_storage(
+        self,
+        backend: BackendIdentifier | StorageInterface = "pickle",
+        only_requested=False,
+        filename: str | Path | None = None,
+        **kwargs,
+    ):
+        """
         Loads the node file and set the loaded state as the node's own.
 
         Args:
@@ -995,8 +1037,8 @@ def load(
             only_requested (bool): Whether to _only_ try loading from the specified
                 backend, or to loop through all available backends. (Default is False,
                 try to load whatever you can find.)
-            filename (str | Path | None): The name of the file (without extensions) at
-                which to save the node. (Default is None, which uses the node's
+            filename (str | Path | None): The name of the file (without extensions)
+                from which to load the node. (Default is None, which uses the node's
                 lexical path.)
             **kwargs: back end-specific arguments (only likely to work in combination
                 with :param:`only_requested`, otherwise there's nothing to be specific
@@ -1012,16 +1054,13 @@ def load(
                 "is the correct thing to do, you can set `self.running=True` where "
                 "`self` is this node object."
             )
-        for selected_backend in available_backends(
-            backend=backend, only_requested=only_requested
-        ):
-            inst = selected_backend.load(
-                node=self if filename is None else None, filename=filename, **kwargs
-            )
-            if inst is not None:
-                break
-        if inst is None:
-            raise FileNotFoundError(f"{self.label} could not find saved content.")
+        inst = self.__class__._new_instance_from_storage(
+            backend=backend,
+            only_requested=only_requested,
+            filename=filename,
+            _node=self if filename is None else None,
+            **kwargs,
+        )
 
         if inst.__class__ != self.__class__:
             raise TypeError(
@@ -1031,6 +1070,44 @@ def load(
             )
         self.__setstate__(inst.__getstate__())
 
+    @overloading.overloaded_classmethod(class_method=_new_instance_from_storage)
+    def load(
+        self,
+        backend: BackendIdentifier | StorageInterface = "pickle",
+        only_requested=False,
+        filename: str | Path | None = None,
+        **kwargs,
+    ):
+        """
+        Load a node from storage, either as a new instance (when used as a class
+        method) or by updating the current instance (when called as a regular instance
+        method).
+
+        Args:
+            backend (str | StorageInterface): The interface to use for serializing the
+                node. (Default is "pickle", which loads the standard pickling back end.)
+            only_requested (bool): Whether to _only_ try loading from the specified
+                backend, or to loop through all available backends. (Default is False,
+                try to load whatever you can find.)
+            filename (str | Path | None): The name of the file (without extensions)
+                from which to load the node. (Default is None, which uses the node's
+                lexical path.)
+            **kwargs: back end-specific arguments (only likely to work in combination
+                with :param:`only_requested`, otherwise there's nothing to be specific
+                _to_.)
+
+        Raises:
+            FileNotFoundError: when nothing got loaded.
+            TypeError: when loading into an exisiting instance and the saved node has a
+                different class name.
+        """
+        return self._update_instance_from_storage(
+            backend=backend,
+            only_requested=only_requested,
+            filename=filename,
+            **kwargs,
+        )
+
     load.__doc__ = cast(str, load.__doc__) + _save_load_warnings
 
     def delete_storage(
@@ -1048,12 +1125,11 @@ def delete_storage(
         Args:
             backend (str | StorageInterface): The interface to use for serializing the
                 node. (Default is "pickle", which loads the standard pickling back end.)
-            only_requested (bool): Whether to _only_ try loading from the specified
-                backend, or to loop through all available backends. (Default is False,
-                try to load whatever you can find.)
-            filename (str | Path | None): The name of the file (without extensions) at
-                which to save the node. (Default is None, which uses the node's
-                lexical path.)
+            only_requested (bool): Whether to _only_ search for files using the
+                specifiedmbackend, or to loop through all available backends. (Default
+                is False, try to remove whatever you can find.)
+            filename (str | Path | None): The name of the file (without extensions) to
+                remove. (Default is None, which uses the node's lexical path.)
             delete_even_if_not_empty (bool): Whether to delete the file even if it is
                 not empty. (Default is False, which will only delete the file if it is
                 empty, i.e. has no content in it.)
@@ -1084,12 +1160,11 @@ def has_saved_content(
         Args:
             backend (str | StorageInterface): The interface to use for serializing the
                 node. (Default is "pickle", which loads the standard pickling back end.)
-            only_requested (bool): Whether to _only_ try loading from the specified
-                backend, or to loop through all available backends. (Default is False,
-                try to load whatever you can find.)
-            filename (str | Path | None): The name of the file (without extensions) at
-                which to save the node. (Default is None, which uses the node's
-                lexical path.)
+            only_requested (bool): Whether to _only_ search for files using the
+                specified backend, or to loop through all available backends. (Default
+                is False, try to finding whatever you can find.)
+            filename (str | Path | None): The name of the file (without extensions) to
+                look for. (Default is None, which uses the node's lexical path.)
             **kwargs: back end-specific arguments (only likely to work in combination
                 with :param:`only_requested`, otherwise there's nothing to be specific
                 _to_.)

pyiron_workflow/overloading.py

@@ -0,0 +1,73 @@
+import functools
+
+
+def overloaded_classmethod(class_method):
+    """
+    Decorator to define a method that behaves like both a classmethod and an
+    instancemethod under the same name.
+
+    Args:
+        instance_method: A method defined on the same object as the decorated method,
+            to be used when an instance of the object calls the decorated method (
+            instead of a class call)
+
+    Returns
+    -------
+    descriptor
+        A descriptor that dispatches to the classmethod when accessed
+        via the class, and to the given instance method when accessed
+        via an instance.
+
+    Examples:
+        >>> class Foo:
+        ...     def __init__(self, y):
+        ...         self.y = y
+        ...
+        ...     @classmethod
+        ...     def _doit_classmethod(cls, x):
+        ...         return f"Class {cls.__name__} doing {x}"
+        ...
+        ...     @overloaded_classmethod(class_method=_doit_classmethod)
+        ...     def doit(self, x):
+        ...         return f"Instance of type {type(self).__name__} doing {x} + {self.y}"
+        ...
+        >>> Foo.doit(10)
+        'Class Foo doing 10'
+        >>> Foo(5).doit(20)
+        'Instance of type Foo doing 20 + 5'
+    """
+
+    class Overloaded:
+        def __init__(self, f_instance, f_class):
+            self.f_instance = f_instance
+            self.f_class = f_class
+            functools.update_wrapper(self, f_instance)
+
+        def __get__(self, obj, cls):
+            if obj is None:
+                f_class = (
+                    cls.__dict__[self.f_class]
+                    if isinstance(self.f_class, str)
+                    else self.f_class
+                )
+
+                if isinstance(f_class, classmethod):
+                    f_class = f_class.__func__
+
+                @functools.wraps(self.f_class)
+                def bound(*args, **kwargs):
+                    return f_class(cls, *args, **kwargs)
+
+                return bound
+            else:
+
+                @functools.wraps(self.f_instance)
+                def bound(*args, **kwargs):
+                    return self.f_instance(obj, *args, **kwargs)
+
+                return bound
+
+    def wrapper(f_instance):
+        return Overloaded(f_instance, class_method)
+
+    return wrapper

tests/unit/test_overloading.py

@@ -0,0 +1,55 @@
+import unittest
+
+from pyiron_workflow.overloading import (
+    overloaded_classmethod,
+)  # replace with actual module path
+
+
+def class_string(obj, x):
+    return f"Class {obj.__name__} doing {x}"
+
+
+def instance_string(obj, x):
+    return f"Instance of type {type(obj).__name__} doing {x}"
+
+
+class Foo:
+    @overloaded_classmethod(class_method="_pseudo_classmethod")
+    def undecorated_string(self, x):
+        return instance_string(self, x)
+
+    @overloaded_classmethod(class_method="_classmethod")
+    def decorated_string(self, x):
+        return instance_string(self, x)
+
+    def _pseudo_classmethod(cls, x):
+        return class_string(cls, x)
+
+    @classmethod
+    def _classmethod(cls, y):
+        return class_string(cls, y)
+
+    @overloaded_classmethod(class_method=_pseudo_classmethod)
+    def undecorated_direct(self, x):
+        return instance_string(self, x)
+
+    @overloaded_classmethod(class_method=_classmethod)
+    def decorated_direct(self, x):
+        return instance_string(self, x)
+
+
+class TestOverloadedClassMethod(unittest.TestCase):
+    def test_instance_and_class_calls(self):
+        self.assertEqual(Foo.undecorated_string(1), class_string(Foo, 1))
+        self.assertEqual(Foo.decorated_string(2), class_string(Foo, 2))
+        self.assertEqual(Foo.undecorated_direct(3), class_string(Foo, 3))
+        self.assertEqual(Foo.decorated_direct(4), class_string(Foo, 4))
+
+        self.assertEqual(Foo().undecorated_string(1), instance_string(Foo(), 1))
+        self.assertEqual(Foo().decorated_string(2), instance_string(Foo(), 2))
+        self.assertEqual(Foo().undecorated_direct(3), instance_string(Foo(), 3))
+        self.assertEqual(Foo().decorated_direct(4), instance_string(Foo(), 4))
+
+
+if __name__ == "__main__":
+    unittest.main()