checkpointer 2.10.1__py3-none-any.whl → 2.11.0__py3-none-any.whl

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/checkpoint.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
-import inspect
 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 (
-  Awaitable, Callable, Concatenate, Generic, Iterable, Literal,
-  ParamSpec, Self, Type, TypedDict, TypeVar, Unpack, cast, overload,
+  Annotated, Callable, Concatenate, Coroutine, Generic,
+  Iterable, Literal, ParamSpec, Self, Type, TypedDict,
+  TypeVar, Unpack, cast, get_args, get_origin, overload,
 )
-from .fn_ident import get_fn_ident
+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 .utils import AwaitableValue, unwrap_fn
+from .types import AwaitableValue, HashBy
+from .utils import unwrap_fn
 
 Fn = TypeVar("Fn", bound=Callable)
 P = ParamSpec("P")
@@ -29,10 +31,9 @@ class CheckpointerOpts(TypedDict, total=False):
   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
+  fn_hash_from: object
 
 class Checkpointer:
   def __init__(self, **opts: Unpack[CheckpointerOpts]):
@@ -40,10 +41,9 @@ class Checkpointer:
     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")
+    self.fn_hash_from = opts.get("fn_hash_from")
 
   @overload
   def __call__(self, fn: Fn, **override_opts: Unpack[CheckpointerOpts]) -> CachedFunction[Fn]: ...
@@ -56,6 +56,35 @@ class Checkpointer:
 
     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)
@@ -70,6 +99,14 @@ class CachedFunction(Generic[Fn]):
     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
@@ -82,41 +119,32 @@ class CachedFunction(Generic[Fn]):
     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]
+    return self.ident.raw_ident.depends
 
   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
+    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_id(self, args: tuple, kw: dict) -> str:
+  def get_call_id(self, args: tuple, kw: dict[str, object]) -> 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
+    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_id: str, coroutine: Coroutine):
+    return self.storage.store(call_id, AwaitableValue(await coroutine)).value
 
   def _call(self: CachedFunction[Callable[P, R]], args: tuple, kw: dict, rerun=False) -> R:
     full_args = self.bound + args
@@ -125,7 +153,7 @@ class CachedFunction(Generic[Fn]):
       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}"
+    call_id_long = f"{self.fn_dir}/{self.ident.fn_hash}/{call_id}"
 
     refresh = rerun \
       or not self.storage.exists(call_id) \
@@ -134,8 +162,8 @@ class CachedFunction(Generic[Fn]):
     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)
+      if iscoroutine(data):
+        return self._resolve_coroutine(call_id, data)
       return self.storage.store(call_id, data)
 
     try:
@@ -154,7 +182,7 @@ class CachedFunction(Generic[Fn]):
     return self._call(args, kw, True)
 
   @overload
-  def get(self: Callable[P, Awaitable[R]], *args: P.args, **kw: P.kwargs) -> R: ...
+  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):
@@ -172,7 +200,7 @@ class CachedFunction(Generic[Fn]):
     self.storage.delete(self.get_call_id(args, kw))
 
   def __repr__(self) -> str:
-    return f"<CachedFunction {self.fn.__name__} {self.fn_hash[:6]}>"
+    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:
@@ -182,3 +210,20 @@ class CachedFunction(Generic[Fn]):
       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/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/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/storages/storage.py

@@ -15,7 +15,7 @@ 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()

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/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:
+  def __init__(self, value):
+    self.value = value
+
+  def __await__(self):
+    yield
+    return self.value

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.dist-info → checkpointer-2.11.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.10.1
+Version: 2.11.0
 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,13 +130,45 @@ 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`.
 
-#### Example: Custom Storage Backend
+**Example: Custom Storage Backend**
 
 ```python
 from checkpointer import checkpoint, Storage

checkpointer-2.11.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+checkpointer/__init__.py,sha256=ayjFyHwvl_HRHwocY-hOJvAx0Ko5X9IMZrNT4CwfoMU,824
+checkpointer/checkpoint.py,sha256=zU67_PGrVCMP90clPTR7AA4vfxECqmeFr0jJElUL5iQ,9051
+checkpointer/fn_ident.py,sha256=-5XbovQowVyYCFc7JdT9z1NoIEiL8h9fi7alF_34Ils,4470
+checkpointer/object_hash.py,sha256=YlyFupQrg3V2mpzTLfOqpqlZWhoSCHliScQ4cKd36T0,8133
+checkpointer/print_checkpoint.py,sha256=aJCeWMRJiIR3KpyPk_UOKTaD906kArGrmLGQ3LqcVgo,1369
+checkpointer/test_checkpointer.py,sha256=-EvsMMNOOiIxhTcG97LLX0jUMWp534ko7qCKDSFWiA0,3802
+checkpointer/types.py,sha256=rAdjZNn1-jk35df7UVtby_qlp-8_18ucXfVCtS4RI_M,323
+checkpointer/utils.py,sha256=0cGVSlTnABgs3jI1uHoTfz353kkGa-qtTfe7jG4NCr0,2192
+checkpointer/storages/__init__.py,sha256=en32nTUltpCSgz8RVGS_leIHC1Y1G89IqG1ZqAb6qUo,236
+checkpointer/storages/memory_storage.py,sha256=Br30b1AyNOcNjjAaDui1mBjDKfhDbu--jV4WmJenzaE,1109
+checkpointer/storages/pickle_storage.py,sha256=xS96q8TvwxH_TOZsiKmrMrhFwQKZCcaH7XOiECmuwl8,1628
+checkpointer/storages/storage.py,sha256=5NWmnfsU2QY24NwKC5MZNe4h7pcYhADw-y800v17aYE,895
+checkpointer-2.11.0.dist-info/METADATA,sha256=GFRKdyTOIQwvmzIIrkV2owT8NO-X2ONRXdnadY9j5xY,11617
+checkpointer-2.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+checkpointer-2.11.0.dist-info/licenses/LICENSE,sha256=9xVsdtv_-uSyY9Xl9yujwAPm4-mjcCLeVy-ljwXEWbo,1059
+checkpointer-2.11.0.dist-info/RECORD,,

checkpointer-2.10.1.dist-info/RECORD

@@ -1,15 +0,0 @@
-checkpointer/__init__.py,sha256=FBdEwdyQS2tc4eQMwrb9qsFgjPu9QOyLOqgDK2c1rAI,837
-checkpointer/checkpoint.py,sha256=yysni7CKzQ2XRX0nvEd34ktNjUcHVZ0CnKMaqQX-f_8,6827
-checkpointer/fn_ident.py,sha256=hiPvm1lw9Aol6v-d5bjteWdg0VhrPQH4r1Wiow-Wzpo,4311
-checkpointer/object_hash.py,sha256=kfyhFFxvcjwkexw5TKkAtCyWU3PZXgsWGJ5dsOaQXmY,7695
-checkpointer/print_checkpoint.py,sha256=aJCeWMRJiIR3KpyPk_UOKTaD906kArGrmLGQ3LqcVgo,1369
-checkpointer/test_checkpointer.py,sha256=mvJu5EItLM95S6IdawgINVgqVCPBNx1bnMHJmwF_cwo,3731
-checkpointer/utils.py,sha256=_zvbkno7221_4l_FbROJMRrs-enx8BRS4FGxn8et9Ns,2751
-checkpointer/storages/__init__.py,sha256=en32nTUltpCSgz8RVGS_leIHC1Y1G89IqG1ZqAb6qUo,236
-checkpointer/storages/memory_storage.py,sha256=Br30b1AyNOcNjjAaDui1mBjDKfhDbu--jV4WmJenzaE,1109
-checkpointer/storages/pickle_storage.py,sha256=xS96q8TvwxH_TOZsiKmrMrhFwQKZCcaH7XOiECmuwl8,1628
-checkpointer/storages/storage.py,sha256=dR84h-IN644XLGe7invSlzLuXlFlU8QIbsDZa1xdtPM,889
-checkpointer-2.10.1.dist-info/METADATA,sha256=n7eiGDEoqRH-bEncwy27EM108HMwqVLHpDbYGY_PFZk,10908
-checkpointer-2.10.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-checkpointer-2.10.1.dist-info/licenses/LICENSE,sha256=9xVsdtv_-uSyY9Xl9yujwAPm4-mjcCLeVy-ljwXEWbo,1059
-checkpointer-2.10.1.dist-info/RECORD,,