WIP: forming Savable protocol

unkcpz · Dec 9, 2024 · 4836edc · 4836edc
1 parent 55bc734
commit 4836edc
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 17 deletions.
diff --git a/src/plumpy/event_helper.py b/src/plumpy/event_helper.py
@@ -5,7 +5,7 @@
 from plumpy.utils import SAVED_STATE_TYPE
 
 from . import persistence
-from plumpy.persistence import Savable, LoadSaveContext, _ensure_object_loader, auto_load
+from plumpy.persistence import Savable, LoadSaveContext, _ensure_object_loader, auto_load, auto_save
 
 if TYPE_CHECKING:
     from typing import Set, Type
@@ -48,6 +48,11 @@ def recreate_from(cls, saved_state: SAVED_STATE_TYPE, load_context: Optional[Loa
         auto_load(obj, saved_state, load_context)
         return obj
 
+    def save(self, save_context: Optional[LoadSaveContext] = None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
+
     @property
     def listeners(self) -> 'Set[ProcessListener]':
         return self._listeners

diff --git a/src/plumpy/persistence.py b/src/plumpy/persistence.py
@@ -474,9 +474,7 @@ def recreate_from(cls, saved_state: SAVED_STATE_TYPE, load_context: Optional[Loa
         ...
 
     def save(self, save_context: Optional[LoadSaveContext] = None) -> SAVED_STATE_TYPE:
-        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
-
-        return out_state
+        ...
 
     def _ensure_persist_configured(self) -> None:
         if not self._persist_configured:

diff --git a/src/plumpy/process_states.py b/src/plumpy/process_states.py
@@ -1,8 +1,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-import copy
-import inspect
 import sys
 import traceback
 from enum import Enum
@@ -13,19 +11,17 @@
     Callable,
     ClassVar,
     Optional,
-    Protocol,
     Tuple,
     Type,
     Union,
     cast,
     final,
-    runtime_checkable,
+    override,
 )
 
 import yaml
 from yaml.loader import Loader
 
-from plumpy import loaders
 from plumpy.process_comms import KillMessage, MessageType
 from plumpy.persistence import _ensure_object_loader
 
@@ -40,9 +36,6 @@
 from .base import state_machine as st
 from .lang import NULL
 from .persistence import (
-    META__OBJECT_LOADER,
-    META__TYPE__METHOD,
-    META__TYPE__SAVABLE,
     LoadSaveContext,
     Savable,
     auto_load,
@@ -94,8 +87,26 @@ class PauseInterruption(Interruption):
 
 
 class Command(persistence.Savable):
-    pass
 
+    @classmethod
+    def recreate_from(cls, saved_state: SAVED_STATE_TYPE, load_context: Optional[LoadSaveContext] = None) -> 'Savable':
+        """
+        Recreate a :class:`Savable` from a saved state using an optional load context.
+
+        :param saved_state: The saved state
+        :param load_context: An optional load context
+
+        :return: The recreated instance
+
+        """
+        obj = cls.__new__(cls)
+        auto_load(obj, saved_state, load_context)
+        return obj
+
+    def save(self, save_context: Optional[LoadSaveContext] = None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
 
 @auto_persist('msg')
 class Kill(Command):
@@ -140,12 +151,14 @@ def __init__(self, continue_fn: Callable[..., Any], *args: Any, **kwargs: Any):
         self.args = args
         self.kwargs = kwargs
 
+    @override
     def save(self, save_context: Optional[persistence.LoadSaveContext] = None) -> SAVED_STATE_TYPE:
         out_state: SAVED_STATE_TYPE = persistence.auto_save(self, save_context)
         out_state[self.CONTINUE_FN] = self.continue_fn.__name__
 
         return out_state
 
+    @override
     @classmethod
     def recreate_from(cls, saved_state: SAVED_STATE_TYPE, load_context: Optional[LoadSaveContext] = None) -> 'Savable':
         """
@@ -605,6 +618,11 @@ def recreate_from(cls, saved_state: SAVED_STATE_TYPE, load_context: Optional[Loa
         auto_load(obj, saved_state, load_context)
         return obj
 
+    def save(self, save_context: Optional[LoadSaveContext] = None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
+
     def enter(self) -> None: ...
 
     def exit(self) -> None: ...
@@ -649,6 +667,11 @@ def recreate_from(cls, saved_state: SAVED_STATE_TYPE, load_context: Optional[Loa
         auto_load(obj, saved_state, load_context)
         return obj
 
+    def save(self, save_context: Optional[LoadSaveContext] = None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
+
     def enter(self) -> None: ...
 
     def exit(self) -> None: ...

diff --git a/tests/test_persistence.py b/tests/test_persistence.py
@@ -5,7 +5,8 @@
 import yaml
 
 import plumpy
-from plumpy.persistence import auto_load
+from plumpy.persistence import auto_load, auto_save
+from plumpy.utils import SAVED_STATE_TYPE
 
 from . import utils
 
@@ -14,7 +15,7 @@ class SaveEmpty(plumpy.Savable):
     pass
 
     @classmethod
-    def recreate_from(cls, saved_state, load_context= None):
+    def recreate_from(cls, saved_state, load_context=None):
         """
         Recreate a :class:`Savable` from a saved state using an optional load context.
 
@@ -28,6 +29,11 @@ def recreate_from(cls, saved_state, load_context= None):
         auto_load(obj, saved_state, load_context)
         return obj
 
+    def save(self, save_context=None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
+
 
 @plumpy.auto_persist('test', 'test_method')
 class Save1(plumpy.Savable):
@@ -39,7 +45,7 @@ def m():
         pass
 
     @classmethod
-    def recreate_from(cls, saved_state, load_context= None):
+    def recreate_from(cls, saved_state, load_context=None):
         """
         Recreate a :class:`Savable` from a saved state using an optional load context.
 
@@ -53,14 +59,19 @@ def recreate_from(cls, saved_state, load_context= None):
         auto_load(obj, saved_state, load_context)
         return obj
 
+    def save(self, save_context=None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
+
 
 @plumpy.auto_persist('test')
 class Save(plumpy.Savable):
     def __init__(self):
         self.test = Save1()
 
     @classmethod
-    def recreate_from(cls, saved_state, load_context= None):
+    def recreate_from(cls, saved_state, load_context=None):
         """
         Recreate a :class:`Savable` from a saved state using an optional load context.
 
@@ -74,6 +85,11 @@ def recreate_from(cls, saved_state, load_context= None):
         auto_load(obj, saved_state, load_context)
         return obj
 
+    def save(self, save_context=None) -> SAVED_STATE_TYPE:
+        out_state: SAVED_STATE_TYPE = auto_save(self, save_context)
+
+        return out_state
+
 
 class TestSavable(unittest.TestCase):
     def test_empty_savable(self):