checkpointer 2.10.1__tar.gz → 2.11.1__tar.gz

{checkpointer-2.10.1 → checkpointer-2.11.1}/.gitignore

@@ -1,3 +1,4 @@
 __pycache__/
 /dist/
 .DS_Store
+/*.py

{checkpointer-2.10.1 → checkpointer-2.11.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.10.1
+Version: 2.11.1
 Summary: checkpointer adds code-aware caching to Python functions, maintaining correctness and speeding up execution as your code changes.
 Project-URL: Repository, https://github.com/Reddan/checkpointer.git
 Author: Hampus Hallman
@@ -121,11 +121,8 @@ The `@checkpoint` decorator accepts the following parameters to customize its be
 * **`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`)\
-    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`)\
-    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.
+* **`fn_hash_from`** (Type: `Any`, Default: `None`)\
+    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 object relevant to your invalidation logic (e.g., version strings, config parameters). The object you provide will be hashed internally by `checkpointer`.
 
 * **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)\
     Controls the level of logging output from `checkpointer`.
@@ -133,31 +130,63 @@ The `@checkpoint` decorator accepts the following parameters to customize its be
     * `1`: Shows when functions are computed and cached.
     * `2`: Also shows when cached results are remembered (loaded from cache).
 
-### 🗄️ Custom Storage Backends
+## 🔬 Customize Argument Hashing
+
+You can customize how individual function arguments are hashed without changing their actual values when passed in.
+
+* **`Annotated[T, HashBy[fn]]`**:\
+    Hashes the argument by applying `fn(argument)` before hashing. This enables custom normalization (e.g., sorting lists to ignore order) or optimized hashing for complex types, improving cache hit rates or speeding up hashing.
+
+* **`NoHash[T]`**:\
+    Completely excludes the argument from hashing, so changes to it won’t trigger cache invalidation.
+
+**Example:**
+
+```python
+from typing import Annotated
+from checkpointer import checkpoint, HashBy, NoHash
+from pathlib import Path
+import logging
+
+def file_bytes(path: Path) -> bytes:
+    return path.read_bytes()
+
+@checkpoint
+def process(
+    numbers: Annotated[list[int], HashBy[sorted]],   # Hash by sorted list
+    data_file: Annotated[Path, HashBy[file_bytes]],  # Hash by file content
+    log: NoHash[logging.Logger],                     # Exclude logger from hashing
+):
+    ...
+```
+
+In this example, the cache key for `numbers` ignores order, `data_file` is hashed based on its contents rather than path, and changes to `log` don’t affect caching.
+
+## 🗄️ Custom Storage Backends
 
 For integration with databases, cloud storage, or custom serialization, implement your own storage backend by inheriting from `checkpointer.Storage` and implementing its abstract methods.
 
-Within custom storage methods, `call_id` identifies calls by arguments. Use `self.fn_id()` to get the function's unique identity (name + hash/version), crucial for organizing stored checkpoints (e.g., by function version). Access global `Checkpointer` config via `self.checkpointer`.
+Within custom storage methods, `call_hash` identifies calls by arguments. Use `self.fn_id()` to get the function's unique identity (name + hash/version), crucial for organizing stored checkpoints (e.g., by function version). Access global `Checkpointer` config via `self.checkpointer`.
 
-#### Example: Custom Storage Backend
+**Example: Custom Storage Backend**
 
 ```python
 from checkpointer import checkpoint, Storage
 from datetime import datetime
 
 class MyCustomStorage(Storage):
-    def exists(self, call_id):
+    def exists(self, call_hash):
         # Example: Constructing a path based on function ID and call ID
         fn_dir = self.checkpointer.root_path / self.fn_id()
-        return (fn_dir / call_id).exists()
+        return (fn_dir / call_hash).exists()
 
-    def store(self, call_id, data):
-        ... # Store the serialized data for `call_id`
+    def store(self, call_hash, data):
+        ... # Store the serialized data for `call_hash`
         return data # This method must return the data back to checkpointer
 
-    def checkpoint_date(self, call_id): ...
-    def load(self, call_id): ...
-    def delete(self, call_id): ...
+    def checkpoint_date(self, call_hash): ...
+    def load(self, call_hash): ...
+    def delete(self, call_hash): ...
 
 @checkpoint(format=MyCustomStorage)
 def custom_cached_function(x: int):

{checkpointer-2.10.1 → checkpointer-2.11.1}/README.md

@@ -101,11 +101,8 @@ The `@checkpoint` decorator accepts the following parameters to customize its be
 * **`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`)\
-    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`)\
-    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.
+* **`fn_hash_from`** (Type: `Any`, Default: `None`)\
+    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 object relevant to your invalidation logic (e.g., version strings, config parameters). The object you provide will be hashed internally by `checkpointer`.
 
 * **`verbosity`** (Type: `int` (`0`, `1`, or `2`), Default: `1`)\
     Controls the level of logging output from `checkpointer`.
@@ -113,31 +110,63 @@ The `@checkpoint` decorator accepts the following parameters to customize its be
     * `1`: Shows when functions are computed and cached.
     * `2`: Also shows when cached results are remembered (loaded from cache).
 
-### 🗄️ Custom Storage Backends
+## 🔬 Customize Argument Hashing
+
+You can customize how individual function arguments are hashed without changing their actual values when passed in.
+
+* **`Annotated[T, HashBy[fn]]`**:\
+    Hashes the argument by applying `fn(argument)` before hashing. This enables custom normalization (e.g., sorting lists to ignore order) or optimized hashing for complex types, improving cache hit rates or speeding up hashing.
+
+* **`NoHash[T]`**:\
+    Completely excludes the argument from hashing, so changes to it won’t trigger cache invalidation.
+
+**Example:**
+
+```python
+from typing import Annotated
+from checkpointer import checkpoint, HashBy, NoHash
+from pathlib import Path
+import logging
+
+def file_bytes(path: Path) -> bytes:
+    return path.read_bytes()
+
+@checkpoint
+def process(
+    numbers: Annotated[list[int], HashBy[sorted]],   # Hash by sorted list
+    data_file: Annotated[Path, HashBy[file_bytes]],  # Hash by file content
+    log: NoHash[logging.Logger],                     # Exclude logger from hashing
+):
+    ...
+```
+
+In this example, the cache key for `numbers` ignores order, `data_file` is hashed based on its contents rather than path, and changes to `log` don’t affect caching.
+
+## 🗄️ Custom Storage Backends
 
 For integration with databases, cloud storage, or custom serialization, implement your own storage backend by inheriting from `checkpointer.Storage` and implementing its abstract methods.
 
-Within custom storage methods, `call_id` identifies calls by arguments. Use `self.fn_id()` to get the function's unique identity (name + hash/version), crucial for organizing stored checkpoints (e.g., by function version). Access global `Checkpointer` config via `self.checkpointer`.
+Within custom storage methods, `call_hash` identifies calls by arguments. Use `self.fn_id()` to get the function's unique identity (name + hash/version), crucial for organizing stored checkpoints (e.g., by function version). Access global `Checkpointer` config via `self.checkpointer`.
 
-#### Example: Custom Storage Backend
+**Example: Custom Storage Backend**
 
 ```python
 from checkpointer import checkpoint, Storage
 from datetime import datetime
 
 class MyCustomStorage(Storage):
-    def exists(self, call_id):
+    def exists(self, call_hash):
         # Example: Constructing a path based on function ID and call ID
         fn_dir = self.checkpointer.root_path / self.fn_id()
-        return (fn_dir / call_id).exists()
+        return (fn_dir / call_hash).exists()
 
-    def store(self, call_id, data):
-        ... # Store the serialized data for `call_id`
+    def store(self, call_hash, data):
+        ... # Store the serialized data for `call_hash`
         return data # This method must return the data back to checkpointer
 
-    def checkpoint_date(self, call_id): ...
-    def load(self, call_id): ...
-    def delete(self, call_id): ...
+    def checkpoint_date(self, call_hash): ...
+    def load(self, call_hash): ...
+    def delete(self, call_hash): ...
 
 @checkpoint(format=MyCustomStorage)
 def custom_cached_function(x: int):

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/__init__.py

@@ -4,18 +4,18 @@ from typing import Callable
 from .checkpoint import CachedFunction, Checkpointer, CheckpointError
 from .object_hash import ObjectHash
 from .storages import MemoryStorage, PickleStorage, Storage
-from .utils import AwaitableValue
+from .types import AwaitableValue, HashBy, NoHash
 
 checkpoint = Checkpointer()
 capture_checkpoint = Checkpointer(capture=True)
 memory_checkpoint = Checkpointer(format="memory", verbosity=0)
 tmp_checkpoint = Checkpointer(root_path=f"{tempfile.gettempdir()}/checkpoints")
-static_checkpoint = Checkpointer(fn_hash=ObjectHash())
+static_checkpoint = Checkpointer(fn_hash_from=())
 
 def cleanup_all(invalidated=True, expired=True):
   for obj in gc.get_objects():
     if isinstance(obj, CachedFunction):
       obj.cleanup(invalidated=invalidated, expired=expired)
 
-def get_function_hash(fn: Callable, capture=False) -> str:
-  return CachedFunction(Checkpointer(capture=capture), fn).fn_hash
+def get_function_hash(fn: Callable) -> str:
+  return CachedFunction(Checkpointer(), fn).ident.fn_hash

checkpointer-2.11.1/checkpointer/checkpoint.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+import re
+from datetime import datetime
+from functools import cached_property, update_wrapper
+from inspect import Parameter, Signature, iscoroutine, signature
+from pathlib import Path
+from typing import (
+  Annotated, Callable, Concatenate, Coroutine, Generic,
+  Iterable, Literal, ParamSpec, Self, Type, TypedDict,
+  TypeVar, Unpack, cast, get_args, get_origin, overload,
+)
+from .fn_ident import RawFunctionIdent, get_fn_ident
+from .object_hash import ObjectHash
+from .print_checkpoint import print_checkpoint
+from .storages import STORAGE_MAP, Storage
+from .types import AwaitableValue, HashBy
+from .utils import unwrap_fn
+
+Fn = TypeVar("Fn", bound=Callable)
+P = ParamSpec("P")
+R = TypeVar("R")
+C = TypeVar("C")
+
+DEFAULT_DIR = Path.home() / ".cache/checkpoints"
+
+class CheckpointError(Exception):
+  pass
+
+class CheckpointerOpts(TypedDict, total=False):
+  format: Type[Storage] | Literal["pickle", "memory", "bcolz"]
+  root_path: Path | str | None
+  when: bool
+  verbosity: Literal[0, 1, 2]
+  should_expire: Callable[[datetime], bool] | None
+  capture: bool
+  fn_hash_from: object
+
+class Checkpointer:
+  def __init__(self, **opts: Unpack[CheckpointerOpts]):
+    self.format = opts.get("format", "pickle")
+    self.root_path = Path(opts.get("root_path", DEFAULT_DIR) or ".")
+    self.when = opts.get("when", True)
+    self.verbosity = opts.get("verbosity", 1)
+    self.should_expire = opts.get("should_expire")
+    self.capture = opts.get("capture", False)
+    self.fn_hash_from = opts.get("fn_hash_from")
+
+  @overload
+  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 | CachedFunction[Fn]:
+    if override_opts:
+      opts = CheckpointerOpts(**{**self.__dict__, **override_opts})
+      return Checkpointer(**opts)(fn)
+
+    return CachedFunction(self, fn) if callable(fn) else self
+
+class FunctionIdent:
+  """
+  Represents the identity and hash state of a cached function.
+  Separated from CachedFunction to prevent hash desynchronization
+  among bound instances when `.reinit()` is called.
+  """
+  def __init__(self, cached_fn: CachedFunction):
+    self.__dict__.clear()
+    self.cached_fn = cached_fn
+
+  @cached_property
+  def raw_ident(self) -> RawFunctionIdent:
+    return get_fn_ident(unwrap_fn(self.cached_fn.fn), self.cached_fn.checkpointer.capture)
+
+  @cached_property
+  def fn_hash(self) -> str:
+    if self.cached_fn.checkpointer.fn_hash_from is not None:
+      return str(ObjectHash(self.cached_fn.checkpointer.fn_hash_from, digest_size=16))
+    deep_hashes = [depend.ident.raw_ident.fn_hash for depend in self.cached_fn.deep_depends()]
+    return str(ObjectHash(digest_size=16).write_text(iter=deep_hashes))
+
+  @cached_property
+  def captured_hash(self) -> str:
+    deep_hashes = [depend.ident.raw_ident.captured_hash for depend in self.cached_fn.deep_depends()]
+    return str(ObjectHash().write_text(iter=deep_hashes))
+
+  def clear(self):
+    self.__init__(self.cached_fn)
+
+class CachedFunction(Generic[Fn]):
+  def __init__(self, checkpointer: Checkpointer, fn: 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)
+    self.checkpointer = checkpointer
+    self.fn = fn
+    self.fn_dir = f"{fn_file}/{fn_name}"
+    self.storage = Storage(self)
+    self.cleanup = self.storage.cleanup
+    self.bound = ()
+
+    sig = signature(wrapped)
+    params = list(sig.parameters.items())
+    pos_params = (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD)
+    self.arg_names = [name for name, param in params if param.kind in pos_params]
+    self.default_args = {name: param.default for name, param in params if param.default is not Parameter.empty}
+    self.hash_by_map = get_hash_by_map(sig)
+    self.ident = FunctionIdent(self)
+
+  @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
+
+  @property
+  def depends(self) -> list[Callable]:
+    return self.ident.raw_ident.depends
+
+  def reinit(self, recursive=False) -> CachedFunction[Fn]:
+    depend_idents = [depend.ident for depend in self.deep_depends()] if recursive else [self.ident]
+    for ident in depend_idents: ident.clear()
+    for ident in depend_idents: ident.fn_hash
+    return self
+
+  def get_call_hash(self, args: tuple, kw: dict[str, object]) -> str:
+    args = self.bound + args
+    pos_args = args[len(self.arg_names):]
+    named_pos_args = dict(zip(self.arg_names, args))
+    named_args = {**self.default_args, **named_pos_args, **kw}
+    if hash_by_map := self.hash_by_map:
+      rest_hash_by = hash_by_map.get(b"**")
+      for key, value in named_args.items():
+        if hash_by := hash_by_map.get(key, rest_hash_by):
+          named_args[key] = hash_by(value)
+      if pos_hash_by := hash_by_map.get(b"*"):
+        pos_args = tuple(map(pos_hash_by, pos_args))
+    return str(ObjectHash(named_args, pos_args, self.ident.captured_hash, digest_size=16))
+
+  async def _resolve_coroutine(self, call_hash: str, coroutine: Coroutine):
+    return self.storage.store(call_hash, AwaitableValue(await coroutine)).value
+
+  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(*full_args, **kw)
+
+    call_hash = self.get_call_hash(args, kw)
+    call_hash_long = f"{self.fn_dir}/{self.ident.fn_hash}/{call_hash}"
+
+    refresh = rerun \
+      or not self.storage.exists(call_hash) \
+      or (params.should_expire and params.should_expire(self.storage.checkpoint_date(call_hash)))
+
+    if refresh:
+      print_checkpoint(params.verbosity >= 1, "MEMORIZING", call_hash_long, "blue")
+      data = self.fn(*full_args, **kw)
+      if iscoroutine(data):
+        return self._resolve_coroutine(call_hash, data)
+      return self.storage.store(call_hash, data)
+
+    try:
+      data = self.storage.load(call_hash)
+      print_checkpoint(params.verbosity >= 2, "REMEMBERED", call_hash_long, "green")
+      return data
+    except (EOFError, FileNotFoundError):
+      pass
+    print_checkpoint(params.verbosity >= 1, "CORRUPTED", call_hash_long, "yellow")
+    return 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)
+
+  def exists(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> bool:
+    return self.storage.exists(self.get_call_hash(args, kw))
+
+  def delete(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs):
+    self.storage.delete(self.get_call_hash(args, kw))
+
+  @overload
+  def get(self: Callable[P, Coroutine[object, object, 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_hash = self.get_call_hash(args, kw)
+    try:
+      data = self.storage.load(call_hash)
+      return data.value if isinstance(data, AwaitableValue) else data
+    except Exception as ex:
+      raise CheckpointError("Could not load checkpoint") from ex
+
+  @overload
+  def _set(self: Callable[P, Coroutine[object, object, R]], value: AwaitableValue[R], *args: P.args, **kw: P.kwargs): ...
+  @overload
+  def _set(self: Callable[P, R], value: R, *args: P.args, **kw: P.kwargs): ...
+  def _set(self, value, *args, **kw):
+    self.storage.store(self.get_call_hash(args, kw), value)
+
+  def __repr__(self) -> str:
+    return f"<CachedFunction {self.fn.__name__} {self.ident.fn_hash[:6]}>"
+
+  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, CachedFunction):
+          yield from depend.deep_depends(visited)
+
+def hash_by_from_annotation(annotation: type) -> Callable[[object], object] | None:
+  if get_origin(annotation) is Annotated:
+    args = get_args(annotation)
+    metadata = args[1] if len(args) > 1 else None
+    if get_origin(metadata) is HashBy:
+      return get_args(metadata)[0]
+
+def get_hash_by_map(sig: Signature) -> dict[str | bytes, Callable[[object], object]]:
+  hash_by_map = {}
+  for name, param in sig.parameters.items():
+    if param.kind == Parameter.VAR_POSITIONAL:
+      name = b"*"
+    elif param.kind == Parameter.VAR_KEYWORD:
+      name = b"**"
+    hash_by_map[name] = hash_by_from_annotation(param.annotation)
+  return hash_by_map if any(hash_by_map.values()) else {}

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/fn_ident.py

@@ -1,15 +1,19 @@
 import dis
 import inspect
-from collections.abc import Callable
 from itertools import takewhile
 from pathlib import Path
 from types import CodeType, FunctionType, MethodType
-from typing import Any, Iterable, Type, TypeGuard
+from typing import Callable, Iterable, NamedTuple, Type, TypeGuard
 from .object_hash import ObjectHash
-from .utils import AttrDict, distinct, get_cell_contents, iterate_and_upcoming, transpose, unwrap_fn
+from .utils import AttrDict, distinct, get_cell_contents, iterate_and_upcoming, unwrap_fn
 
 cwd = Path.cwd().resolve()
 
+class RawFunctionIdent(NamedTuple):
+  fn_hash: str
+  captured_hash: str
+  depends: list[Callable]
+
 def is_class(obj) -> TypeGuard[Type]:
   # isinstance works too, but needlessly triggers _lazyinit()
   return issubclass(type(obj), type)
@@ -33,7 +37,7 @@ def extract_classvars(code: CodeType, scope_vars: AttrDict) -> dict[str, dict[st
       scope_obj = None
   return classvars
 
-def extract_scope_values(code: CodeType, scope_vars: AttrDict) -> Iterable[tuple[tuple[str, ...], Any]]:
+def extract_scope_values(code: CodeType, scope_vars: AttrDict) -> Iterable[tuple[tuple[str, ...], object]]:
   classvars = extract_classvars(code, scope_vars)
   scope_vars = scope_vars.set({k: scope_vars[k].set(v) for k, v in classvars.items()})
   for instr, upcoming_instrs in iterate_and_upcoming(dis.get_instructions(code)):
@@ -55,7 +59,7 @@ def get_self_value(fn: Callable) -> type | object | None:
   if is_class(cls):
     return cls
 
-def get_fn_captured_vals(fn: Callable) -> list[Any]:
+def get_fn_captured_vals(fn: Callable) -> list[object]:
   self_value = get_self_value(fn)
   scope_vars = AttrDict({
     "LOAD_FAST": AttrDict({"self": self_value} if self_value else {}),
@@ -71,24 +75,29 @@ def is_user_fn(candidate_fn) -> TypeGuard[Callable]:
   fn_path = Path(inspect.getfile(candidate_fn)).resolve()
   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]]:
+def get_depend_fns(fn: Callable, captured_vals_by_fn: dict[Callable, list[object]] = {}) -> dict[Callable, list[object]]:
   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, cached_fn=True) for val in captured_vals if callable(val))
-  for child_fn in child_fns:
+  captured_vals_by_fn = captured_vals_by_fn or {}
+  captured_vals_by_fn[fn] = [val for val in captured_vals if not callable(val)]
+  for val in captured_vals:
+    if not callable(val):
+      continue
+    child_fn = unwrap_fn(val, cached_fn=True)
     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)
+      get_depend_fns(child_fn, captured_vals_by_fn)
   return captured_vals_by_fn
 
-def get_fn_ident(fn: Callable, capture: bool) -> tuple[str, list[Callable]]:
+def get_fn_ident(fn: Callable, capture: bool) -> RawFunctionIdent:
   from .checkpoint import CachedFunction
-  captured_vals_by_fn = get_depend_fns(fn, capture)
-  depends, depend_captured_vals = transpose(captured_vals_by_fn.items(), 2)
+  captured_vals_by_fn = get_depend_fns(fn)
+  depend_captured_vals = list(captured_vals_by_fn.values()) * capture
+  depends = captured_vals_by_fn.keys()
   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, CachedFunction)]
-  fn_hash = str(ObjectHash(fn, unwrapped_depends).update(depend_captured_vals, tolerate_errors=True))
-  return fn_hash, depends
+  assert fn == unwrapped_depends[0]
+  fn_hash = str(ObjectHash(iter=unwrapped_depends))
+  captured_hash = str(ObjectHash(iter=depend_captured_vals, tolerate_errors=True))
+  return RawFunctionIdent(fn_hash, captured_hash, depends)

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/object_hash.py

@@ -1,16 +1,19 @@
 import ctypes
 import hashlib
+import inspect
 import io
 import re
 import sys
+import tokenize
 from collections.abc import Iterable
 from contextlib import nullcontext, suppress
 from decimal import Decimal
+from io import StringIO
 from itertools import chain
-from pickle import HIGHEST_PROTOCOL as PROTOCOL
+from pickle import HIGHEST_PROTOCOL as PICKLE_PROTOCOL
 from types import BuiltinFunctionType, FunctionType, GeneratorType, MethodType, ModuleType, UnionType
-from typing import Any, TypeVar
-from .utils import ContextVar, get_fn_body
+from typing import Callable, TypeVar
+from .utils import ContextVar
 
 np, torch = None, None
 
@@ -31,16 +34,16 @@ else:
 def encode_type(t: type | FunctionType) -> str:
   return f"{t.__module__}:{t.__qualname__}"
 
-def encode_type_of(v: Any) -> str:
+def encode_type_of(v: object) -> str:
   return encode_type(type(v))
 
 class ObjectHashError(Exception):
-  def __init__(self, obj: Any, cause: Exception):
+  def __init__(self, obj: object, cause: Exception):
     super().__init__(f"{type(cause).__name__} error when hashing {obj}")
     self.obj = obj
 
 class ObjectHash:
-  def __init__(self, *objs: Any, iter: Iterable[Any] = (), digest_size=64, tolerate_errors=False) -> None:
+  def __init__(self, *objs: object, iter: Iterable[object] = (), digest_size=64, tolerate_errors=False) -> None:
     self.hash = hashlib.blake2b(digest_size=digest_size)
     self.current: dict[int, int] = {}
     self.tolerate_errors = ContextVar(tolerate_errors)
@@ -59,7 +62,7 @@ class ObjectHash:
   def __eq__(self, value: object) -> bool:
     return isinstance(value, ObjectHash) and str(self) == str(value)
 
-  def nested_hash(self, *objs: Any) -> str:
+  def nested_hash(self, *objs: object) -> str:
     return ObjectHash(iter=objs, tolerate_errors=self.tolerate_errors.value).hexdigest()
 
   def write_bytes(self, *data: bytes, iter: Iterable[bytes] = ()) -> "ObjectHash":
@@ -70,10 +73,10 @@ class ObjectHash:
   def write_text(self, *data: str, iter: Iterable[str] = ()) -> "ObjectHash":
     return self.write_bytes(iter=(d.encode() for d in chain(data, iter)))
 
-  def header(self, *args: Any) -> "ObjectHash":
+  def header(self, *args: object) -> "ObjectHash":
     return self.write_bytes(":".join(map(str, args)).encode())
 
-  def update(self, *objs: Any, iter: Iterable[Any] = (), tolerate_errors: bool | None=None) -> "ObjectHash":
+  def update(self, *objs: object, iter: Iterable[object] = (), tolerate_errors: bool | None=None) -> "ObjectHash":
     with nullcontext() if tolerate_errors is None else self.tolerate_errors.set(tolerate_errors):
       for obj in chain(objs, iter):
         try:
@@ -81,11 +84,11 @@ class ObjectHash:
         except Exception as ex:
           if self.tolerate_errors.value:
             self.header("error").update(type(ex))
-            continue
-          raise ObjectHashError(obj, ex) from ex
+          else:
+            raise ObjectHashError(obj, ex) from ex
       return self
 
-  def _update_one(self, obj: Any) -> None:
+  def _update_one(self, obj: object) -> None:
     match obj:
       case None:
         self.header("null")
@@ -142,7 +145,7 @@ class ObjectHash:
       case _ if np and isinstance(obj, np.ndarray):
         self.header("ndarray", encode_type_of(obj), obj.shape, obj.strides).update(obj.dtype)
         if obj.dtype.hasobject:
-          self.update(obj.__reduce_ex__(PROTOCOL))
+          self.update(obj.__reduce_ex__(PICKLE_PROTOCOL))
         else:
           array = np.ascontiguousarray(obj if obj.base is None else obj.base).view(np.uint8)
           self.write_bytes(array.data)
@@ -180,13 +183,14 @@ 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: Any) -> "ObjectHash":
+  def _update_object(self, obj: object) -> "ObjectHash":
     self.header("instance", encode_type_of(obj))
-    if hasattr(obj, "__objecthash__") and callable(obj.__objecthash__):
-      return self.header("objecthash").update(obj.__objecthash__())
+    get_hash = hasattr(obj, "__objecthash__") and getattr(obj, "__objecthash__")
+    if callable(get_hash):
+      return self.header("objecthash").update(get_hash())
     reduced = None
     with suppress(Exception):
-      reduced = obj.__reduce_ex__(PROTOCOL)
+      reduced = obj.__reduce_ex__(PICKLE_PROTOCOL)
     with suppress(Exception):
       reduced = reduced or obj.__reduce__()
     if isinstance(reduced, str):
@@ -206,3 +210,12 @@ class ObjectHash:
       return self._update_iterator(obj)
     repr_str = re.sub(r"\s+(at\s+0x[0-9a-fA-F]+)(>)$", r"\2", repr(obj))
     return self.header("repr").update(repr_str)
+
+def get_fn_body(fn: Callable) -> str:
+  try:
+    source = inspect.getsource(fn)
+  except OSError:
+    return ""
+  tokens = tokenize.generate_tokens(StringIO(source).readline)
+  ignore_types = (tokenize.COMMENT, tokenize.NL)
+  return "".join("\0" + token.string for token in tokens if token.type not in ignore_types)

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/storages/memory_storage.py

@@ -9,21 +9,21 @@ class MemoryStorage(Storage):
   def get_dict(self):
     return item_map.setdefault(self.fn_dir(), {})
 
-  def store(self, call_id, data):
-    self.get_dict()[call_id] = (datetime.now(), data)
+  def store(self, call_hash, data):
+    self.get_dict()[call_hash] = (datetime.now(), data)
     return data
 
-  def exists(self, call_id):
-    return call_id in self.get_dict()
+  def exists(self, call_hash):
+    return call_hash in self.get_dict()
 
-  def checkpoint_date(self, call_id):
-    return self.get_dict()[call_id][0]
+  def checkpoint_date(self, call_hash):
+    return self.get_dict()[call_hash][0]
 
-  def load(self, call_id):
-    return self.get_dict()[call_id][1]
+  def load(self, call_hash):
+    return self.get_dict()[call_hash][1]
 
-  def delete(self, call_id):
-    self.get_dict().pop(call_id, None)
+  def delete(self, call_hash):
+    self.get_dict().pop(call_hash, None)
 
   def cleanup(self, invalidated=True, expired=True):
     curr_key = self.fn_dir()
@@ -32,6 +32,6 @@ class MemoryStorage(Storage):
         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()):
+          for call_hash, (date, _) in list(calldict.items()):
             if self.checkpointer.should_expire(date):
-              del calldict[call_id]
+              del calldict[call_hash]

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/storages/pickle_storage.py

@@ -8,29 +8,29 @@ 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"
+  def get_path(self, call_hash: str):
+    return self.fn_dir() / f"{call_hash}.pkl"
 
-  def store(self, call_id, data):
-    path = self.get_path(call_id)
+  def store(self, call_hash, data):
+    path = self.get_path(call_hash)
     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 exists(self, call_hash):
+    return self.get_path(call_hash).exists()
 
-  def checkpoint_date(self, call_id):
+  def checkpoint_date(self, call_hash):
     # Should use st_atime/access time?
-    return filedate(self.get_path(call_id))
+    return filedate(self.get_path(call_hash))
 
-  def load(self, call_id):
-    with self.get_path(call_id).open("rb") as file:
+  def load(self, call_hash):
+    with self.get_path(call_hash).open("rb") as file:
       return pickle.load(file)
 
-  def delete(self, call_id):
-    self.get_path(call_id).unlink(missing_ok=True)
+  def delete(self, call_hash):
+    self.get_path(call_hash).unlink(missing_ok=True)
 
   def cleanup(self, invalidated=True, expired=True):
     version_path = self.fn_dir()

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/storages/storage.py

@@ -15,19 +15,19 @@ class Storage:
     self.cached_fn = cached_fn
 
   def fn_id(self) -> str:
-    return f"{self.cached_fn.fn_dir}/{self.cached_fn.fn_hash}"
+    return f"{self.cached_fn.fn_dir}/{self.cached_fn.ident.fn_hash}"
 
   def fn_dir(self) -> Path:
     return self.checkpointer.root_path / self.fn_id()
 
-  def store(self, call_id: str, data: Any) -> Any: ...
+  def store(self, call_hash: str, data: Any) -> Any: ...
 
-  def exists(self, call_id: str) -> bool: ...
+  def exists(self, call_hash: str) -> bool: ...
 
-  def checkpoint_date(self, call_id: str) -> datetime: ...
+  def checkpoint_date(self, call_hash: str) -> datetime: ...
 
-  def load(self, call_id: str) -> Any: ...
+  def load(self, call_hash: str) -> Any: ...
 
-  def delete(self, call_id: str) -> None: ...
+  def delete(self, call_hash: str) -> None: ...
 
   def cleanup(self, invalidated=True, expired=True): ...

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/test_checkpointer.py

@@ -1,6 +1,6 @@
 import asyncio
 import pytest
-from . import CheckpointError, checkpoint
+from checkpointer import CheckpointError, checkpoint
 from .utils import AttrDict
 
 def global_multiply(a: int, b: int) -> int:
@@ -110,16 +110,16 @@ def test_capture():
   def test_a():
     return item_dict.a + 1
 
-  init_hash_a = test_a.fn_hash
-  init_hash_whole = test_whole.fn_hash
+  init_hash_a = test_a.ident.captured_hash
+  init_hash_whole = test_whole.ident.captured_hash
   item_dict.b += 1
   test_whole.reinit()
   test_a.reinit()
-  assert test_whole.fn_hash != init_hash_whole
-  assert test_a.fn_hash == init_hash_a
+  assert test_whole.ident.captured_hash != init_hash_whole
+  assert test_a.ident.captured_hash == init_hash_a
   item_dict.a += 1
   test_a.reinit()
-  assert test_a.fn_hash != init_hash_a
+  assert test_a.ident.captured_hash != init_hash_a
 
 def test_depends():
   def multiply_wrapper(a: int, b: int) -> int:

checkpointer-2.11.1/checkpointer/types.py

@@ -0,0 +1,17 @@
+from typing import Annotated, Callable, Generic, TypeVar
+
+T = TypeVar("T")
+Fn = TypeVar("Fn", bound=Callable)
+
+class HashBy(Generic[Fn]):
+  pass
+
+NoHash = Annotated[T, HashBy[lambda _: None]]
+
+class AwaitableValue(Generic[T]):
+  def __init__(self, value: T):
+    self.value = value
+
+  def __await__(self):
+    yield
+    return self.value

{checkpointer-2.10.1 → checkpointer-2.11.1}/checkpointer/utils.py

@@ -1,7 +1,4 @@
-import inspect
-import tokenize
 from contextlib import contextmanager
-from io import StringIO
 from typing import Any, Callable, Generic, Iterable, TypeVar, cast
 
 T = TypeVar("T")
@@ -10,21 +7,6 @@ Fn = TypeVar("Fn", bound=Callable)
 def distinct(seq: Iterable[T]) -> list[T]:
   return list(dict.fromkeys(seq))
 
-def transpose(tuples, default_num_returns=0):
-  output = tuple(zip(*tuples))
-  if not output:
-    return ([],) * default_num_returns
-  return tuple(map(list, output))
-
-def get_fn_body(fn: Callable) -> str:
-  try:
-    source = inspect.getsource(fn)
-  except OSError:
-    return ""
-  tokens = tokenize.generate_tokens(StringIO(source).readline)
-  ignore_types = (tokenize.COMMENT, tokenize.NL)
-  return "".join("\0" + token.string for token in tokens if token.type not in ignore_types)
-
 def get_cell_contents(fn: Callable) -> Iterable[tuple[str, Any]]:
   for key, cell in zip(fn.__code__.co_freevars, fn.__closure__ or []):
     try:
@@ -39,14 +21,6 @@ def unwrap_fn(fn: Fn, cached_fn=False) -> Fn:
       return cast(Fn, fn)
     fn = getattr(fn, "__wrapped__")
 
-class AwaitableValue:
-  def __init__(self, value):
-    self.value = value
-
-  def __await__(self):
-    yield
-    return self.value
-
 class AttrDict(dict):
   def __init__(self, *args, **kwargs):
     super().__init__(*args, **kwargs)
@@ -91,14 +65,19 @@ class iterate_and_upcoming(Generic[T]):
   def __init__(self, it: Iterable[T]) -> None:
     self.it = iter(it)
     self.previous: tuple[()] | tuple[T] = ()
+    self.tracked = self._tracked_iter()
 
   def __iter__(self):
     return self
 
   def __next__(self) -> tuple[T, Iterable[T]]:
-    item = self.previous[0] if self.previous else next(self.it)
-    self.previous = ()
-    return item, self._tracked_iter()
+    try:
+      item = self.previous[0] if self.previous else next(self.it)
+      self.previous = ()
+      return item, self.tracked
+    except StopIteration:
+      self.tracked.close()
+      raise
 
   def _tracked_iter(self):
     for x in self.it:

{checkpointer-2.10.1 → checkpointer-2.11.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.10.1"
+version = "2.11.1"
 requires-python = ">=3.11"
 dependencies = []
 authors = [
@@ -21,10 +21,11 @@ Repository = "https://github.com/Reddan/checkpointer.git"
 [dependency-groups]
 dev = [
   "numpy>=2.2.1",
-  "omg>=1.3.4",
+  "omg>=1.3.6",
   "poethepoet>=0.30.0",
   "pytest>=8.3.5",
   "pytest-asyncio>=0.26.0",
+  "rich>=14.0.0",
   "torch>=2.5.1",
 ]
 

{checkpointer-2.10.1 → checkpointer-2.11.1}/uv.lock

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.10.1"
+version = "2.11.1"
 source = { editable = "." }
 
 [package.dev-dependencies]
@@ -18,6 +18,7 @@ dev = [
     { name = "poethepoet" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
+    { name = "rich" },
     { name = "torch" },
 ]
 
@@ -26,10 +27,11 @@ dev = [
 [package.metadata.requires-dev]
 dev = [
     { name = "numpy", specifier = ">=2.2.1" },
-    { name = "omg", specifier = ">=1.3.4" },
+    { name = "omg", specifier = ">=1.3.6" },
     { name = "poethepoet", specifier = ">=0.30.0" },
     { name = "pytest", specifier = ">=8.3.5" },
     { name = "pytest-asyncio", specifier = ">=0.26.0" },
+    { name = "rich", specifier = ">=14.0.0" },
     { name = "torch", specifier = ">=2.5.1" },
 ]
 
@@ -81,6 +83,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bd/0f/2ba5fbcd631e3e88689309dbe978c5769e883e4b84ebfe7da30b43275c5a/jinja2-3.1.5-py3-none-any.whl", hash = "sha256:aba0f4dc9ed8013c424088f68a5c226f7d6097ed89b246d7749c2ec4175c6adb", size = 134596 },
 ]
 
+[[package]]
+name = "markdown-it-py"
+version = "3.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mdurl" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/42/d7/1ec15b46af6af88f19b8e5ffea08fa375d433c998b8a7639e76935c14f1f/markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", size = 87528 },
+]
+
 [[package]]
 name = "markupsafe"
 version = "3.0.2"
@@ -129,6 +143,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/4f/65/6079a46068dfceaeabb5dcad6d674f5f5c61a6fa5673746f42a9f4c233b3/MarkupSafe-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:e444a31f8db13eb18ada366ab3cf45fd4b31e4db1236a4448f68778c1d1a5a2f", size = 15739 },
 ]
 
+[[package]]
+name = "mdurl"
+version = "0.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
+]
+
 [[package]]
 name = "mpmath"
 version = "1.3.0"
@@ -307,14 +330,14 @@ wheels = [
 
 [[package]]
 name = "omg"
-version = "1.3.4"
+version = "1.3.6"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "watchdog" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/0b/12/9037613ad41ad71f6342e18a7fdf8f19826dfe61555c45541d87045cd084/omg-1.3.4.tar.gz", hash = "sha256:31e5fd82c2d2fdc53e944b1cb9fb7f630919f9441b8822fb5ae108962fabb49b", size = 14193 }
+sdist = { url = "https://files.pythonhosted.org/packages/65/06/da0a3778b7ff8f1333ed7ddc0931ffff3c86ab5cb8bc4a96a1d0edb8671b/omg-1.3.6.tar.gz", hash = "sha256:465a51b7576fa31ef313e2b9a77d57f5d4816fb0a14dca0fc5c09ff471074fe6", size = 14268 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/fd/88/4a4e1cc054f141744a60e17d76c8ea1aea7af7e125cae0713d9b0d5ec12f/omg-1.3.4-py3-none-any.whl", hash = "sha256:449e00d341b63afa23633f47cd751a1874ad2545162e5f9abbdd115525fd7a71", size = 7926 },
+    { url = "https://files.pythonhosted.org/packages/dd/d2/87346e94dbecd3a65a09e2156c1adf30c162f31e69d0936343c3eff53e7a/omg-1.3.6-py3-none-any.whl", hash = "sha256:8e3ac99a18d5284ceef2ed98492d288d5f22ee2bb417591654a7d2433e196607", size = 7988 },
 ]
 
 [[package]]
@@ -357,6 +380,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/06/e1/04f56c9d848d6135ca3328c5a2ca84d3303c358ad7828db290385e36a8cc/poethepoet-0.31.1-py3-none-any.whl", hash = "sha256:7fdfa0ac6074be9936723e7231b5bfaad2923e96c674a9857e81d326cf8ccdc2", size = 80238 },
 ]
 
+[[package]]
+name = "pygments"
+version = "2.19.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7c/2d/c3338d48ea6cc0feb8446d8e6937e1408088a72a39937982cc6111d17f84/pygments-2.19.1.tar.gz", hash = "sha256:61c16d2a8576dc0649d9f39e089b5f02bcd27fba10d8fb4dcc28173f7a45151f", size = 4968581 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8a/0b/9fcc47d19c48b59121088dd6da2488a49d5f72dacf8262e2790a1d2c7d15/pygments-2.19.1-py3-none-any.whl", hash = "sha256:9ea1544ad55cecf4b8242fab6dd35a93bbce657034b0611ee383099054ab6d8c", size = 1225293 },
+]
+
 [[package]]
 name = "pytest"
 version = "8.3.5"
@@ -419,6 +451,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446 },
 ]
 
+[[package]]
+name = "rich"
+version = "14.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markdown-it-py" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a1/53/830aa4c3066a8ab0ae9a9955976fb770fe9c6102117c8ec4ab3ea62d89e8/rich-14.0.0.tar.gz", hash = "sha256:82f1bc23a6a21ebca4ae0c45af9bdbc492ed20231dcb63f297d6d1021a9d5725", size = 224078 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0d/9b/63f4c7ebc259242c89b3acafdb37b41d1185c07ff0011164674e9076b491/rich-14.0.0-py3-none-any.whl", hash = "sha256:1c9491e1951aac09caffd42f448ee3d04e58923ffe14993f6e83068dc395d7e0", size = 243229 },
+]
+
 [[package]]
 name = "setuptools"
 version = "75.6.0"

checkpointer-2.10.1/checkpointer/checkpoint.py

@@ -1,184 +0,0 @@
-from __future__ import annotations
-import inspect
-import re
-from datetime import datetime
-from functools import cached_property, update_wrapper
-from pathlib import Path
-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
-from .storages import STORAGE_MAP, Storage
-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"
-
-class CheckpointError(Exception):
-  pass
-
-class CheckpointerOpts(TypedDict, total=False):
-  format: Type[Storage] | Literal["pickle", "memory", "bcolz"]
-  root_path: Path | str | None
-  when: bool
-  verbosity: Literal[0, 1, 2]
-  hash_by: Callable | None
-  should_expire: Callable[[datetime], bool] | None
-  capture: bool
-  fn_hash: ObjectHash | None
-
-class Checkpointer:
-  def __init__(self, **opts: Unpack[CheckpointerOpts]):
-    self.format = opts.get("format", "pickle")
-    self.root_path = Path(opts.get("root_path", DEFAULT_DIR) or ".")
-    self.when = opts.get("when", True)
-    self.verbosity = opts.get("verbosity", 1)
-    self.hash_by = opts.get("hash_by")
-    self.should_expire = opts.get("should_expire")
-    self.capture = opts.get("capture", False)
-    self.fn_hash = opts.get("fn_hash")
-
-  @overload
-  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 | CachedFunction[Fn]:
-    if override_opts:
-      opts = CheckpointerOpts(**{**self.__dict__, **override_opts})
-      return Checkpointer(**opts)(fn)
-
-    return CachedFunction(self, fn) if callable(fn) else self
-
-class CachedFunction(Generic[Fn]):
-  def __init__(self, checkpointer: Checkpointer, fn: 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)
-    self.checkpointer = checkpointer
-    self.fn = 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]]:
-    return get_fn_ident(unwrap_fn(self.fn), self.checkpointer.capture)
-
-  @property
-  def fn_hash_raw(self) -> str:
-    return self.ident_tuple[0]
-
-  @property
-  def depends(self) -> list[Callable]:
-    return self.ident_tuple[1]
-
-  @cached_property
-  def fn_hash(self) -> str:
-    deep_hashes = [depend.fn_hash_raw for depend in self.deep_depends()]
-    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) -> CachedFunction[Fn]:
-    depends = list(self.deep_depends()) if recursive else [self]
-    for depend in depends:
-      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, call_id: str, awaitable: Awaitable):
-    return self.storage.store(call_id, AwaitableValue(await awaitable)).value
-
-  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(*full_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(call_id) \
-      or (params.should_expire and params.should_expire(self.storage.checkpoint_date(call_id)))
-
-    if refresh:
-      print_checkpoint(params.verbosity >= 1, "MEMORIZING", call_id_long, "blue")
-      data = self.fn(*full_args, **kw)
-      if inspect.isawaitable(data):
-        return self._resolve_awaitable(call_id, data)
-      return self.storage.store(call_id, data)
-
-    try:
-      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", call_id_long, "yellow")
-    return 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:
-      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: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> bool:
-    return self.storage.exists(self.get_call_id(args, kw))
-
-  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"<CachedFunction {self.fn.__name__} {self.fn_hash[:6]}>"
-
-  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, CachedFunction):
-          yield from depend.deep_depends(visited)