checkpointer 2.9.1__tar.gz → 2.10.0__tar.gz

{checkpointer-2.9.1 → checkpointer-2.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.9.1
+Version: 2.10.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
@@ -48,19 +48,14 @@ result = expensive_function(4)  # Loads from the cache
 
 ## 🧠 How It Works
 
-When a function decorated with `@checkpoint` is called:
+When a `@checkpoint`-decorated function is called, `checkpointer` first computes a unique identifier (hash) for the call. This hash is derived from the function's source code, its dependencies, and the arguments passed.
 
-1.  `checkpointer` computes a unique identifier (hash) for the function call based on its source code, its dependencies, and the arguments passed.
-2.  It attempts to retrieve a cached result using this identifier.
-3.  If a cached result is found, it's returned immediately.
-4.  If no cached result exists or the cache has expired, the original function is executed, its result is stored, and then returned.
+It then tries to retrieve a cached result using this ID. If a valid cached result is found, it's returned immediately. Otherwise, the original function executes, its result is stored, and then returned.
 
-### ♻️ Automatic Cache Invalidation
+Cache validity is determined by this function's hash, which automatically updates if:
 
-`checkpointer` ensures caches are invalidated automatically when the underlying computation changes. A function's hash, which determines cache validity, updates if:
-
-* **Function Code Changes**: The source code of the decorated function itself is modified.
-* **Dependencies Change**: Any user-defined function in its dependency tree (direct or indirect, even across modules or not decorated with `@checkpoint`) is modified.
+* **Function Code Changes**: The decorated function's source code is modified.
+* **Dependencies Change**: Any user-defined function in its dependency tree (direct or indirect, even across modules or not decorated) is modified.
 * **Captured Variables Change** (with `capture=True`): Global or closure-based variables used within the function are altered.
 
 **Example: Dependency Invalidation**
@@ -76,63 +71,63 @@ def helper(x):
 
 @checkpoint
 def compute(a, b):
-    # Depends on `helper`
+    # Depends on `helper` and `multiply`
     return helper(a) + helper(b)
 ```
 
-If `multiply` is modified, caches for both `helper` and `compute` will automatically be invalidated and recomputed upon their next call.
+If `multiply` is modified, caches for both `helper` and `compute` are automatically invalidated and recomputed.
 
 ## 💡 Usage
 
 Once a function is decorated with `@checkpoint`, you can interact with its caching behavior using the following methods:
 
-* **`expensive_function(...)`**:
+* **`expensive_function(...)`**:\
     Call the function normally. This will either compute and cache the result or load it from the cache if available.
 
-* **`expensive_function.rerun(...)`**:
+* **`expensive_function.rerun(...)`**:\
     Forces the original function to execute, compute a new result, and overwrite any existing cached value for the given arguments.
 
-* **`expensive_function.fn(...)`**:
+* **`expensive_function.fn(...)`**:\
     Calls the original, undecorated function directly, bypassing the cache entirely. This is particularly useful within recursive functions to prevent caching intermediate steps.
 
-* **`expensive_function.get(...)`**:
+* **`expensive_function.get(...)`**:\
     Attempts to retrieve the cached result for the given arguments without executing the original function. Raises `CheckpointError` if no valid cached result exists.
 
-* **`expensive_function.exists(...)`**:
+* **`expensive_function.exists(...)`**:\
     Checks if a cached result exists for the given arguments without attempting to compute or load it. Returns `True` if a valid checkpoint exists, `False` otherwise.
 
-* **`expensive_function.delete(...)`**:
+* **`expensive_function.delete(...)`**:\
     Removes the cached entry for the specified arguments.
 
-* **`expensive_function.reinit()`**:
+* **`expensive_function.reinit()`**:\
     Recalculates the function's internal hash. This is primarily used when `capture=True` and you need to update the cache based on changes to external variables within the same Python session.
 
 ## ⚙️ Configuration & Customization
 
 The `@checkpoint` decorator accepts the following parameters to customize its behavior:
 
-* **`format`** (Type: `str` or `checkpointer.Storage`, Default: `"pickle"`)
+* **`format`** (Type: `str` or `checkpointer.Storage`, Default: `"pickle"`)\
     Defines the storage backend to use. Built-in options are `"pickle"` (disk-based, persistent) and `"memory"` (in-memory, non-persistent). You can also provide a custom `Storage` class.
 
-* **`root_path`** (Type: `str` or `pathlib.Path` or `None`, Default: `~/.cache/checkpoints`)
+* **`root_path`** (Type: `str` or `pathlib.Path` or `None`, Default: `~/.cache/checkpoints`)\
     The base directory for storing disk-based checkpoints. This parameter is only relevant when `format` is set to `"pickle"`.
 
-* **`when`** (Type: `bool`, Default: `True`)
+* **`when`** (Type: `bool`, Default: `True`)\
     A boolean flag to enable or disable checkpointing for the decorated function. This is particularly useful for toggling caching based on environment variables (e.g., `when=os.environ.get("ENABLE_CACHING", "false").lower() == "true"`).
 
-* **`capture`** (Type: `bool`, Default: `False`)
+* **`capture`** (Type: `bool`, Default: `False`)\
     If set to `True`, `checkpointer` includes global or closure-based variables used by the function in its hash calculation. This ensures that changes to these external variables also trigger cache invalidation and recomputation.
 
-* **`should_expire`** (Type: `Callable[[datetime.datetime], bool]`, Default: `None`)
+* **`should_expire`** (Type: `Callable[[datetime.datetime], bool]`, Default: `None`)\
     A custom callable that receives the `datetime` timestamp of a cached result. It should return `True` if the cached result is considered expired and needs recomputation, or `False` otherwise.
 
-* **`hash_by`** (Type: `Callable[..., Any]`, Default: `None`)
+* **`hash_by`** (Type: `Callable[..., Any]`, Default: `None`)\
     A custom callable that takes the function's arguments (`*args`, `**kwargs`) and returns a hashable object (or tuple of objects). This allows for custom argument normalization (e.g., sorting lists before hashing) or optimized hashing for complex input types, which can improve cache hit rates or speed up the hashing process.
 
-* **`fn_hash`** (Type: `checkpointer.ObjectHash`, Default: `None`)
+* **`fn_hash`** (Type: `checkpointer.ObjectHash`, Default: `None`)\
     An optional parameter that takes an instance of `checkpointer.ObjectHash`. This allows you to override the automatically computed function hash, giving you explicit control over when the function's cache should be invalidated. You can pass any values relevant to your invalidation logic to `ObjectHash` (e.g., `ObjectHash(version_string, config_id, ...)`, as it can consistently hash most Python values.
 
-* **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)
+* **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)\
     Controls the level of logging output from `checkpointer`.
     * `0`: No output.
     * `1`: Shows when functions are computed and cached.
@@ -156,8 +151,11 @@ class MyCustomStorage(Storage):
         fn_dir = self.checkpointer.root_path / self.fn_id()
         return (fn_dir / call_id).exists()
 
+    def store(self, call_id, data):
+        ... # Store the serialized data for `call_id`
+        return data # This method must return the data back to checkpointer
+
     def checkpoint_date(self, call_id): ...
-    def store(self, call_id, data): ...
     def load(self, call_id): ...
     def delete(self, call_id): ...
 

{checkpointer-2.9.1 → checkpointer-2.10.0}/README.md

@@ -28,19 +28,14 @@ result = expensive_function(4)  # Loads from the cache
 
 ## 🧠 How It Works
 
-When a function decorated with `@checkpoint` is called:
+When a `@checkpoint`-decorated function is called, `checkpointer` first computes a unique identifier (hash) for the call. This hash is derived from the function's source code, its dependencies, and the arguments passed.
 
-1.  `checkpointer` computes a unique identifier (hash) for the function call based on its source code, its dependencies, and the arguments passed.
-2.  It attempts to retrieve a cached result using this identifier.
-3.  If a cached result is found, it's returned immediately.
-4.  If no cached result exists or the cache has expired, the original function is executed, its result is stored, and then returned.
+It then tries to retrieve a cached result using this ID. If a valid cached result is found, it's returned immediately. Otherwise, the original function executes, its result is stored, and then returned.
 
-### ♻️ Automatic Cache Invalidation
+Cache validity is determined by this function's hash, which automatically updates if:
 
-`checkpointer` ensures caches are invalidated automatically when the underlying computation changes. A function's hash, which determines cache validity, updates if:
-
-* **Function Code Changes**: The source code of the decorated function itself is modified.
-* **Dependencies Change**: Any user-defined function in its dependency tree (direct or indirect, even across modules or not decorated with `@checkpoint`) is modified.
+* **Function Code Changes**: The decorated function's source code is modified.
+* **Dependencies Change**: Any user-defined function in its dependency tree (direct or indirect, even across modules or not decorated) is modified.
 * **Captured Variables Change** (with `capture=True`): Global or closure-based variables used within the function are altered.
 
 **Example: Dependency Invalidation**
@@ -56,63 +51,63 @@ def helper(x):
 
 @checkpoint
 def compute(a, b):
-    # Depends on `helper`
+    # Depends on `helper` and `multiply`
     return helper(a) + helper(b)
 ```
 
-If `multiply` is modified, caches for both `helper` and `compute` will automatically be invalidated and recomputed upon their next call.
+If `multiply` is modified, caches for both `helper` and `compute` are automatically invalidated and recomputed.
 
 ## 💡 Usage
 
 Once a function is decorated with `@checkpoint`, you can interact with its caching behavior using the following methods:
 
-* **`expensive_function(...)`**:
+* **`expensive_function(...)`**:\
     Call the function normally. This will either compute and cache the result or load it from the cache if available.
 
-* **`expensive_function.rerun(...)`**:
+* **`expensive_function.rerun(...)`**:\
     Forces the original function to execute, compute a new result, and overwrite any existing cached value for the given arguments.
 
-* **`expensive_function.fn(...)`**:
+* **`expensive_function.fn(...)`**:\
     Calls the original, undecorated function directly, bypassing the cache entirely. This is particularly useful within recursive functions to prevent caching intermediate steps.
 
-* **`expensive_function.get(...)`**:
+* **`expensive_function.get(...)`**:\
     Attempts to retrieve the cached result for the given arguments without executing the original function. Raises `CheckpointError` if no valid cached result exists.
 
-* **`expensive_function.exists(...)`**:
+* **`expensive_function.exists(...)`**:\
     Checks if a cached result exists for the given arguments without attempting to compute or load it. Returns `True` if a valid checkpoint exists, `False` otherwise.
 
-* **`expensive_function.delete(...)`**:
+* **`expensive_function.delete(...)`**:\
     Removes the cached entry for the specified arguments.
 
-* **`expensive_function.reinit()`**:
+* **`expensive_function.reinit()`**:\
     Recalculates the function's internal hash. This is primarily used when `capture=True` and you need to update the cache based on changes to external variables within the same Python session.
 
 ## ⚙️ Configuration & Customization
 
 The `@checkpoint` decorator accepts the following parameters to customize its behavior:
 
-* **`format`** (Type: `str` or `checkpointer.Storage`, Default: `"pickle"`)
+* **`format`** (Type: `str` or `checkpointer.Storage`, Default: `"pickle"`)\
     Defines the storage backend to use. Built-in options are `"pickle"` (disk-based, persistent) and `"memory"` (in-memory, non-persistent). You can also provide a custom `Storage` class.
 
-* **`root_path`** (Type: `str` or `pathlib.Path` or `None`, Default: `~/.cache/checkpoints`)
+* **`root_path`** (Type: `str` or `pathlib.Path` or `None`, Default: `~/.cache/checkpoints`)\
     The base directory for storing disk-based checkpoints. This parameter is only relevant when `format` is set to `"pickle"`.
 
-* **`when`** (Type: `bool`, Default: `True`)
+* **`when`** (Type: `bool`, Default: `True`)\
     A boolean flag to enable or disable checkpointing for the decorated function. This is particularly useful for toggling caching based on environment variables (e.g., `when=os.environ.get("ENABLE_CACHING", "false").lower() == "true"`).
 
-* **`capture`** (Type: `bool`, Default: `False`)
+* **`capture`** (Type: `bool`, Default: `False`)\
     If set to `True`, `checkpointer` includes global or closure-based variables used by the function in its hash calculation. This ensures that changes to these external variables also trigger cache invalidation and recomputation.
 
-* **`should_expire`** (Type: `Callable[[datetime.datetime], bool]`, Default: `None`)
+* **`should_expire`** (Type: `Callable[[datetime.datetime], bool]`, Default: `None`)\
     A custom callable that receives the `datetime` timestamp of a cached result. It should return `True` if the cached result is considered expired and needs recomputation, or `False` otherwise.
 
-* **`hash_by`** (Type: `Callable[..., Any]`, Default: `None`)
+* **`hash_by`** (Type: `Callable[..., Any]`, Default: `None`)\
     A custom callable that takes the function's arguments (`*args`, `**kwargs`) and returns a hashable object (or tuple of objects). This allows for custom argument normalization (e.g., sorting lists before hashing) or optimized hashing for complex input types, which can improve cache hit rates or speed up the hashing process.
 
-* **`fn_hash`** (Type: `checkpointer.ObjectHash`, Default: `None`)
+* **`fn_hash`** (Type: `checkpointer.ObjectHash`, Default: `None`)\
     An optional parameter that takes an instance of `checkpointer.ObjectHash`. This allows you to override the automatically computed function hash, giving you explicit control over when the function's cache should be invalidated. You can pass any values relevant to your invalidation logic to `ObjectHash` (e.g., `ObjectHash(version_string, config_id, ...)`, as it can consistently hash most Python values.
 
-* **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)
+* **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)\
     Controls the level of logging output from `checkpointer`.
     * `0`: No output.
     * `1`: Shows when functions are computed and cached.
@@ -136,8 +131,11 @@ class MyCustomStorage(Storage):
         fn_dir = self.checkpointer.root_path / self.fn_id()
         return (fn_dir / call_id).exists()
 
+    def store(self, call_id, data):
+        ... # Store the serialized data for `call_id`
+        return data # This method must return the data back to checkpointer
+
     def checkpoint_date(self, call_id): ...
-    def store(self, call_id, data): ...
     def load(self, call_id): ...
     def delete(self, call_id): ...
 

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/__init__.py

@@ -1,11 +1,11 @@
 import gc
 import tempfile
 from typing import Callable
-from .checkpoint import Checkpointer, CheckpointError, CheckpointFn
+from .checkpoint import CachedFunction, Checkpointer, CheckpointError
 from .object_hash import ObjectHash
 from .storages import MemoryStorage, PickleStorage, Storage
+from .utils import AwaitableValue
 
-create_checkpointer = Checkpointer
 checkpoint = Checkpointer()
 capture_checkpoint = Checkpointer(capture=True)
 memory_checkpoint = Checkpointer(format="memory", verbosity=0)
@@ -14,8 +14,8 @@ static_checkpoint = Checkpointer(fn_hash=ObjectHash())
 
 def cleanup_all(invalidated=True, expired=True):
   for obj in gc.get_objects():
-    if isinstance(obj, CheckpointFn):
+    if isinstance(obj, CachedFunction):
       obj.cleanup(invalidated=invalidated, expired=expired)
 
 def get_function_hash(fn: Callable, capture=False) -> str:
-  return CheckpointFn(Checkpointer(capture=capture), fn).fn_hash
+  return CachedFunction(Checkpointer(capture=capture), fn).fn_hash

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/checkpoint.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 import inspect
 import re
-from contextlib import suppress
 from datetime import datetime
 from functools import cached_property, update_wrapper
 from pathlib import Path
-from typing import Awaitable, Callable, Generic, Iterable, Literal, ParamSpec, Type, TypedDict, TypeVar, Unpack, cast, overload
+from typing import (
+  Awaitable, Callable, Concatenate, Generic, Iterable, Literal,
+  ParamSpec, Self, Type, TypedDict, TypeVar, Unpack, cast, overload,
+)
 from .fn_ident import get_fn_ident
 from .object_hash import ObjectHash
 from .print_checkpoint import print_checkpoint
@@ -15,6 +17,7 @@ from .utils import AwaitableValue, unwrap_fn
 Fn = TypeVar("Fn", bound=Callable)
 P = ParamSpec("P")
 R = TypeVar("R")
+C = TypeVar("C")
 
 DEFAULT_DIR = Path.home() / ".cache/checkpoints"
 
@@ -43,17 +46,17 @@ class Checkpointer:
     self.fn_hash = opts.get("fn_hash")
 
   @overload
-  def __call__(self, fn: Fn, **override_opts: Unpack[CheckpointerOpts]) -> CheckpointFn[Fn]: ...
+  def __call__(self, fn: Fn, **override_opts: Unpack[CheckpointerOpts]) -> CachedFunction[Fn]: ...
   @overload
   def __call__(self, fn: None=None, **override_opts: Unpack[CheckpointerOpts]) -> Checkpointer: ...
-  def __call__(self, fn: Fn | None=None, **override_opts: Unpack[CheckpointerOpts]) -> Checkpointer | CheckpointFn[Fn]:
+  def __call__(self, fn: Fn | None=None, **override_opts: Unpack[CheckpointerOpts]) -> Checkpointer | CachedFunction[Fn]:
     if override_opts:
       opts = CheckpointerOpts(**{**self.__dict__, **override_opts})
       return Checkpointer(**opts)(fn)
 
-    return CheckpointFn(self, fn) if callable(fn) else self
+    return CachedFunction(self, fn) if callable(fn) else self
 
-class CheckpointFn(Generic[Fn]):
+class CachedFunction(Generic[Fn]):
   def __init__(self, checkpointer: Checkpointer, fn: Fn):
     wrapped = unwrap_fn(fn)
     fn_file = Path(wrapped.__code__.co_filename).name
@@ -65,6 +68,19 @@ class CheckpointFn(Generic[Fn]):
     self.fn_dir = f"{fn_file}/{fn_name}"
     self.storage = Storage(self)
     self.cleanup = self.storage.cleanup
+    self.bound = ()
+
+  @overload
+  def __get__(self: Self, instance: None, owner: Type[C]) -> Self: ...
+  @overload
+  def __get__(self: CachedFunction[Callable[Concatenate[C, P], R]], instance: C, owner: Type[C]) -> CachedFunction[Callable[P, R]]: ...
+  def __get__(self, instance, owner):
+    if instance is None:
+      return self
+    bound_fn = object.__new__(CachedFunction)
+    bound_fn.__dict__ |= self.__dict__
+    bound_fn.bound = (instance,)
+    return bound_fn
 
   @cached_property
   def ident_tuple(self) -> tuple[str, list[Callable]]:
@@ -80,33 +96,33 @@ class CheckpointFn(Generic[Fn]):
 
   @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]
+    fn_hash = ObjectHash(digest_size=16).write_text(self.fn_hash_raw, *deep_hashes)
+    return str(self.checkpointer.fn_hash or fn_hash)[:32]
 
-  def reinit(self, recursive=False) -> CheckpointFn[Fn]:
+  def reinit(self, recursive=False) -> CachedFunction[Fn]:
     depends = list(self.deep_depends()) if recursive else [self]
     for depend in depends:
-      with suppress(AttributeError):
-        del depend.ident_tuple, depend.fn_hash
+      self.__dict__.pop("fn_hash", None)
+      self.__dict__.pop("ident_tuple", None)
     for depend in depends:
       depend.fn_hash
     return self
 
   def get_call_id(self, args: tuple, kw: dict) -> str:
+    args = self.bound + args
     hash_by = self.checkpointer.hash_by
     hash_params = hash_by(*args, **kw) if hash_by else (args, kw)
     return str(ObjectHash(hash_params, digest_size=16))
 
-  async def _resolve_awaitable(self, checkpoint_id: str, awaitable: Awaitable):
-    data = await awaitable
-    self.storage.store(checkpoint_id, AwaitableValue(data))
-    return data
+  async def _resolve_awaitable(self, call_id: str, awaitable: Awaitable):
+    return await self.storage.store(call_id, AwaitableValue(await awaitable))
 
-  def _call(self, args: tuple, kw: dict, rerun=False):
+  def _call(self: CachedFunction[Callable[P, R]], args: tuple, kw: dict, rerun=False) -> R:
+    full_args = self.bound + args
     params = self.checkpointer
     if not params.when:
-      return self.fn(*args, **kw)
+      return self.fn(*full_args, **kw)
 
     call_id = self.get_call_id(args, kw)
     call_id_long = f"{self.fn_dir}/{self.fn_hash}/{call_id}"
@@ -117,12 +133,10 @@ class CheckpointFn(Generic[Fn]):
 
     if refresh:
       print_checkpoint(params.verbosity >= 1, "MEMORIZING", call_id_long, "blue")
-      data = self.fn(*args, **kw)
+      data = self.fn(*full_args, **kw)
       if inspect.isawaitable(data):
         return self._resolve_awaitable(call_id, data)
-      else:
-        self.storage.store(call_id, data)
-        return data
+      return self.storage.store(call_id, data)
 
     try:
       data = self.storage.load(call_id)
@@ -133,14 +147,16 @@ class CheckpointFn(Generic[Fn]):
     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))
+  def __call__(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> R:
+    return self._call(args, kw)
+
+  def rerun(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> R:
+    return self._call(args, kw, True)
 
   @overload
   def get(self: Callable[P, Awaitable[R]], *args: P.args, **kw: P.kwargs) -> R: ...
   @overload
   def get(self: Callable[P, R], *args: P.args, **kw: P.kwargs) -> R: ...
-
   def get(self, *args, **kw):
     call_id = self.get_call_id(args, kw)
     try:
@@ -149,22 +165,20 @@ class CheckpointFn(Generic[Fn]):
     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)
+  def exists(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> bool:
     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)
+  def delete(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs):
     self.storage.delete(self.get_call_id(args, kw))
 
   def __repr__(self) -> str:
-    return f"<CheckpointFn {self.fn.__name__} {self.fn_hash[:6]}>"
+    return f"<CachedFunction {self.fn.__name__} {self.fn_hash[:6]}>"
 
-  def deep_depends(self, visited: set[CheckpointFn] = set()) -> Iterable[CheckpointFn]:
+  def deep_depends(self, visited: set[CachedFunction] = set()) -> Iterable[CachedFunction]:
     if self not in visited:
       yield self
       visited = visited or set()
       visited.add(self)
       for depend in self.depends:
-        if isinstance(depend, CheckpointFn):
+        if isinstance(depend, CachedFunction):
           yield from depend.deep_depends(visited)

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/fn_ident.py

@@ -8,7 +8,7 @@ from typing import Any, Iterable, Type, TypeGuard
 from .object_hash import ObjectHash
 from .utils import AttrDict, distinct, get_cell_contents, iterate_and_upcoming, transpose, unwrap_fn
 
-cwd = Path.cwd()
+cwd = Path.cwd().resolve()
 
 def is_class(obj) -> TypeGuard[Type]:
   # isinstance works too, but needlessly triggers _lazyinit()
@@ -72,23 +72,23 @@ def is_user_fn(candidate_fn) -> TypeGuard[Callable]:
   return cwd in fn_path.parents and ".venv" not in fn_path.parts
 
 def get_depend_fns(fn: Callable, capture: bool, captured_vals_by_fn: dict[Callable, list[Any]] = {}) -> dict[Callable, list[Any]]:
-  from .checkpoint import CheckpointFn
+  from .checkpoint import CachedFunction
   captured_vals_by_fn = captured_vals_by_fn or {}
   captured_vals = get_fn_captured_vals(fn)
   captured_vals_by_fn[fn] = [val for val in captured_vals if not callable(val)] * capture
-  child_fns = (unwrap_fn(val, checkpoint_fn=True) for val in captured_vals if callable(val))
+  child_fns = (unwrap_fn(val, cached_fn=True) for val in captured_vals if callable(val))
   for child_fn in child_fns:
-    if isinstance(child_fn, CheckpointFn):
+    if isinstance(child_fn, CachedFunction):
       captured_vals_by_fn[child_fn] = []
     elif child_fn not in captured_vals_by_fn and is_user_fn(child_fn):
       get_depend_fns(child_fn, capture, captured_vals_by_fn)
   return captured_vals_by_fn
 
 def get_fn_ident(fn: Callable, capture: bool) -> tuple[str, list[Callable]]:
-  from .checkpoint import CheckpointFn
+  from .checkpoint import CachedFunction
   captured_vals_by_fn = get_depend_fns(fn, capture)
   depends, depend_captured_vals = transpose(captured_vals_by_fn.items(), 2)
   depends = distinct(fn.__func__ if isinstance(fn, MethodType) else fn for fn in depends)
-  unwrapped_depends = [fn for fn in depends if not isinstance(fn, CheckpointFn)]
+  unwrapped_depends = [fn for fn in depends if not isinstance(fn, CachedFunction)]
   fn_hash = str(ObjectHash(fn, unwrapped_depends).update(depend_captured_vals, tolerate_errors=True))
   return fn_hash, depends

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/object_hash.py

@@ -180,8 +180,10 @@ class ObjectHash:
   def _update_iterator(self, obj: Iterable) -> "ObjectHash":
     return self.header("iterator", encode_type_of(obj)).update(iter=obj).header("iterator-end")
 
-  def _update_object(self, obj: object) -> "ObjectHash":
+  def _update_object(self, obj: Any) -> "ObjectHash":
     self.header("instance", encode_type_of(obj))
+    if hasattr(obj, "__objecthash__") and callable(obj.__objecthash__):
+      return self.header("objecthash").update(obj.__objecthash__())
     reduced = None
     with suppress(Exception):
       reduced = obj.__reduce_ex__(PROTOCOL)

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/storages/__init__.py

@@ -2,10 +2,8 @@ from typing import Type
 from .storage import Storage
 from .pickle_storage import PickleStorage
 from .memory_storage import MemoryStorage
-from .bcolz_storage import BcolzStorage
 
 STORAGE_MAP: dict[str, Type[Storage]] = {
   "pickle": PickleStorage,
   "memory": MemoryStorage,
-  "bcolz": BcolzStorage,
 }

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/storages/memory_storage.py

@@ -11,6 +11,7 @@ class MemoryStorage(Storage):
 
   def store(self, call_id, data):
     self.get_dict()[call_id] = (datetime.now(), data)
+    return data
 
   def exists(self, call_id):
     return call_id in self.get_dict()

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/storages/pickle_storage.py

@@ -1,8 +1,12 @@
 import pickle
 import shutil
 from datetime import datetime
+from pathlib import Path
 from .storage import Storage
 
+def filedate(path: Path) -> datetime:
+  return datetime.fromtimestamp(path.stat().st_mtime)
+
 class PickleStorage(Storage):
   def get_path(self, call_id: str):
     return self.fn_dir() / f"{call_id}.pkl"
@@ -12,13 +16,14 @@ class PickleStorage(Storage):
     path.parent.mkdir(parents=True, exist_ok=True)
     with path.open("wb") as file:
       pickle.dump(data, file, -1)
+    return data
 
   def exists(self, call_id):
     return self.get_path(call_id).exists()
 
   def checkpoint_date(self, call_id):
     # Should use st_atime/access time?
-    return datetime.fromtimestamp(self.get_path(call_id).stat().st_mtime)
+    return filedate(self.get_path(call_id))
 
   def load(self, call_id):
     with self.get_path(call_id).open("rb") as file:
@@ -34,11 +39,11 @@ class PickleStorage(Storage):
       old_dirs = [path for path in fn_path.iterdir() if path.is_dir() and path != version_path]
       for path in old_dirs:
         shutil.rmtree(path)
-      print(f"Removed {len(old_dirs)} invalidated directories for {self.checkpoint_fn.__qualname__}")
+      print(f"Removed {len(old_dirs)} invalidated directories for {self.cached_fn.__qualname__}")
     if expired and self.checkpointer.should_expire:
       count = 0
-      for pkl_path in fn_path.rglob("*.pkl"):
-        if self.checkpointer.should_expire(self.checkpoint_date(pkl_path.stem)):
+      for pkl_path in fn_path.glob("**/*.pkl"):
+        if self.checkpointer.should_expire(filedate(pkl_path)):
           count += 1
-          self.delete(pkl_path.stem)
-      print(f"Removed {count} expired checkpoints for {self.checkpoint_fn.__qualname__}")
+          pkl_path.unlink(missing_ok=True)
+      print(f"Removed {count} expired checkpoints for {self.cached_fn.__qualname__}")

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/storages/storage.py

@@ -4,23 +4,23 @@ from pathlib import Path
 from datetime import datetime
 
 if TYPE_CHECKING:
-  from ..checkpoint import Checkpointer, CheckpointFn
+  from ..checkpoint import Checkpointer, CachedFunction
 
 class Storage:
   checkpointer: Checkpointer
-  checkpoint_fn: CheckpointFn
+  cached_fn: CachedFunction
 
-  def __init__(self, checkpoint_fn: CheckpointFn):
-    self.checkpointer = checkpoint_fn.checkpointer
-    self.checkpoint_fn = checkpoint_fn
+  def __init__(self, cached_fn: CachedFunction):
+    self.checkpointer = cached_fn.checkpointer
+    self.cached_fn = cached_fn
 
   def fn_id(self) -> str:
-    return f"{self.checkpoint_fn.fn_dir}/{self.checkpoint_fn.fn_hash}"
+    return f"{self.cached_fn.fn_dir}/{self.cached_fn.fn_hash}"
 
   def fn_dir(self) -> Path:
     return self.checkpointer.root_path / self.fn_id()
 
-  def store(self, call_id: str, data: Any) -> None: ...
+  def store(self, call_id: str, data: Any) -> Any: ...
 
   def exists(self, call_id: str) -> bool: ...
 

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/test_checkpointer.py

@@ -1,7 +1,6 @@
 import asyncio
 import pytest
 from riprint import riprint as print
-from types import MethodType, MethodWrapperType
 from . import checkpoint
 from .checkpoint import CheckpointError
 from .utils import AttrDict

{checkpointer-2.9.1 → checkpointer-2.10.0}/checkpointer/utils.py

@@ -32,10 +32,10 @@ def get_cell_contents(fn: Callable) -> Iterable[tuple[str, Any]]:
     except ValueError:
       pass
 
-def unwrap_fn(fn: Fn, checkpoint_fn=False) -> Fn:
-  from .checkpoint import CheckpointFn
+def unwrap_fn(fn: Fn, cached_fn=False) -> Fn:
+  from .checkpoint import CachedFunction
   while True:
-    if (checkpoint_fn and isinstance(fn, CheckpointFn)) or not hasattr(fn, "__wrapped__"):
+    if (cached_fn and isinstance(fn, CachedFunction)) or not hasattr(fn, "__wrapped__"):
       return cast(Fn, fn)
     fn = getattr(fn, "__wrapped__")
 

{checkpointer-2.9.1 → checkpointer-2.10.0}/pyproject.toml

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

{checkpointer-2.9.1 → checkpointer-2.10.0}/uv.lock

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

checkpointer-2.9.1/checkpointer/storages/bcolz_storage.py

@@ -1,80 +0,0 @@
-import shutil
-from pathlib import Path
-from datetime import datetime
-from .storage import Storage
-
-def get_data_type_str(x):
-  if isinstance(x, tuple):
-    return "tuple"
-  elif isinstance(x, dict):
-    return "dict"
-  elif isinstance(x, list):
-    return "list"
-  elif isinstance(x, str) or not hasattr(x, "__len__"):
-    return "other"
-  else:
-    return "ndarray"
-
-def get_metapath(path: Path):
-  return path.with_name(f"{path.name}_meta")
-
-def insert_data(path: Path, data):
-  import bcolz
-  c = bcolz.carray(data, rootdir=path, mode="w")
-  c.flush()
-
-class BcolzStorage(Storage):
-  def exists(self, path):
-    return path.exists()
-
-  def checkpoint_date(self, path):
-    return datetime.fromtimestamp(path.stat().st_mtime)
-
-  def store(self, path, data):
-    metapath = get_metapath(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    data_type_str = get_data_type_str(data)
-    if data_type_str == "tuple":
-      fields = list(range(len(data)))
-    elif data_type_str == "dict":
-      fields = sorted(data.keys())
-    else:
-      fields = []
-    meta_data = {"data_type_str": data_type_str, "fields": fields}
-    insert_data(metapath, meta_data)
-    if data_type_str in ["tuple", "dict"]:
-      for i in range(len(fields)):
-        child_path = Path(f"{path} ({i})")
-        self.store(child_path, data[fields[i]])
-    else:
-      insert_data(path, data)
-
-  def load(self, path):
-    import bcolz
-    metapath = get_metapath(path)
-    meta_data = bcolz.open(metapath)[:][0]
-    data_type_str = meta_data["data_type_str"]
-    if data_type_str in ["tuple", "dict"]:
-      fields = meta_data["fields"]
-      partitions = range(len(fields))
-      data = [self.load(Path(f"{path} ({i})")) for i in partitions]
-      if data_type_str == "tuple":
-        return tuple(data)
-      else:
-        return dict(zip(fields, data))
-    else:
-      data = bcolz.open(path)
-      if data_type_str == "list":
-        return list(data)
-      elif data_type_str == "other":
-        return data[0]
-      else:
-        return data[:]
-
-  def delete(self, path):
-    # NOTE: Not recursive
-    shutil.rmtree(get_metapath(path), ignore_errors=True)
-    shutil.rmtree(path, ignore_errors=True)
-
-  def cleanup(self, invalidated=True, expired=True):
-    raise NotImplementedError("cleanup() not implemented for bcolz storage")