checkpointer 2.8.1__tar.gz → 2.9.0__tar.gz

{checkpointer-2.8.1 → checkpointer-2.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.8.1
+Version: 2.9.0
 Summary: A Python library for memoizing function results with support for multiple storage backends, async runtimes, and automatic cache invalidation
 Project-URL: Repository, https://github.com/Reddan/checkpointer.git
 Author: Hampus Hallman
@@ -203,11 +203,11 @@ from checkpointer import checkpoint, Storage
 from datetime import datetime
 
 class CustomStorage(Storage):
-    def exists(self, path) -> bool: ...  # Check if a checkpoint exists
-    def checkpoint_date(self, path) -> datetime: ...  # Get the checkpoint's timestamp
-    def store(self, path, data): ...  # Save data to the checkpoint
-    def load(self, path): ...  # Load data from the checkpoint
-    def delete(self, path): ...  # Delete the checkpoint
+    def exists(self, call_id) -> bool: ...  # Check if a checkpoint exists
+    def checkpoint_date(self, call_id) -> datetime: ...  # Get the checkpoint's timestamp
+    def store(self, call_id, data): ...  # Save data to the checkpoint
+    def load(self, call_id): ...  # Load data from the checkpoint
+    def delete(self, call_id): ...  # Delete the checkpoint
 
 @checkpoint(format=CustomStorage)
 def custom_cached(x: int):

{checkpointer-2.8.1 → checkpointer-2.9.0}/README.md

@@ -183,11 +183,11 @@ from checkpointer import checkpoint, Storage
 from datetime import datetime
 
 class CustomStorage(Storage):
-    def exists(self, path) -> bool: ...  # Check if a checkpoint exists
-    def checkpoint_date(self, path) -> datetime: ...  # Get the checkpoint's timestamp
-    def store(self, path, data): ...  # Save data to the checkpoint
-    def load(self, path): ...  # Load data from the checkpoint
-    def delete(self, path): ...  # Delete the checkpoint
+    def exists(self, call_id) -> bool: ...  # Check if a checkpoint exists
+    def checkpoint_date(self, call_id) -> datetime: ...  # Get the checkpoint's timestamp
+    def store(self, call_id, data): ...  # Save data to the checkpoint
+    def load(self, call_id): ...  # Load data from the checkpoint
+    def delete(self, call_id): ...  # Delete the checkpoint
 
 @checkpoint(format=CustomStorage)
 def custom_cached(x: int):

{checkpointer-2.8.1 → checkpointer-2.9.0}/checkpointer/checkpoint.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 import inspect
 import re
+from contextlib import suppress
 from datetime import datetime
-from functools import update_wrapper
+from functools import cached_property, update_wrapper
 from pathlib import Path
-from typing import Any, Awaitable, Callable, Generic, Iterable, Literal, ParamSpec, Type, TypedDict, TypeVar, Unpack, cast, overload
+from typing import Awaitable, Callable, Generic, Iterable, Literal, ParamSpec, Type, TypedDict, TypeVar, Unpack, cast, overload
 from .fn_ident import get_fn_ident
 from .object_hash import ObjectHash
 from .print_checkpoint import print_checkpoint
@@ -54,84 +55,83 @@ class Checkpointer:
 
 class CheckpointFn(Generic[Fn]):
   def __init__(self, checkpointer: Checkpointer, fn: Fn):
-    self.checkpointer = checkpointer
-    self.fn = fn
-
-  def _set_ident(self, force=False):
-    if not hasattr(self, "fn_hash_raw") or force:
-      self.fn_hash_raw, self.depends = get_fn_ident(unwrap_fn(self.fn), self.checkpointer.capture)
-    return self
-
-  def _lazyinit(self):
-    params = self.checkpointer
-    wrapped = unwrap_fn(self.fn)
+    wrapped = unwrap_fn(fn)
     fn_file = Path(wrapped.__code__.co_filename).name
     fn_name = re.sub(r"[^\w.]", "", wrapped.__qualname__)
+    Storage = STORAGE_MAP[checkpointer.format] if isinstance(checkpointer.format, str) else checkpointer.format
     update_wrapper(cast(Callable, self), wrapped)
-    Storage = STORAGE_MAP[params.format] if isinstance(params.format, str) else params.format
-    deep_hashes = [child._set_ident().fn_hash_raw for child in iterate_checkpoint_fns(self)]
-    self.fn_hash = str(params.fn_hash or ObjectHash(digest_size=16).write_text(self.fn_hash_raw, *deep_hashes))
-    self.fn_subdir = f"{fn_file}/{fn_name}/{self.fn_hash[:32]}"
+    self.checkpointer = checkpointer
+    self.fn = fn
     self.storage = Storage(self)
     self.cleanup = self.storage.cleanup
+    self.fn_dir = f"{fn_file}/{fn_name}"
+
+  @cached_property
+  def ident_tuple(self) -> tuple[str, list[Callable]]:
+    return get_fn_ident(unwrap_fn(self.fn), self.checkpointer.capture)
+
+  @property
+  def fn_hash_raw(self) -> str:
+    return self.ident_tuple[0]
 
-  def __getattribute__(self, name: str) -> Any:
-    return object.__getattribute__(self, "_getattribute")(name)
+  @property
+  def depends(self) -> list[Callable]:
+    return self.ident_tuple[1]
 
-  def _getattribute(self, name: str) -> Any:
-    setattr(self, "_getattribute", super().__getattribute__)
-    self._lazyinit()
-    return self._getattribute(name)
+  @cached_property
+  def fn_hash(self) -> str:
+    fn_hash = self.checkpointer.fn_hash
+    deep_hashes = [depend.fn_hash_raw for depend in self.deep_depends()]
+    return str(fn_hash or ObjectHash(digest_size=16).write_text(self.fn_hash_raw, *deep_hashes))[:32]
 
   def reinit(self, recursive=False) -> CheckpointFn[Fn]:
-    pointfns = list(iterate_checkpoint_fns(self)) if recursive else [self]
-    for pointfn in pointfns:
-      pointfn._set_ident(True)
-    for pointfn in pointfns:
-      pointfn._lazyinit()
+    depends = list(self.deep_depends()) if recursive else [self]
+    for depend in depends:
+      with suppress(AttributeError):
+        del depend.ident_tuple, depend.fn_hash
+    for depend in depends:
+      depend.fn_hash
     return self
 
-  def get_checkpoint_id(self, args: tuple, kw: dict) -> str:
+  def get_call_id(self, args: tuple, kw: dict) -> str:
     hash_by = self.checkpointer.hash_by
     hash_params = hash_by(*args, **kw) if hash_by else (args, kw)
-    call_hash = ObjectHash(hash_params, digest_size=16)
-    return f"{self.fn_subdir}/{call_hash}"
+    return str(ObjectHash(hash_params, digest_size=16))
 
-  async def _resolve_awaitable(self, checkpoint_path: Path, awaitable: Awaitable):
+  async def _resolve_awaitable(self, checkpoint_id: str, awaitable: Awaitable):
     data = await awaitable
-    self.storage.store(checkpoint_path, AwaitableValue(data))
+    self.storage.store(checkpoint_id, AwaitableValue(data))
     return data
 
-  def _store_on_demand(self, args: tuple, kw: dict, rerun: bool):
+  def _call(self, args: tuple, kw: dict, rerun=False):
     params = self.checkpointer
-    checkpoint_id = self.get_checkpoint_id(args, kw)
-    checkpoint_path = params.root_path / checkpoint_id
+    if not params.when:
+      return self.fn(*args, **kw)
+
+    call_id = self.get_call_id(args, kw)
+    call_id_long = f"{self.fn_dir}/{self.fn_hash}/{call_id}"
+
     refresh = rerun \
-      or not self.storage.exists(checkpoint_path) \
-      or (params.should_expire and params.should_expire(self.storage.checkpoint_date(checkpoint_path)))
+      or not self.storage.exists(call_id) \
+      or (params.should_expire and params.should_expire(self.storage.checkpoint_date(call_id)))
 
     if refresh:
-      print_checkpoint(params.verbosity >= 1, "MEMORIZING", checkpoint_id, "blue")
+      print_checkpoint(params.verbosity >= 1, "MEMORIZING", call_id_long, "blue")
       data = self.fn(*args, **kw)
       if inspect.isawaitable(data):
-        return self._resolve_awaitable(checkpoint_path, data)
+        return self._resolve_awaitable(call_id, data)
       else:
-        self.storage.store(checkpoint_path, data)
+        self.storage.store(call_id, data)
         return data
 
     try:
-      data = self.storage.load(checkpoint_path)
-      print_checkpoint(params.verbosity >= 2, "REMEMBERED", checkpoint_id, "green")
+      data = self.storage.load(call_id)
+      print_checkpoint(params.verbosity >= 2, "REMEMBERED", call_id_long, "green")
       return data
     except (EOFError, FileNotFoundError):
       pass
-    print_checkpoint(params.verbosity >= 1, "CORRUPTED", checkpoint_id, "yellow")
-    return self._store_on_demand(args, kw, True)
-
-  def _call(self, args: tuple, kw: dict, rerun=False):
-    if not self.checkpointer.when:
-      return self.fn(*args, **kw)
-    return self._store_on_demand(args, kw, rerun)
+    print_checkpoint(params.verbosity >= 1, "CORRUPTED", call_id_long, "yellow")
+    return self._call(args, kw, True)
 
   __call__: Fn = cast(Fn, lambda self, *args, **kw: self._call(args, kw))
   rerun: Fn = cast(Fn, lambda self, *args, **kw: self._call(args, kw, True))
@@ -142,25 +142,29 @@ class CheckpointFn(Generic[Fn]):
   def get(self: Callable[P, R], *args: P.args, **kw: P.kwargs) -> R: ...
 
   def get(self, *args, **kw):
-    checkpoint_path = self.checkpointer.root_path / self.get_checkpoint_id(args, kw)
+    call_id = self.get_call_id(args, kw)
     try:
-      data = self.storage.load(checkpoint_path)
+      data = self.storage.load(call_id)
       return data.value if isinstance(data, AwaitableValue) else data
     except Exception as ex:
       raise CheckpointError("Could not load checkpoint") from ex
 
   def exists(self: Callable[P, R], *args: P.args, **kw: P.kwargs) -> bool: # type: ignore
     self = cast(CheckpointFn, self)
-    return self.storage.exists(self.checkpointer.root_path / self.get_checkpoint_id(args, kw))
+    return self.storage.exists(self.get_call_id(args, kw))
+
+  def delete(self: Callable[P, R], *args: P.args, **kw: P.kwargs): # type: ignore
+    self = cast(CheckpointFn, self)
+    self.storage.delete(self.get_call_id(args, kw))
 
   def __repr__(self) -> str:
     return f"<CheckpointFn {self.fn.__name__} {self.fn_hash[:6]}>"
 
-def iterate_checkpoint_fns(pointfn: CheckpointFn, visited: set[CheckpointFn] = set()) -> Iterable[CheckpointFn]:
-  visited = visited or set()
-  if pointfn not in visited:
-    yield pointfn
-    visited.add(pointfn)
-    for depend in pointfn.depends:
-      if isinstance(depend, CheckpointFn):
-        yield from iterate_checkpoint_fns(depend, visited)
+  def deep_depends(self, visited: set[CheckpointFn] = set()) -> Iterable[CheckpointFn]:
+    if self not in visited:
+      yield self
+      visited = visited or set()
+      visited.add(self)
+      for depend in self.depends:
+        if isinstance(depend, CheckpointFn):
+          yield from depend.deep_depends(visited)

checkpointer-2.9.0/checkpointer/storages/memory_storage.py

@@ -0,0 +1,36 @@
+from typing import Any
+from pathlib import Path
+from datetime import datetime
+from .storage import Storage
+
+item_map: dict[Path, dict[str, tuple[datetime, Any]]] = {}
+
+class MemoryStorage(Storage):
+  def get_dict(self):
+    return item_map.setdefault(self.dir(), {})
+
+  def store(self, call_id, data):
+    self.get_dict()[call_id] = (datetime.now(), data)
+
+  def exists(self, call_id):
+    return call_id in self.get_dict()
+
+  def checkpoint_date(self, call_id):
+    return self.get_dict()[call_id][0]
+
+  def load(self, call_id):
+    return self.get_dict()[call_id][1]
+
+  def delete(self, call_id):
+    self.get_dict().pop(call_id, None)
+
+  def cleanup(self, invalidated=True, expired=True):
+    curr_key = self.dir()
+    for key, calldict in list(item_map.items()):
+      if key.parent == curr_key.parent:
+        if invalidated and key != curr_key:
+          del item_map[key]
+        elif expired and self.checkpointer.should_expire:
+          for call_id, (date, _) in list(calldict.items()):
+            if self.checkpointer.should_expire(date):
+              del calldict[call_id]

{checkpointer-2.8.1 → checkpointer-2.9.0}/checkpointer/storages/pickle_storage.py

@@ -1,35 +1,34 @@
 import pickle
 import shutil
-from pathlib import Path
 from datetime import datetime
 from .storage import Storage
 
-def get_path(path: Path):
-  return path.with_name(f"{path.name}.pkl")
-
 class PickleStorage(Storage):
-  def store(self, path, data):
-    full_path = get_path(path)
-    full_path.parent.mkdir(parents=True, exist_ok=True)
-    with full_path.open("wb") as file:
+  def get_path(self, call_id: str):
+    return self.dir() / f"{call_id}.pkl"
+
+  def store(self, call_id, data):
+    path = self.get_path(call_id)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("wb") as file:
       pickle.dump(data, file, -1)
 
-  def exists(self, path):
-    return get_path(path).exists()
+  def exists(self, call_id):
+    return self.get_path(call_id).exists()
 
-  def checkpoint_date(self, path):
+  def checkpoint_date(self, call_id):
     # Should use st_atime/access time?
-    return datetime.fromtimestamp(get_path(path).stat().st_mtime)
+    return datetime.fromtimestamp(self.get_path(call_id).stat().st_mtime)
 
-  def load(self, path):
-    with get_path(path).open("rb") as file:
+  def load(self, call_id):
+    with self.get_path(call_id).open("rb") as file:
       return pickle.load(file)
 
-  def delete(self, path):
-    get_path(path).unlink(missing_ok=True)
+  def delete(self, call_id):
+    self.get_path(call_id).unlink(missing_ok=True)
 
   def cleanup(self, invalidated=True, expired=True):
-    version_path = self.checkpointer.root_path.resolve() / self.checkpoint_fn.fn_subdir
+    version_path = self.dir()
     fn_path = version_path.parent
     if invalidated:
       old_dirs = [path for path in fn_path.iterdir() if path.is_dir() and path != version_path]
@@ -39,8 +38,7 @@ class PickleStorage(Storage):
     if expired and self.checkpointer.should_expire:
       count = 0
       for pkl_path in fn_path.rglob("*.pkl"):
-        path = pkl_path.with_suffix("")
-        if self.checkpointer.should_expire(self.checkpoint_date(path)):
+        if self.checkpointer.should_expire(self.checkpoint_date(pkl_path.stem)):
           count += 1
-          self.delete(path)
+          self.delete(pkl_path.stem)
       print(f"Removed {count} expired checkpoints for {self.checkpoint_fn.__qualname__}")

{checkpointer-2.8.1 → checkpointer-2.9.0}/checkpointer/storages/storage.py

@@ -14,14 +14,17 @@ class Storage:
     self.checkpointer = checkpoint_fn.checkpointer
     self.checkpoint_fn = checkpoint_fn
 
-  def store(self, path: Path, data: Any) -> None: ...
+  def dir(self) -> Path:
+    return self.checkpointer.root_path / self.checkpoint_fn.fn_dir / self.checkpoint_fn.fn_hash
 
-  def exists(self, path: Path) -> bool: ...
+  def store(self, call_id: str, data: Any) -> None: ...
 
-  def checkpoint_date(self, path: Path) -> datetime: ...
+  def exists(self, call_id: str) -> bool: ...
 
-  def load(self, path: Path) -> Any: ...
+  def checkpoint_date(self, call_id: str) -> datetime: ...
 
-  def delete(self, path: Path) -> None: ...
+  def load(self, call_id: str) -> Any: ...
+
+  def delete(self, call_id: str) -> None: ...
 
   def cleanup(self, invalidated=True, expired=True): ...

{checkpointer-2.8.1 → checkpointer-2.9.0}/checkpointer/test_checkpointer.py

@@ -1,3 +1,10 @@
+"""
+TODO: Add tests for:
+- Checkpointing with different formats (pickle, memory, etc.)
+- Classes and methods - instances and classes
+- reinit deep depends
+"""
+
 import asyncio
 import pytest
 from riprint import riprint as print
@@ -142,7 +149,7 @@ def test_depends():
   assert set(test_a.depends) == {test_a.fn, helper, multiply_wrapper, global_multiply}
   assert set(test_b.depends) == {test_b.fn, test_a, multiply_wrapper, global_multiply}
 
-def test_lazy_init():
+def test_lazy_init_1():
   @checkpoint
   def fn1(x: object) -> object:
     return fn2(x)
@@ -151,10 +158,21 @@ def test_lazy_init():
   def fn2(x: object) -> object:
     return fn1(x)
 
-  assert type(object.__getattribute__(fn1, "_getattribute")) is MethodType
-  with pytest.raises(AttributeError):
-    object.__getattribute__(fn1, "fn_hash")
-  assert fn1.fn_hash == object.__getattribute__(fn1, "fn_hash")
-  assert type(object.__getattribute__(fn1, "_getattribute")) is MethodWrapperType
+  assert set(fn1.depends) == {fn1.fn, fn2}
+  assert set(fn2.depends) == {fn1, fn2.fn}
+
+def test_lazy_init_2():
+  @checkpoint
+  def fn1(x: object) -> object:
+    return fn2(x)
+
+  assert set(fn1.depends) == {fn1.fn}
+
+  @checkpoint
+  def fn2(x: object) -> object:
+    return fn1(x)
+
+  assert set(fn1.depends) == {fn1.fn}
+  fn1.reinit()
   assert set(fn1.depends) == {fn1.fn, fn2}
   assert set(fn2.depends) == {fn1, fn2.fn}

{checkpointer-2.8.1 → checkpointer-2.9.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.8.1"
+version = "2.9.0"
 requires-python = ">=3.11"
 dependencies = []
 authors = [

{checkpointer-2.8.1 → checkpointer-2.9.0}/uv.lock

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.8.1"
+version = "2.9.0"
 source = { editable = "." }
 
 [package.dev-dependencies]

checkpointer-2.8.1/checkpointer/storages/memory_storage.py

@@ -1,39 +0,0 @@
-from typing import Any
-from pathlib import Path
-from datetime import datetime
-from .storage import Storage
-
-item_map: dict[Path, dict[str, tuple[datetime, Any]]] = {}
-
-def get_short_path(path: Path):
-  return path.parts[-1]
-
-class MemoryStorage(Storage):
-  def get_dict(self):
-    return item_map.setdefault(self.checkpointer.root_path / self.checkpoint_fn.fn_subdir, {})
-
-  def store(self, path, data):
-    self.get_dict()[get_short_path(path)] = (datetime.now(), data)
-
-  def exists(self, path):
-    return get_short_path(path) in self.get_dict()
-
-  def checkpoint_date(self, path):
-    return self.get_dict()[get_short_path(path)][0]
-
-  def load(self, path):
-    return self.get_dict()[get_short_path(path)][1]
-
-  def delete(self, path):
-    del self.get_dict()[get_short_path(path)]
-
-  def cleanup(self, invalidated=True, expired=True):
-    curr_key = self.checkpointer.root_path / self.checkpoint_fn.fn_subdir
-    for key, calldict in list(item_map.items()):
-      if key.parent == curr_key.parent:
-        if invalidated and key != curr_key:
-          del item_map[key]
-        elif expired and self.checkpointer.should_expire:
-          for callid, (date, _) in list(calldict.items()):
-            if self.checkpointer.should_expire(date):
-              del calldict[callid]