python: make restore_fn optional and symmetric with save_fn

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

python/README.md

@@ -571,7 +571,7 @@ Checkpoint.restore("my-snapshot", restore_fn=lambda data: rebuild(data))
 |--------|-------------|
 | `cp.save(name, store=None)` | Persist checkpoint to disk |
 | `Checkpoint.load(name, store=None)` | Load from disk |
-| `Checkpoint.restore(name, restore_fn, store=None)` | Load and call restore_fn with app_state |
+| `Checkpoint.restore(name, restore_fn=None, store=None)` | Load; if app_state was saved (via save_fn) call restore_fn with it. restore_fn and save_fn must both be present or both absent. |
 | `Checkpoint.list(store=None)` | List saved checkpoint names |
 | `Checkpoint.delete(name, store=None)` | Delete a saved checkpoint |
 

python/src/sandlock/_sdk.py

@@ -861,18 +861,25 @@ def load(cls, name: str, *, store: Path | str | None = None) -> "Checkpoint":
     def restore(
         cls,
         name: str,
-        restore_fn: "Callable[[bytes], None]",
+        restore_fn: "Callable[[bytes], None] | None" = None,
         *,
         store: "Path | str | None" = None,
     ) -> "Checkpoint":
-        """Load a checkpoint and pass its app state to restore_fn.
+        """Load a checkpoint and, if it carries app state, pass it to restore_fn.
 
-        Convenience for ``load()`` + calling ``restore_fn(cp.app_state)``.
+        Convenience for ``load()`` plus calling ``restore_fn(cp.app_state)``.
+
+        ``restore_fn`` mirrors ``save_fn`` on ``Sandbox.checkpoint``: a checkpoint
+        taken without ``save_fn`` has no app state and is restored without
+        ``restore_fn``; a checkpoint taken with ``save_fn`` carries app state and
+        must be restored with ``restore_fn``. Supplying exactly one of the pair
+        is an error.
 
         Args:
             name: Checkpoint name.
-            restore_fn: Callback that receives the saved application-level
-                state bytes. Use this to rebuild state that ptrace can't
+            restore_fn: Optional callback receiving the saved application-level
+                state bytes. Required if (and only if) the checkpoint was taken
+                with ``save_fn``. Use it to rebuild state that ptrace cannot
                 capture (caches, session data, etc.).
             store: Storage root. Defaults to ``~/.sandlock/checkpoints/``.
 
@@ -881,16 +888,26 @@ def restore(
 
         Raises:
             FileNotFoundError: If the checkpoint does not exist.
-            ValueError: If the checkpoint has no app_state.
+            ValueError: If the checkpoint has app_state but no restore_fn was
+                given, or a restore_fn was given but the checkpoint has no
+                app_state.
         """
         cp = cls.load(name, store=store)
         state = cp.app_state
-        if state is None:
+        has_state = state is not None
+        has_fn = restore_fn is not None
+        if has_state and not has_fn:
+            raise ValueError(
+                f"Checkpoint {name!r} has app_state; pass restore_fn to "
+                "consume it (it was created with save_fn)"
+            )
+        if has_fn and not has_state:
             raise ValueError(
-                f"Checkpoint {name!r} has no app_state — "
-                "was it created with save_fn?"
+                f"Checkpoint {name!r} has no app_state; restore_fn was given "
+                "but there is nothing to pass (was it created with save_fn?)"
             )
-        restore_fn(state)
+        if has_fn:
+            restore_fn(state)
         return cp
 
     @classmethod

python/tests/test_checkpoint.py

@@ -146,12 +146,33 @@ def test_restore_returns_checkpoint(self, running_sandbox, tmp_dir):
         assert isinstance(result, Checkpoint)
         assert result.name == "ret-test"
 
-    def test_restore_no_app_state_raises(self, running_sandbox, tmp_dir):
+    def test_restore_no_app_state_no_fn_succeeds(self, running_sandbox, tmp_dir):
+        # save_fn absent at checkpoint, restore_fn absent at restore: both
+        # absent is the symmetric, valid case and must not raise.
         cp = running_sandbox.checkpoint()
         cp.save("no-app", store=tmp_dir)
 
+        result = Checkpoint.restore("no-app", store=tmp_dir)
+        assert isinstance(result, Checkpoint)
+        assert result.app_state is None
+
+    def test_restore_fn_without_app_state_raises(self, running_sandbox, tmp_dir):
+        # restore_fn given but the checkpoint has no app_state: a one-sided
+        # pairing, which is rejected.
+        cp = running_sandbox.checkpoint()
+        cp.save("no-app-fn", store=tmp_dir)
+
         with pytest.raises(ValueError, match="no app_state"):
-            Checkpoint.restore("no-app", lambda d: None, store=tmp_dir)
+            Checkpoint.restore("no-app-fn", lambda d: None, store=tmp_dir)
+
+    def test_restore_app_state_without_fn_raises(self, running_sandbox, tmp_dir):
+        # app_state present but restore_fn omitted: the other one-sided pairing,
+        # also rejected.
+        cp = running_sandbox.checkpoint(save_fn=lambda: b"state")
+        cp.save("app-no-fn", store=tmp_dir)
+
+        with pytest.raises(ValueError, match="has app_state"):
+            Checkpoint.restore("app-no-fn", store=tmp_dir)
 
     def test_restore_nonexistent_raises(self, tmp_dir):
         with pytest.raises(FileNotFoundError):