safe-state 0.1.0__py3-none-any.whl

safe_state/__init__.py

@@ -0,0 +1,64 @@
+"""
+safe-state: resumable execution for Python.
+
+Drop ``@safe_state`` on top of any function that iterates through work. If the
+function crashes — network blip, rate limit, power loss — the checkpoint is on
+disk. Run the script again and it picks up where it left off. Live objects
+(database connections, sockets, file handles) are restored via a reconnect
+registry instead of being naively pickled.
+
+Quick start:
+
+    from safe_state import safe_state
+
+    @safe_state
+    def remediate(clients, db):
+        for client in clients:
+            secure_endpoint(client, db)
+
+    remediate(load_clients(), open_db())  # crashes at client 15? just rerun.
+"""
+
+from .core import safe_state, checkpoint, SafeIterator
+from .checkpoint import Checkpoint, CheckpointStore
+from .reconnect import (
+    ReconnectRegistry,
+    reconnect_handler,
+    register_reconnector,
+    get_default_registry,
+)
+from .serialization import freeze_state, thaw_state, can_freeze
+from .exceptions import (
+    SafeStateError,
+    StateNotFoundError,
+    StateCorruptedError,
+    ReconnectError,
+    NoReconnectorError,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # decorator + helpers
+    "safe_state",
+    "checkpoint",
+    "SafeIterator",
+    # checkpoint store
+    "Checkpoint",
+    "CheckpointStore",
+    # reconnect machinery
+    "ReconnectRegistry",
+    "reconnect_handler",
+    "register_reconnector",
+    "get_default_registry",
+    # serialization
+    "freeze_state",
+    "thaw_state",
+    "can_freeze",
+    # exceptions
+    "SafeStateError",
+    "StateNotFoundError",
+    "StateCorruptedError",
+    "ReconnectError",
+    "NoReconnectorError",
+]

safe_state/checkpoint.py

@@ -0,0 +1,168 @@
+"""
+Checkpoint store: on-disk persistence of execution progress.
+
+A Checkpoint records:
+    - which items in an iteration have completed (by index)
+    - results returned for each completed item (optional)
+    - the last failure (exception type, message, traceback) if any
+    - frozen local state at the point of failure
+    - a timestamp and run identifier
+
+Checkpoints are stored as a single dill blob per job name, atomically replaced
+via write-temp-then-rename to avoid corruption on crash.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+import time
+import traceback
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set
+
+from .exceptions import StateCorruptedError, StateNotFoundError
+from .reconnect import ReconnectRegistry, get_default_registry
+from .serialization import freeze_state, thaw_state
+
+
+@dataclass
+class Checkpoint:
+    """A snapshot of progress through a safe-state-decorated function."""
+
+    job_name: str
+    completed_indices: Set[int] = field(default_factory=set)
+    results: Dict[int, Any] = field(default_factory=dict)
+    total_items: Optional[int] = None
+    last_failure: Optional[Dict[str, Any]] = None
+    frozen_state: Optional[bytes] = None  # raw bytes from freeze_state()
+    started_at: float = field(default_factory=time.time)
+    updated_at: float = field(default_factory=time.time)
+    run_count: int = 1
+
+    def mark_complete(self, index: int, result: Any = None, store_result: bool = False) -> None:
+        self.completed_indices.add(index)
+        if store_result:
+            self.results[index] = result
+        self.updated_at = time.time()
+
+    def record_failure(self, index: int, exc: BaseException, locals_snapshot: Dict[str, Any]) -> None:
+        self.last_failure = {
+            "index": index,
+            "exception_type": type(exc).__name__,
+            "exception_module": type(exc).__module__,
+            "message": str(exc),
+            "traceback": "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
+            "at": time.time(),
+        }
+        # Try to freeze the whole snapshot first; if that fails, freeze per-key,
+        # dropping values that can't be serialized.
+        try:
+            self.frozen_state = freeze_state(locals_snapshot)
+        except Exception:
+            partial: Dict[str, Any] = {}
+            dropped: List[str] = []
+            for k, v in locals_snapshot.items():
+                try:
+                    freeze_state({k: v})
+                    partial[k] = v
+                except Exception:
+                    dropped.append(k)
+            try:
+                self.frozen_state = freeze_state(partial)
+                if dropped:
+                    self.last_failure["dropped_locals"] = dropped
+            except Exception as freeze_err:
+                self.frozen_state = None
+                self.last_failure["freeze_error"] = (
+                    f"{type(freeze_err).__name__}: {freeze_err}"
+                )
+        self.updated_at = time.time()
+
+    def progress(self) -> Dict[str, Any]:
+        """Human-readable progress summary."""
+        done = len(self.completed_indices)
+        total = self.total_items if self.total_items is not None else "?"
+        return {
+            "job": self.job_name,
+            "completed": done,
+            "total": total,
+            "remaining": (self.total_items - done) if self.total_items is not None else "?",
+            "has_failure": self.last_failure is not None,
+            "run_count": self.run_count,
+        }
+
+
+class CheckpointStore:
+    """File-backed persistence for a single Checkpoint."""
+
+    def __init__(self, path: str | os.PathLike, registry: Optional[ReconnectRegistry] = None):
+        self.path = Path(path)
+        self.registry = registry or get_default_registry()
+
+    def exists(self) -> bool:
+        return self.path.exists()
+
+    def load(self) -> Checkpoint:
+        if not self.path.exists():
+            raise StateNotFoundError(f"No checkpoint at {self.path}")
+        try:
+            with open(self.path, "rb") as f:
+                blob = f.read()
+            payload = thaw_state(blob, registry=self.registry)
+        except Exception as e:
+            raise StateCorruptedError(f"Could not load checkpoint {self.path}: {e}") from e
+
+        cp = Checkpoint(
+            job_name=payload["job_name"],
+            completed_indices=set(payload.get("completed_indices", [])),
+            results=payload.get("results", {}),
+            total_items=payload.get("total_items"),
+            last_failure=payload.get("last_failure"),
+            frozen_state=payload.get("frozen_state"),
+            started_at=payload.get("started_at", time.time()),
+            updated_at=payload.get("updated_at", time.time()),
+            run_count=payload.get("run_count", 1),
+        )
+        return cp
+
+    def save(self, cp: Checkpoint) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        payload = {
+            "job_name": cp.job_name,
+            "completed_indices": list(cp.completed_indices),
+            "results": cp.results,
+            "total_items": cp.total_items,
+            "last_failure": cp.last_failure,
+            "frozen_state": cp.frozen_state,  # already bytes
+            "started_at": cp.started_at,
+            "updated_at": cp.updated_at,
+            "run_count": cp.run_count,
+        }
+        blob = freeze_state(payload, registry=self.registry)
+
+        # Atomic write: temp file in same directory, then rename.
+        tmp_fd, tmp_path = tempfile.mkstemp(
+            prefix=f".{self.path.name}.", suffix=".tmp", dir=str(self.path.parent)
+        )
+        try:
+            with os.fdopen(tmp_fd, "wb") as f:
+                f.write(blob)
+            os.replace(tmp_path, self.path)
+        except Exception:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+            raise
+
+    def clear(self) -> None:
+        if self.path.exists():
+            self.path.unlink()
+
+    def thaw_locals(self, cp: Checkpoint) -> Optional[Dict[str, Any]]:
+        """Rebuild the frozen locals snapshot, reconnecting any live objects."""
+        if cp.frozen_state is None:
+            return None
+        return thaw_state(cp.frozen_state, registry=self.registry)

safe_state/core.py

@@ -0,0 +1,358 @@
+"""
+The `safe_state` decorator and its supporting iterator.
+
+The decorator wraps a function that processes an iterable. The first positional
+argument (by default) is intercepted and replaced with a SafeIterator that:
+
+    1. Tracks progress through the iteration to disk after every item.
+    2. On exception, saves a checkpoint of completed indices + frozen locals.
+    3. On the next invocation with the same job name, loads the checkpoint
+       and skips items that were already completed.
+
+A `safe_state.checkpoint()` helper is also provided for use inside the function
+body (manual checkpoint mode).
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import os
+import sys
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
+
+from .checkpoint import Checkpoint, CheckpointStore
+from .exceptions import SafeStateError, StateNotFoundError
+from .reconnect import ReconnectRegistry, get_default_registry
+
+
+# Default directory for checkpoints. Override via env var or decorator arg.
+DEFAULT_STATE_DIR = Path(os.environ.get("SAFE_STATE_DIR", ".safe_state"))
+
+
+# Thread-local-ish current context, exposed so user code can call manual checkpoint().
+# We use a module-level dict keyed by function id — simple and avoids threading concerns
+# for the common single-threaded use case. For multithreaded use, users should pass
+# explicit CheckpointStore objects.
+_active_contexts: Dict[int, "_RunContext"] = {}
+
+
+class _RunContext:
+    """Per-call context used by the SafeIterator and manual checkpoint()."""
+
+    def __init__(
+        self,
+        store: CheckpointStore,
+        checkpoint: Checkpoint,
+        verbose: bool,
+    ):
+        self.store = store
+        self.checkpoint = checkpoint
+        self.verbose = verbose
+        self.current_index: Optional[int] = None
+        self.current_locals: Dict[str, Any] = {}
+
+
+class SafeIterator:
+    """
+    Wraps an iterable so iteration progress is checkpointed automatically.
+
+    Yields items whose index is not already in checkpoint.completed_indices.
+    After each yielded item is consumed (i.e., the next iteration starts),
+    the previous item's index is marked complete and the checkpoint is saved.
+    """
+
+    def __init__(
+        self,
+        iterable: Iterable[Any],
+        context: _RunContext,
+        save_every: int = 1,
+        store_results: bool = False,
+    ):
+        self._items = list(iterable)  # materialize to get a stable length and indexing
+        self._context = context
+        self._save_every = max(1, int(save_every))
+        self._store_results = store_results
+        self._since_save = 0
+
+        # Record total on first run.
+        if self._context.checkpoint.total_items is None:
+            self._context.checkpoint.total_items = len(self._items)
+
+    def __iter__(self) -> Iterator[Tuple[int, Any]]:
+        cp = self._context.checkpoint
+        for idx, item in enumerate(self._items):
+            if idx in cp.completed_indices:
+                if self._context.verbose:
+                    print(f"[safe_state] skip {idx} (already done)", file=sys.stderr)
+                continue
+            self._context.current_index = idx
+            yield item
+            # Mark the just-yielded item complete.
+            cp.mark_complete(idx, store_result=self._store_results)
+            self._since_save += 1
+            if self._since_save >= self._save_every:
+                self._context.store.save(cp)
+                self._since_save = 0
+        # Final flush.
+        if self._since_save > 0:
+            self._context.store.save(cp)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+
+def checkpoint(**locals_to_save: Any) -> None:
+    """
+    Manually trigger a checkpoint save from inside a decorated function.
+
+    Pass any local variables you want frozen into the checkpoint:
+
+        @safe_state
+        def long_job(tasks, db):
+            results = []
+            for t in tasks:
+                results.append(do(t))
+                checkpoint(results=results)  # snapshot progress
+
+    Has no effect if called outside a safe-state-decorated function.
+    """
+    # Walk up the call stack to find a frame whose function is decorated.
+    frame = inspect.currentframe()
+    try:
+        while frame is not None:
+            ctx = _active_contexts.get(id(frame.f_code))
+            if ctx is not None:
+                ctx.current_locals = dict(locals_to_save)
+                # Freeze locals into the checkpoint snapshot.
+                cp = ctx.checkpoint
+                from .serialization import freeze_state
+                try:
+                    cp.frozen_state = freeze_state(locals_to_save)
+                except Exception:
+                    cp.frozen_state = None
+                ctx.store.save(cp)
+                return
+            frame = frame.f_back
+    finally:
+        del frame
+
+
+def safe_state(
+    _fn: Optional[Callable] = None,
+    *,
+    name: Optional[str] = None,
+    state_dir: Union[str, os.PathLike, None] = None,
+    iterable_arg: Union[str, int] = 0,
+    save_every: int = 1,
+    store_results: bool = False,
+    keep_on_success: bool = False,
+    verbose: bool = False,
+    registry: Optional[ReconnectRegistry] = None,
+    auto_iterate: bool = True,
+) -> Callable:
+    """
+    Decorator that makes a function resumable.
+
+    Args:
+        name: Job name (used as the checkpoint filename). Defaults to the function's
+            qualified name.
+        state_dir: Directory for checkpoint files. Defaults to ``.safe_state/`` or
+            ``$SAFE_STATE_DIR``.
+        iterable_arg: Position (int) or name (str) of the iterable argument that
+            should be intercepted and checkpointed. Default: first positional arg.
+        save_every: Persist progress every N completed items. Default: 1 (every item).
+        store_results: If True, store the iterable's items in the checkpoint so
+            results from already-completed items survive across restarts. Note: items
+            must be serializable.
+        keep_on_success: If True, leave the checkpoint file in place after the
+            function completes successfully. Default: delete it.
+        verbose: Print progress messages to stderr.
+        registry: Optional custom ReconnectRegistry for live-object handling.
+        auto_iterate: If True (default), the decorated function receives a
+            SafeIterator in place of its iterable argument. If False, the function
+            is left untouched — use this with manual checkpoint() calls.
+
+    Usage:
+
+        @safe_state
+        def remediate(clients, db):
+            for client in clients:
+                secure_endpoint(client, db)
+
+        # Or with config:
+        @safe_state(name="newsletter-blast", state_dir="state/", save_every=5)
+        def remediate(clients, db):
+            ...
+
+    On the first run, the checkpoint file is created. On any subsequent run,
+    if the previous run crashed, the function resumes from where it stopped.
+    """
+    state_dir_path = Path(state_dir) if state_dir is not None else DEFAULT_STATE_DIR
+    reg = registry or get_default_registry()
+
+    def decorate(fn: Callable) -> Callable:
+        job_name = name or f"{fn.__module__}.{fn.__qualname__}"
+        # Sanitize for filename use.
+        safe_name = job_name.replace("/", "_").replace(os.sep, "_")
+        path = state_dir_path / f"{safe_name}.safestate"
+
+        sig = inspect.signature(fn)
+
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            store = CheckpointStore(path, registry=reg)
+
+            # Load or create the checkpoint.
+            if store.exists():
+                try:
+                    cp = store.load()
+                    cp.run_count += 1
+                    if verbose:
+                        prog = cp.progress()
+                        print(
+                            f"[safe_state] resuming {job_name!r}: "
+                            f"{prog['completed']}/{prog['total']} done "
+                            f"(run #{cp.run_count})",
+                            file=sys.stderr,
+                        )
+                except SafeStateError as e:
+                    if verbose:
+                        print(
+                            f"[safe_state] could not load checkpoint ({e}); starting fresh",
+                            file=sys.stderr,
+                        )
+                    cp = Checkpoint(job_name=job_name)
+            else:
+                cp = Checkpoint(job_name=job_name)
+                if verbose:
+                    print(f"[safe_state] starting fresh job {job_name!r}", file=sys.stderr)
+
+            ctx = _RunContext(store=store, checkpoint=cp, verbose=verbose)
+            _active_contexts[id(fn.__code__)] = ctx
+
+            # Bind args and intercept the iterable if auto_iterate.
+            bound = sig.bind_partial(*args, **kwargs)
+            bound.apply_defaults()
+
+            if auto_iterate:
+                arg_name: Optional[str]
+                if isinstance(iterable_arg, int):
+                    params = list(sig.parameters)
+                    if iterable_arg >= len(params):
+                        raise SafeStateError(
+                            f"iterable_arg={iterable_arg} but {fn.__name__} only "
+                            f"has {len(params)} parameters"
+                        )
+                    arg_name = params[iterable_arg]
+                else:
+                    arg_name = iterable_arg
+                if arg_name not in bound.arguments:
+                    raise SafeStateError(
+                        f"iterable arg {arg_name!r} not provided to {fn.__name__}"
+                    )
+                original_iter = bound.arguments[arg_name]
+                bound.arguments[arg_name] = (
+                    item for _, item in _enumerate_safe(
+                        original_iter, ctx, save_every, store_results
+                    )
+                )
+
+            try:
+                result = fn(*bound.args, **bound.kwargs)
+                # Success: optionally clear.
+                if not keep_on_success:
+                    store.clear()
+                else:
+                    store.save(cp)
+                return result
+            except BaseException as exc:
+                # Capture the failure with locals from the failing frame.
+                locals_snapshot = _grab_failure_locals(exc, fn)
+                cp.record_failure(
+                    index=ctx.current_index if ctx.current_index is not None else -1,
+                    exc=exc,
+                    locals_snapshot=locals_snapshot,
+                )
+                store.save(cp)
+                if verbose:
+                    prog = cp.progress()
+                    print(
+                        f"[safe_state] {job_name!r} failed at item "
+                        f"{ctx.current_index}: {type(exc).__name__}: {exc}. "
+                        f"Progress {prog['completed']}/{prog['total']} saved to {path}",
+                        file=sys.stderr,
+                    )
+                raise
+            finally:
+                _active_contexts.pop(id(fn.__code__), None)
+
+        wrapper.checkpoint_path = path  # type: ignore[attr-defined]
+        wrapper.job_name = job_name  # type: ignore[attr-defined]
+        wrapper.peek_checkpoint = lambda: CheckpointStore(path, registry=reg).load() if path.exists() else None  # type: ignore[attr-defined]
+        wrapper.clear_checkpoint = lambda: CheckpointStore(path, registry=reg).clear()  # type: ignore[attr-defined]
+        return wrapper
+
+    if _fn is not None and callable(_fn):
+        return decorate(_fn)
+    return decorate
+
+
+def _enumerate_safe(
+    iterable: Iterable[Any],
+    ctx: _RunContext,
+    save_every: int,
+    store_results: bool,
+) -> Iterator[Tuple[int, Any]]:
+    """Internal: yield (idx, item) pairs, skipping completed indices and checkpointing."""
+    cp = ctx.checkpoint
+    items = list(iterable)
+    if cp.total_items is None:
+        cp.total_items = len(items)
+    since_save = 0
+    for idx, item in enumerate(items):
+        if idx in cp.completed_indices:
+            if ctx.verbose:
+                print(f"[safe_state] skip index {idx} (done)", file=sys.stderr)
+            continue
+        ctx.current_index = idx
+        yield idx, item
+        cp.mark_complete(idx, store_result=store_results)
+        since_save += 1
+        if since_save >= save_every:
+            ctx.store.save(cp)
+            since_save = 0
+    if since_save > 0:
+        ctx.store.save(cp)
+
+
+def _grab_failure_locals(exc: BaseException, decorated_fn: Callable) -> Dict[str, Any]:
+    """
+    Pull locals from the frame where the exception was raised, filtered to
+    serializable-looking values. Falls back to the decorated function's frame
+    if available.
+    """
+    tb = exc.__traceback__
+    target_frame = None
+    while tb is not None:
+        if tb.tb_frame.f_code is decorated_fn.__code__:
+            target_frame = tb.tb_frame
+        tb = tb.tb_next
+
+    if target_frame is None:
+        # Use the deepest frame as fallback.
+        tb = exc.__traceback__
+        while tb is not None:
+            target_frame = tb.tb_frame
+            tb = tb.tb_next
+
+    if target_frame is None:
+        return {}
+
+    out: Dict[str, Any] = {}
+    for k, v in target_frame.f_locals.items():
+        if k.startswith("__"):
+            continue
+        out[k] = v
+    return out

safe_state/exceptions.py

@@ -0,0 +1,21 @@
+"""Custom exceptions for safe-state."""
+
+
+class SafeStateError(Exception):
+    """Base exception for all safe-state errors."""
+
+
+class StateNotFoundError(SafeStateError):
+    """Raised when attempting to resume from a checkpoint that doesn't exist."""
+
+
+class StateCorruptedError(SafeStateError):
+    """Raised when a checkpoint file cannot be deserialized."""
+
+
+class ReconnectError(SafeStateError):
+    """Raised when a live object cannot be re-established on resume."""
+
+
+class NoReconnectorError(SafeStateError):
+    """Raised when a live object is encountered with no registered reconnect handler."""