checkpointer 2.1.0__py3-none-any.whl → 2.6.0__py3-none-any.whl

checkpointer/__init__.py

@@ -1,10 +1,20 @@
-from .checkpoint import Checkpointer, CheckpointFn, CheckpointError
-from .types import Storage
-from .function_body import get_function_hash
+import gc
 import tempfile
+from typing import Callable
+from .checkpoint import Checkpointer, CheckpointError, CheckpointFn
+from .object_hash import ObjectHash
+from .storages import MemoryStorage, PickleStorage, Storage
 
 create_checkpointer = Checkpointer
 checkpoint = Checkpointer()
 capture_checkpoint = Checkpointer(capture=True)
 memory_checkpoint = Checkpointer(format="memory", verbosity=0)
 tmp_checkpoint = Checkpointer(root_path=tempfile.gettempdir() + "/checkpoints")
+
+def cleanup_all(invalidated=True, expired=True):
+  for obj in gc.get_objects():
+    if isinstance(obj, CheckpointFn):
+      obj.cleanup(invalidated=invalidated, expired=expired)
+
+def get_function_hash(fn: Callable, capture=False) -> str:
+  return CheckpointFn(Checkpointer(capture=capture), fn).fn_hash

checkpointer/checkpoint.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 import inspect
-import relib.hashing as hashing
-from typing import Generic, TypeVar, Type, TypedDict, Callable, Unpack, Literal, Any, cast, overload
-from pathlib import Path
+import re
 from datetime import datetime
 from functools import update_wrapper
-from .types import Storage
-from .function_body import get_function_hash
-from .utils import unwrap_fn, sync_resolve_coroutine, resolved_awaitable
-from .storages import STORAGE_MAP
+from pathlib import Path
+from typing import Any, Callable, Generic, Iterable, Literal, 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 resolved_awaitable, sync_resolve_coroutine, unwrap_fn
 
 Fn = TypeVar("Fn", bound=Callable)
 
@@ -23,7 +23,7 @@ class CheckpointerOpts(TypedDict, total=False):
   root_path: Path | str | None
   when: bool
   verbosity: Literal[0, 1]
-  path: Callable[..., str] | None
+  hash_by: Callable | None
   should_expire: Callable[[datetime], bool] | None
   capture: bool
 
@@ -33,7 +33,7 @@ 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.path = opts.get("path")
+    self.hash_by = opts.get("hash_by")
     self.should_expire = opts.get("should_expire")
     self.capture = opts.get("capture", False)
 
@@ -50,37 +50,58 @@ class Checkpointer:
 
 class CheckpointFn(Generic[Fn]):
   def __init__(self, checkpointer: Checkpointer, fn: Fn):
-    wrapped = unwrap_fn(fn)
-    file_name = Path(wrapped.__code__.co_filename).name
-    update_wrapper(cast(Callable, self), wrapped)
-    storage = STORAGE_MAP[checkpointer.format] if isinstance(checkpointer.format, str) else checkpointer.format
     self.checkpointer = checkpointer
     self.fn = fn
-    self.fn_hash, self.depends = get_function_hash(wrapped, self.checkpointer.capture)
-    self.fn_id = f"{file_name}/{wrapped.__name__}"
+
+  def _set_ident(self, force=False):
+    if not hasattr(self, "fn_hash_raw") or force:
+      self.fn_hash_raw, self.depends = get_fn_ident(unwrap_fn(self.fn), self.checkpointer.capture)
+    return self
+
+  def _lazyinit(self):
+    wrapped = unwrap_fn(self.fn)
+    fn_file = Path(wrapped.__code__.co_filename).name
+    fn_name = re.sub(r"[^\w.]", "", wrapped.__qualname__)
+    update_wrapper(cast(Callable, self), wrapped)
+    store_format = self.checkpointer.format
+    Storage = STORAGE_MAP[store_format] if isinstance(store_format, str) else store_format
+    deep_hashes = [child._set_ident().fn_hash_raw for child in iterate_checkpoint_fns(self)]
+    self.fn_hash = str(ObjectHash().write_text(self.fn_hash_raw, iter=deep_hashes))
+    self.fn_subdir = f"{fn_file}/{fn_name}/{self.fn_hash[:16]}"
     self.is_async = inspect.iscoroutinefunction(wrapped)
-    self.storage = storage(checkpointer)
+    self.storage = Storage(self)
+    self.cleanup = self.storage.cleanup
+
+  def __getattribute__(self, name: str) -> Any:
+    return object.__getattribute__(self, "_getattribute")(name)
+
+  def _getattribute(self, name: str) -> Any:
+    setattr(self, "_getattribute", super().__getattribute__)
+    self._lazyinit()
+    return self._getattribute(name)
+
+  def reinit(self, recursive=False):
+    pointfns = list(iterate_checkpoint_fns(self)) if recursive else [self]
+    for pointfn in pointfns:
+      pointfn._set_ident(True)
+    for pointfn in pointfns:
+      pointfn._lazyinit()
 
   def get_checkpoint_id(self, args: tuple, kw: dict) -> str:
-    if not callable(self.checkpointer.path):
-      # TODO: use digest size before digesting instead of truncating the hash
-      call_hash = hashing.hash((self.fn_hash, args, kw), "blake2b")[:32]
-      return f"{self.fn_id}/{call_hash}"
-    checkpoint_id = self.checkpointer.path(*args, **kw)
-    if not isinstance(checkpoint_id, str):
-      raise CheckpointError(f"path function must return a string, got {type(checkpoint_id)}")
-    return checkpoint_id
+    hash_params = [self.checkpointer.hash_by(*args, **kw)] if self.checkpointer.hash_by else (args, kw)
+    call_hash = ObjectHash(self.fn_hash, *hash_params, digest_size=16)
+    return f"{self.fn_subdir}/{call_hash}"
 
   async def _store_on_demand(self, args: tuple, kw: dict, rerun: bool):
     checkpoint_id = self.get_checkpoint_id(args, kw)
     checkpoint_path = self.checkpointer.root_path / checkpoint_id
-    should_log = self.checkpointer.verbosity > 0
+    verbose = self.checkpointer.verbosity > 0
     refresh = rerun \
       or not self.storage.exists(checkpoint_path) \
       or (self.checkpointer.should_expire and self.checkpointer.should_expire(self.storage.checkpoint_date(checkpoint_path)))
 
     if refresh:
-      print_checkpoint(should_log, "MEMORIZING", checkpoint_id, "blue")
+      print_checkpoint(verbose, "MEMORIZING", checkpoint_id, "blue")
       data = self.fn(*args, **kw)
       if inspect.iscoroutine(data):
         data = await data
@@ -89,12 +110,12 @@ class CheckpointFn(Generic[Fn]):
 
     try:
       data = self.storage.load(checkpoint_path)
-      print_checkpoint(should_log, "REMEMBERED", checkpoint_id, "green")
+      print_checkpoint(verbose, "REMEMBERED", checkpoint_id, "green")
       return data
     except (EOFError, FileNotFoundError):
-      print_checkpoint(should_log, "CORRUPTED", checkpoint_id, "yellow")
-      self.storage.delete(checkpoint_path)
-      return await self._store_on_demand(args, kw, rerun)
+      pass
+    print_checkpoint(verbose, "CORRUPTED", checkpoint_id, "yellow")
+    return await self._store_on_demand(args, kw, True)
 
   def _call(self, args: tuple, kw: dict, rerun=False):
     if not self.checkpointer.when:
@@ -107,8 +128,8 @@ class CheckpointFn(Generic[Fn]):
     try:
       val = self.storage.load(checkpoint_path)
       return resolved_awaitable(val) if self.is_async else val
-    except:
-      raise CheckpointError("Could not load checkpoint")
+    except Exception as ex:
+      raise CheckpointError("Could not load checkpoint") from ex
 
   def exists(self, *args: tuple, **kw: dict) -> bool:
     return self.storage.exists(self.checkpointer.root_path / self.get_checkpoint_id(args, kw))
@@ -116,3 +137,15 @@ class CheckpointFn(Generic[Fn]):
   __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))
   get: Fn = cast(Fn, lambda self, *args, **kw: self._get(args, kw))
+
+  def __repr__(self) -> str:
+    return f"<CheckpointFn {self.fn.__name__} {self.fn_hash[:6]}>"
+
+def iterate_checkpoint_fns(pointfn: CheckpointFn, visited: set[CheckpointFn] = set()) -> Iterable[CheckpointFn]:
+  visited = visited or set()
+  if pointfn not in visited:
+    yield pointfn
+    visited.add(pointfn)
+    for depend in pointfn.depends:
+      if isinstance(depend, CheckpointFn):
+        yield from iterate_checkpoint_fns(depend, visited)

checkpointer/fn_ident.py

@@ -0,0 +1,94 @@
+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, Generator, Type, TypeGuard
+from .object_hash import ObjectHash
+from .utils import AttrDict, distinct, get_cell_contents, iterate_and_upcoming, transpose, unwrap_fn
+
+cwd = Path.cwd()
+
+def is_class(obj) -> TypeGuard[Type]:
+  # isinstance works too, but needlessly triggers _lazyinit()
+  return issubclass(type(obj), type)
+
+def extract_classvars(code: CodeType, scope_vars: AttrDict) -> dict[str, dict[str, Type]]:
+  attr_path: tuple[str, ...] = ()
+  scope_obj = None
+  classvars: dict[str, dict[str, Type]] = {}
+  for instr, upcoming_instrs in iterate_and_upcoming(dis.get_instructions(code)):
+    if instr.opname in scope_vars and not attr_path:
+      attrs = takewhile(lambda instr: instr.opname == "LOAD_ATTR", upcoming_instrs)
+      attr_path = (instr.opname, instr.argval, *(str(x.argval) for x in attrs))
+    elif instr.opname == "CALL":
+      obj = scope_vars.get_at(attr_path)
+      attr_path = ()
+      if is_class(obj):
+        scope_obj = obj
+    elif instr.opname in ("STORE_FAST", "STORE_DEREF") and scope_obj:
+      load_key = instr.opname.replace("STORE", "LOAD")
+      classvars.setdefault(load_key, {})[instr.argval] = scope_obj
+      scope_obj = None
+  return classvars
+
+def extract_scope_values(code: CodeType, scope_vars: AttrDict) -> Generator[tuple[tuple[str, ...], Any], None, None]:
+  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)):
+    if instr.opname in scope_vars:
+      attrs = takewhile(lambda instr: instr.opname == "LOAD_ATTR", upcoming_instrs)
+      attr_path: tuple[str, ...] = (instr.opname, instr.argval, *(str(x.argval) for x in attrs))
+      val = scope_vars.get_at(attr_path)
+      if val is not None:
+        yield attr_path, val
+  for const in code.co_consts:
+    if isinstance(const, CodeType):
+      yield from extract_scope_values(const, scope_vars)
+
+def get_self_value(fn: Callable) -> type | object | None:
+  if isinstance(fn, MethodType):
+    return fn.__self__
+  parts = tuple(fn.__qualname__.split(".")[:-1])
+  cls = parts and AttrDict(fn.__globals__).get_at(parts)
+  if is_class(cls):
+    return cls
+
+def get_fn_captured_vals(fn: Callable) -> list[Any]:
+  self_value = get_self_value(fn)
+  scope_vars = AttrDict({
+    "LOAD_FAST": AttrDict({"self": self_value} if self_value else {}),
+    "LOAD_DEREF": AttrDict(get_cell_contents(fn)),
+    "LOAD_GLOBAL": AttrDict(fn.__globals__),
+  })
+  vals = dict(extract_scope_values(fn.__code__, scope_vars))
+  return list(vals.values())
+
+def is_user_fn(candidate_fn) -> TypeGuard[Callable]:
+  if not isinstance(candidate_fn, (FunctionType, MethodType)):
+    return False
+  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]]:
+  from .checkpoint import CheckpointFn
+  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))
+  for child_fn in child_fns:
+    if isinstance(child_fn, CheckpointFn):
+      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
+  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)]
+  fn_hash = str(ObjectHash(fn, unwrapped_depends).update(depend_captured_vals, tolerate_errors=True))
+  return fn_hash, depends

checkpointer/object_hash.py

@@ -0,0 +1,192 @@
+import ctypes
+import hashlib
+import io
+import re
+from collections.abc import Iterable
+from contextlib import nullcontext, suppress
+from decimal import Decimal
+from itertools import chain
+from pickle import HIGHEST_PROTOCOL as PROTOCOL
+from types import BuiltinFunctionType, FunctionType, GeneratorType, MethodType, ModuleType, UnionType
+from typing import Any, TypeAliasType, TypeVar
+from .utils import ContextVar, get_fn_body
+
+np, torch = None, None
+
+with suppress(Exception):
+  import numpy as np
+
+with suppress(Exception):
+  import torch
+
+def encode_type(t: type | FunctionType) -> str:
+  return f"{t.__module__}:{t.__qualname__}"
+
+def encode_type_of(v: Any) -> str:
+  return encode_type(type(v))
+
+class ObjectHashError(Exception):
+  def __init__(self, obj: Any, 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:
+    self.hash = hashlib.blake2b(digest_size=digest_size)
+    self.current: dict[int, int] = {}
+    self.tolerate_errors = ContextVar(tolerate_errors)
+    self.update(iter=chain(objs, iter))
+
+  def copy(self) -> "ObjectHash":
+    new = ObjectHash(tolerate_errors=self.tolerate_errors.value)
+    new.hash = self.hash.copy()
+    return new
+
+  def hexdigest(self) -> str:
+    return self.hash.hexdigest()
+
+  __str__ = hexdigest
+
+  def nested_hash(self, *objs: Any) -> str:
+    return ObjectHash(iter=objs, tolerate_errors=self.tolerate_errors.value).hexdigest()
+
+  def write_bytes(self, *data: bytes, iter: Iterable[bytes] = ()) -> "ObjectHash":
+    for d in chain(data, iter):
+      self.hash.update(d)
+    return self
+
+  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":
+    return self.write_bytes(":".join(map(str, args)).encode())
+
+  def update(self, *objs: Any, iter: Iterable[Any] = (), 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:
+          self._update_one(obj)
+        except Exception as ex:
+          if self.tolerate_errors.value:
+            self.header("error").update(type(ex))
+            continue
+          raise ObjectHashError(obj, ex) from ex
+      return self
+
+  def _update_one(self, obj: Any) -> None:
+    match obj:
+      case None:
+        self.header("null")
+
+      case bool() | int() | float() | complex() | Decimal() | ObjectHash():
+        self.header("number", encode_type_of(obj), obj)
+
+      case str() | bytes() | bytearray() | memoryview():
+        b = obj.encode() if isinstance(obj, str) else obj
+        self.header("bytes", encode_type_of(obj), len(b)).write_bytes(b)
+
+      case set() | frozenset():
+        try:
+          items = sorted(obj)
+          header = "set"
+        except:
+          items = sorted(map(self.nested_hash, obj))
+          header = "set-unsortable"
+        self.header(header, encode_type_of(obj), len(obj)).update(iter=items)
+
+      case TypeVar():
+        self.header("TypeVar").update(obj.__name__, obj.__bound__, obj.__constraints__, obj.__contravariant__, obj.__covariant__)
+
+      case TypeAliasType():
+        self.header("TypeAliasType").update(obj.__name__, obj.__value__)
+
+      case UnionType():
+        self.header("UnionType").update(obj.__args__)
+
+      case BuiltinFunctionType():
+        self.header("builtin", obj.__qualname__)
+
+      case FunctionType():
+        self.header("function", encode_type(obj)).update(get_fn_body(obj), obj.__defaults__, obj.__kwdefaults__, obj.__annotations__)
+
+      case MethodType():
+        self.header("method").update(obj.__func__, obj.__self__.__class__)
+
+      case ModuleType():
+        self.header("module", obj.__name__, obj.__file__)
+
+      case GeneratorType():
+        self.header("generator", obj.__qualname__)._update_iterator(obj)
+
+      case io.TextIOWrapper() | io.FileIO() | io.BufferedRandom() | io.BufferedWriter() | io.BufferedReader():
+        self.header("file", encode_type_of(obj)).update(obj.name, obj.mode, obj.tell())
+
+      case type():
+        self.header("type", encode_type(obj))
+
+      case _ if np and isinstance(obj, np.dtype):
+        self.header("dtype").update(obj.__class__, obj.descr)
+
+      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))
+        else:
+          array = np.ascontiguousarray(obj if obj.base is None else obj.base).view(np.uint8)
+          self.write_bytes(array.data)
+
+      case _ if torch and isinstance(obj, torch.Tensor):
+        self.header("tensor", encode_type_of(obj), obj.dtype, tuple(obj.shape), obj.stride(), obj.device)
+        if obj.device.type != "cpu":
+          obj = obj.cpu()
+        storage = obj.storage()
+        buffer = (ctypes.c_ubyte * storage.nbytes()).from_address(storage.data_ptr())
+        self.write_bytes(memoryview(buffer))
+
+      case _ if id(obj) in self.current:
+        self.header("circular", self.current[id(obj)])
+
+      case _:
+        try:
+          self.current[id(obj)] = len(self.current)
+          match obj:
+            case list() | tuple():
+              self.header("list", encode_type_of(obj), len(obj)).update(iter=obj)
+            case dict():
+              try:
+                items = sorted(obj.items())
+                header = "dict"
+              except:
+                items = sorted((self.nested_hash(key), val) for key, val in obj.items())
+                header = "dict-unsortable"
+              self.header(header, encode_type_of(obj), len(obj)).update(iter=chain.from_iterable(items))
+            case _:
+              self._update_object(obj)
+        finally:
+          del self.current[id(obj)]
+
+  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":
+    self.header("instance", encode_type_of(obj))
+    reduced = None
+    with suppress(Exception):
+      reduced = obj.__reduce_ex__(PROTOCOL)
+    with suppress(Exception):
+      reduced = reduced or obj.__reduce__()
+    if isinstance(reduced, str):
+      return self.header("reduce-str").update(reduced)
+    if reduced:
+      reduced = list(reduced)
+      it = reduced.pop(3) if len(reduced) >= 4 else None
+      return self.header("reduce").update(reduced)._update_iterator(it or ())
+    if state := hasattr(obj, "__getstate__") and obj.__getstate__():
+      return self.header("getstate").update(state)
+    if len(getattr(obj, "__slots__", [])):
+      slots = {slot: getattr(obj, slot, None) for slot in getattr(obj, "__slots__")}
+      return self.header("slots").update(slots)
+    if d := getattr(obj, "__dict__", {}):
+      return self.header("dict").update(d)
+    repr_str = re.sub(r"\s+(at\s+0x[0-9a-fA-F]+)(>)$", r"\2", repr(obj))
+    return self.header("repr").update(repr_str)

checkpointer/storages/__init__.py

@@ -1,5 +1,5 @@
 from typing import Type
-from ..types import Storage
+from .storage import Storage
 from .pickle_storage import PickleStorage
 from .memory_storage import MemoryStorage
 from .bcolz_storage import BcolzStorage

checkpointer/storages/bcolz_storage.py

@@ -1,7 +1,7 @@
 import shutil
 from pathlib import Path
 from datetime import datetime
-from ..types import Storage
+from .storage import Storage
 
 def get_data_type_str(x):
   if isinstance(x, tuple):
@@ -73,9 +73,8 @@ class BcolzStorage(Storage):
 
   def delete(self, path):
     # NOTE: Not recursive
-    metapath = get_metapath(path)
-    try:
-      shutil.rmtree(metapath)
-      shutil.rmtree(path)
-    except FileNotFoundError:
-      pass
+    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")

checkpointer/storages/memory_storage.py

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

checkpointer/storages/pickle_storage.py

@@ -1,31 +1,45 @@
 import pickle
+import shutil
 from pathlib import Path
 from datetime import datetime
-from ..types import Storage
+from .storage import Storage
 
 def get_path(path: Path):
   return path.with_name(f"{path.name}.pkl")
 
 class PickleStorage(Storage):
-  def exists(self, path):
-    return get_path(path).exists()
-
-  def checkpoint_date(self, path):
-    return datetime.fromtimestamp(get_path(path).stat().st_mtime)
-
   def store(self, path, data):
     full_path = get_path(path)
     full_path.parent.mkdir(parents=True, exist_ok=True)
     with full_path.open("wb") as file:
       pickle.dump(data, file, -1)
 
+  def exists(self, path):
+    return get_path(path).exists()
+
+  def checkpoint_date(self, path):
+    return datetime.fromtimestamp(get_path(path).stat().st_mtime)
+
   def load(self, path):
-    full_path = get_path(path)
-    with full_path.open("rb") as file:
+    with get_path(path).open("rb") as file:
       return pickle.load(file)
 
   def delete(self, path):
-    try:
-      get_path(path).unlink()
-    except FileNotFoundError:
-      pass
+    get_path(path).unlink(missing_ok=True)
+
+  def cleanup(self, invalidated=True, expired=True):
+    version_path = self.checkpointer.root_path.resolve() / self.checkpoint_fn.fn_subdir
+    fn_path = version_path.parent
+    if invalidated:
+      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__}")
+    if expired and self.checkpointer.should_expire:
+      count = 0
+      for pkl_path in fn_path.rglob("*.pkl"):
+        path = pkl_path.with_suffix("")
+        if self.checkpointer.should_expire(self.checkpoint_date(path)):
+          count += 1
+          self.delete(path)
+      print(f"Removed {count} expired checkpoints for {self.checkpoint_fn.__qualname__}")

checkpointer/{types.py → storages/storage.py}

@@ -4,20 +4,24 @@ from pathlib import Path
 from datetime import datetime
 
 if TYPE_CHECKING:
-  from .checkpoint import Checkpointer
+  from ..checkpoint import Checkpointer, CheckpointFn
 
 class Storage:
   checkpointer: Checkpointer
+  checkpoint_fn: CheckpointFn
 
-  def __init__(self, checkpointer: Checkpointer):
-    self.checkpointer = checkpointer
+  def __init__(self, checkpoint_fn: CheckpointFn):
+    self.checkpointer = checkpoint_fn.checkpointer
+    self.checkpoint_fn = checkpoint_fn
+
+  def store(self, path: Path, data: Any) -> None: ...
 
   def exists(self, path: Path) -> bool: ...
 
   def checkpoint_date(self, path: Path) -> datetime: ...
 
-  def store(self, path: Path, data: Any) -> None: ...
-
   def load(self, path: Path) -> Any: ...
 
   def delete(self, path: Path) -> None: ...
+
+  def cleanup(self, invalidated=True, expired=True): ...

checkpointer/test_checkpointer.py

@@ -0,0 +1,159 @@
+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
+
+def global_multiply(a, b):
+  return a * b
+
+@pytest.fixture(autouse=True)
+def run_before_and_after_tests(tmpdir):
+  global checkpoint
+  checkpoint = checkpoint(root_path=tmpdir)
+  yield
+
+def test_basic_caching():
+  @checkpoint
+  def square(x: int) -> int:
+    return x ** 2
+
+  result1 = square(4)
+  result2 = square(4)
+
+  assert result1 == result2 == 16
+
+def test_cache_invalidation():
+  @checkpoint
+  def multiply(a, b):
+    return a * b
+
+  @checkpoint
+  def helper(x):
+    return multiply(x + 1, 2)
+
+  @checkpoint
+  def compute(a, b):
+    return helper(a) + helper(b)
+
+  result1 = compute(3, 4)
+  assert result1 == 18
+
+def test_layered_caching():
+  dev_checkpoint = checkpoint(when=True)
+
+  @checkpoint(format="memory")
+  @dev_checkpoint
+  def expensive_function(x):
+    return x ** 2
+
+  assert expensive_function(4) == 16
+  assert expensive_function(4) == 16
+
+def test_recursive_caching1():
+  @checkpoint
+  def fib(n: int) -> int:
+    return fib(n - 1) + fib(n - 2) if n > 1 else n
+
+  assert fib(10) == 55
+  assert fib.get(10) == 55
+  assert fib.get(5) == 5
+
+def test_recursive_caching2():
+  @checkpoint
+  def fib(n: int) -> int:
+    return fib.fn(n - 1) + fib.fn(n - 2) if n > 1 else n
+
+  assert fib(10) == 55
+  assert fib.get(10) == 55
+  with pytest.raises(CheckpointError):
+    fib.get(5)
+
+@pytest.mark.asyncio
+async def test_async_caching():
+  @checkpoint(format="memory")
+  async def async_square(x: int) -> int:
+    await asyncio.sleep(0.1)
+    return x ** 2
+
+  result1 = await async_square(3)
+  result2 = await async_square.get(3)
+
+  assert result1 == result2 == 9
+
+def test_force_recalculation():
+  @checkpoint
+  def square(x: int) -> int:
+    return x ** 2
+
+  assert square(5) == 25
+  square.rerun(5)
+  assert square.get(5) == 25
+
+def test_multi_layer_decorator():
+  @checkpoint(format="memory")
+  @checkpoint(format="pickle")
+  def add(a, b):
+    return a + b
+
+  assert add(2, 3) == 5
+  assert add.get(2, 3) == 5
+
+def test_capture():
+  item_dict = AttrDict({"a": 1, "b": 1})
+
+  @checkpoint(capture=True)
+  def test_whole():
+    return item_dict
+
+  @checkpoint(capture=True)
+  def test_a():
+    return item_dict.a + 1
+
+  init_hash_a = test_a.fn_hash
+  init_hash_whole = test_whole.fn_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
+  item_dict.a += 1
+  test_a.reinit()
+  assert test_a.fn_hash != init_hash_a
+
+def test_depends():
+  def multiply_wrapper(a, b):
+    return global_multiply(a, b)
+
+  def helper(a, b):
+    return multiply_wrapper(a + 1, b + 1)
+
+  @checkpoint
+  def test_a(a, b):
+    return helper(a, b)
+
+  @checkpoint
+  def test_b(a, b):
+    return test_a(a, b) + multiply_wrapper(a, b)
+
+  assert set(test_a.depends) == {test_a.fn, helper, multiply_wrapper, global_multiply}
+  assert set(test_b.depends) == {test_b.fn, test_a, multiply_wrapper, global_multiply}
+
+def test_lazy_init():
+  @checkpoint
+  def fn1(x):
+    return fn2(x)
+
+  @checkpoint
+  def fn2(x):
+    return fn1(x)
+
+  assert type(object.__getattribute__(fn1, "_getattribute")) == MethodType
+  with pytest.raises(AttributeError):
+    object.__getattribute__(fn1, "fn_hash")
+  assert fn1.fn_hash == object.__getattribute__(fn1, "fn_hash")
+  assert type(object.__getattribute__(fn1, "_getattribute")) == MethodWrapperType
+  assert set(fn1.depends) == {fn1.fn, fn2}
+  assert set(fn2.depends) == {fn1, fn2.fn}

checkpointer/utils.py

@@ -1,22 +1,44 @@
-from typing import Generator, Coroutine, Iterable, Any, cast
-from types import CellType, coroutine
-from itertools import islice
+import inspect
+import tokenize
+from contextlib import contextmanager
+from io import StringIO
+from types import coroutine
+from typing import Any, Callable, Coroutine, Generator, Iterable, cast
 
-class AttrDict(dict):
-  def __init__(self, *args, **kwargs):
-    super(AttrDict, self).__init__(*args, **kwargs)
-    self.__dict__ = self
+def distinct[T](seq: Iterable[T]) -> list[T]:
+  return list(dict.fromkeys(seq))
 
-  def __getattribute__(self, name: str) -> Any:
-    return super().__getattribute__(name)
+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 unwrap_fn[T](fn: T, checkpoint_fn=False) -> T:
+def get_cell_contents(fn: Callable) -> Generator[tuple[str, Any], None, None]:
+  for key, cell in zip(fn.__code__.co_freevars, fn.__closure__ or []):
+    try:
+      yield (key, cell.cell_contents)
+    except ValueError:
+      pass
+
+def unwrap_fn[T: Callable](fn: T, checkpoint_fn=False) -> T:
   from .checkpoint import CheckpointFn
-  while hasattr(fn, "__wrapped__"):
-    if checkpoint_fn and isinstance(fn, CheckpointFn):
-      return fn
+  while True:
+    if (checkpoint_fn and isinstance(fn, CheckpointFn)) or not hasattr(fn, "__wrapped__"):
+      return cast(T, fn)
     fn = getattr(fn, "__wrapped__")
-  return fn
+
+async def resolved_awaitable[T](value: T) -> T:
+  return value
 
 @coroutine
 def coroutine_as_generator[T](coroutine: Coroutine[None, None, T]) -> Generator[None, None, T]:
@@ -26,27 +48,65 @@ def coroutine_as_generator[T](coroutine: Coroutine[None, None, T]) -> Generator[
 def sync_resolve_coroutine[T](coroutine: Coroutine[None, None, T]) -> T:
   gen = cast(Generator, coroutine_as_generator(coroutine))
   try:
-    while True: next(gen)
+    while True:
+      next(gen)
   except StopIteration as ex:
     return ex.value
 
-async def resolved_awaitable[T](value: T) -> T:
-  return value
+class AttrDict(dict):
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.__dict__ = self
 
-def iterate_and_upcoming[T](l: list[T]) -> Iterable[tuple[T, Iterable[T]]]:
-  for i, item in enumerate(l):
-    yield item, islice(l, i + 1, None)
+  def __getattribute__(self, name: str) -> Any:
+    return super().__getattribute__(name)
 
-def get_at_attr(d: dict, keys: tuple[str, ...]) -> Any:
-  try:
-    for key in keys:
-      d = getattr(d, key)
-  except AttributeError:
-    return None
-  return d
+  def __setattr__(self, name: str, value: Any) -> None:
+    return super().__setattr__(name, value)
 
-def get_cell_contents(cell: CellType) -> Any:
-  try:
-    return cell.cell_contents
-  except ValueError:
-    return None
+  def set(self, d: dict) -> "AttrDict":
+    if not d:
+      return self
+    return AttrDict({**self, **d})
+
+  def delete(self, *attrs: str) -> "AttrDict":
+    d = AttrDict(self)
+    for attr in attrs:
+      del d[attr]
+    return d
+
+  def get_at(self, attrs: tuple[str, ...]) -> Any:
+    d = self
+    for attr in attrs:
+      d = getattr(d, attr, None)
+    return d
+
+class ContextVar[T]:
+  def __init__(self, value: T):
+    self.value = value
+
+  @contextmanager
+  def set(self, value: T):
+    self.value, old = value, self.value
+    try:
+      yield
+    finally:
+      self.value = old
+
+class iterate_and_upcoming[T]:
+  def __init__(self, it: Iterable[T]) -> None:
+    self.it = iter(it)
+    self.previous: tuple[()] | tuple[T] = ()
+
+  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()
+
+  def _tracked_iter(self):
+    for x in self.it:
+      self.previous = (x,)
+      yield x

{checkpointer-2.1.0.dist-info → checkpointer-2.6.0.dist-info}/METADATA

@@ -1,18 +1,20 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.1.0
+Version: 2.6.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
-License: Copyright 2024 Hampus Hallman
+License: Copyright 2018-2025 Hampus Hallman
         
         Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
         
         The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
         
         THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+License-File: LICENSE
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.12
-Requires-Dist: relib
 Description-Content-Type: text/markdown
 
 # checkpointer · [![License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/Reddan/checkpointer/blob/master/LICENSE) [![pypi](https://img.shields.io/pypi/v/checkpointer)](https://pypi.org/project/checkpointer/) [![Python 3.12](https://img.shields.io/badge/python-3.12-blue)](https://pypi.org/project/checkpointer/)
@@ -28,6 +30,7 @@ Adding or removing `@checkpoint` doesn't change how your code works. You can app
 - ⏲️ **Custom Expiration Logic**: Automatically invalidate old checkpoints.
 - 📂 **Flexible Path Configuration**: Control where checkpoints are stored.
 - 📦 **Captured Variables Handling**: Optionally include captured variables in cache invalidation.
+- ⚡ **Custom Argument Hashing**: Override argument hashing for speed or specialized hashing logic.
 
 ---
 
@@ -61,9 +64,9 @@ When you use `@checkpoint`, the function's **arguments** (`args`, `kwargs`) are
 
 Additionally, `checkpointer` ensures that caches are invalidated when a function's implementation or any of its dependencies change. Each function is assigned a hash based on:
 
-1. **Its source code**: Changes to the function's code update its hash.
-2. **Dependent functions**: If a function calls others, changes in those dependencies will also update the hash.
-3. **Captured variables**: (Optional) If `capture=True`, changes to captured variables and global variables will also update the hash.
+1. **Function Code**: The hash updates when the function’s own source code changes.
+2. **Dependencies**: If the function calls other user-defined functions, changes in those dependencies also update the hash.
+3. **External Variables** *(with `capture=True`)*: Any global or closure-based variables used by the function are included in its hash, so changes to those variables also trigger cache invalidation.
 
 ### Example: Cache Invalidation
 
@@ -108,7 +111,7 @@ Layer caches by stacking checkpoints:
 @dev_checkpoint  # Adds caching during development
 def some_expensive_function():
     print("Performing a time-consuming operation...")
-    return sum(i * i for i in range(10**6))
+    return sum(i * i for i in range(10**8))
 ```
 
 - **In development**: Both `dev_checkpoint` and `memory` caches are active.
@@ -153,11 +156,21 @@ Access cached results without recalculating:
 stored_result = expensive_function.get(4)
 ```
 
+### Refresh Function Hash
+
+If `capture=True`, you might need to re-hash a function during the same Python session. For that, call `reinit`:
+
+```python
+expensive_function.reinit()
+```
+
+This tells `checkpointer` to recalculate the function hash, reflecting changes in captured variables.
+
 ---
 
 ## Storage Backends
 
-`checkpointer` works with both built-in and custom storage backends, so you can use what's provided or roll your own as needed.
+`checkpointer` works with built-in and custom storage backends, so you can use what's provided or roll your own as needed.
 
 ### Built-In Backends
 
@@ -189,10 +202,10 @@ from checkpointer import checkpoint, Storage
 from datetime import datetime
 
 class CustomStorage(Storage):
-    def exists(self, path) -> bool: ...  # Check if a checkpoint exists at the given path
-    def checkpoint_date(self, path) -> datetime: ...  # Return the date the checkpoint was created
-    def store(self, path, data): ...  # Save the checkpoint data
-    def load(self, path): ...  # Return the checkpoint data
+    def exists(self, path) -> bool: ...  # Check if a checkpoint exists
+    def checkpoint_date(self, path) -> datetime: ...  # Get the checkpoint's timestamp
+    def store(self, path, data): ...  # Save data to the checkpoint
+    def load(self, path): ...  # Load data from the checkpoint
     def delete(self, path): ...  # Delete the checkpoint
 
 @checkpoint(format=CustomStorage)
@@ -200,21 +213,21 @@ def custom_cached(x: int):
     return x ** 2
 ```
 
-Using a custom backend lets you tailor storage to your application, whether it involves databases, cloud storage, or custom file formats.
+Use a custom backend to integrate with databases, cloud storage, or specialized file formats.
 
 ---
 
 ## Configuration Options ⚙️
 
-| Option          | Type                              | Default              | Description                                    |
-|-----------------|-----------------------------------|----------------------|------------------------------------------------|
-| `capture`       | `bool`                            | `False`              | Include captured variables in function hashes. |
-| `format`        | `"pickle"`, `"memory"`, `Storage` | `"pickle"`           | Storage backend format.                        |
-| `root_path`     | `Path`, `str`, or `None`          | ~/.cache/checkpoints | Root directory for storing checkpoints.        |
-| `when`          | `bool`                            | `True`               | Enable or disable checkpointing.               |
-| `verbosity`     | `0` or `1`                        | `1`                  | Logging verbosity.                             |
-| `path`          | `Callable[..., str]`              | `None`               | Custom path for checkpoint storage.            |
-| `should_expire` | `Callable[[datetime], bool]`      | `None`               | Custom expiration logic.                       |
+| Option          | Type                                | Default              | Description                                               |
+|-----------------|-------------------------------------|----------------------|-----------------------------------------------------------|
+| `capture`       | `bool`                              | `False`              | Include captured variables in function hashes.            |
+| `format`        | `"pickle"`, `"memory"`, `Storage`   | `"pickle"`           | Storage backend format.                                   |
+| `root_path`     | `Path`, `str`, or `None`            | ~/.cache/checkpoints | Root directory for storing checkpoints.                   |
+| `when`          | `bool`                              | `True`               | Enable or disable checkpointing.                          |
+| `verbosity`     | `0` or `1`                          | `1`                  | Logging verbosity.                                        |
+| `should_expire` | `Callable[[datetime], bool]`        | `None`               | Custom expiration logic.                                  |
+| `hash_by`       | `Callable[..., Any]`                | `None`               | Custom function that transforms arguments before hashing. |
 
 ---
 

checkpointer-2.6.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+checkpointer/__init__.py,sha256=ZJ6frUNgkklUi85b5uXTyTfRzMvZgQOJY-ZOnu7jh78,777
+checkpointer/checkpoint.py,sha256=2CAlNPnMLeou8SwHWeIjJ4kQhSqnpM9EBP_f-cYdSlE,6123
+checkpointer/fn_ident.py,sha256=TEM_SdjzQ5OgRJnqNgyYZZKONx7tM7Dk4bNM5TB_RyY,4311
+checkpointer/object_hash.py,sha256=o7Qr9VViWPmG1Fy1S3RediCagRydEVgTVZrZ9nRN54E,7215
+checkpointer/print_checkpoint.py,sha256=21aeqgM9CMjNAJyScqFmXCWWfh3jBIn7o7i5zJkZGaA,1369
+checkpointer/test_checkpointer.py,sha256=uJ2Pg9Miq1W0l28eNlRhMjuT_R8c-ygYwp3KP3VW8Os,3600
+checkpointer/utils.py,sha256=Rvm2NaJHtPTusM7fyHz_w9HUy_fqQfx8S1fr5CBWGL0,3047
+checkpointer/storages/__init__.py,sha256=Kl4Og5jhYxn6m3tB_kTMsabf4_eWVLmFVAoC-pikNQE,301
+checkpointer/storages/bcolz_storage.py,sha256=3QkSUSeG5s2kFuVV_LZpzMn1A5E7kqC7jk7w35c0NyQ,2314
+checkpointer/storages/memory_storage.py,sha256=S5ayOZE_CyaFQJ-vSgObTanldPzG3gh3NksjNAc7vsk,1282
+checkpointer/storages/pickle_storage.py,sha256=lJ0ton9ib3eifiny8XtPSNsx-w4Cm8oYUlbmKob34xU,1554
+checkpointer/storages/storage.py,sha256=_m18Z8TKrdAbi6YYYQmuNOnhna4RB2sJDn1v3liaU3U,721
+checkpointer-2.6.0.dist-info/METADATA,sha256=RWzn1RNHN76iF4hNkomVLouWvtjHtOody5doC4ajRk8,10556
+checkpointer-2.6.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+checkpointer-2.6.0.dist-info/licenses/LICENSE,sha256=9xVsdtv_-uSyY9Xl9yujwAPm4-mjcCLeVy-ljwXEWbo,1059
+checkpointer-2.6.0.dist-info/RECORD,,

{checkpointer-2.1.0.dist-info → checkpointer-2.6.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.26.3
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any

{checkpointer-2.1.0.dist-info → checkpointer-2.6.0.dist-info}/licenses/LICENSE

@@ -1,4 +1,4 @@
-Copyright 2024 Hampus Hallman
+Copyright 2018-2025 Hampus Hallman
 
 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 

checkpointer/function_body.py

@@ -1,80 +0,0 @@
-from __future__ import annotations
-import dis
-import inspect
-import tokenize
-from io import StringIO
-from collections.abc import Callable
-from itertools import chain, takewhile
-from operator import itemgetter
-from pathlib import Path
-from typing import Any, TypeGuard, TYPE_CHECKING
-from types import CodeType, FunctionType
-from relib import transpose, hashing, merge_dicts, drop_none
-from .utils import unwrap_fn, iterate_and_upcoming, get_cell_contents, AttrDict, get_at_attr
-
-if TYPE_CHECKING:
-  from .checkpoint import CheckpointFn
-
-cwd = Path.cwd()
-
-def extract_scope_values(code: CodeType, scope_vars: dict[str, Any], closure = False) -> dict[tuple[str, ...], Any]:
-  opname = "LOAD_GLOBAL" if not closure else "LOAD_DEREF"
-  scope_values_by_path: dict[tuple[str, ...], Any] = {}
-  instructions = list(dis.get_instructions(code))
-
-  for instr, upcoming_instrs in iterate_and_upcoming(instructions):
-    if instr.opname == opname:
-      name = instr.argval
-      attrs = takewhile(lambda instr: instr.opname == "LOAD_ATTR", upcoming_instrs)
-      attr_path = (name, *(instr.argval for instr in attrs))
-      scope_values_by_path[attr_path] = get_at_attr(scope_vars, attr_path)
-
-  children = (extract_scope_values(const, scope_vars, closure) for const in code.co_consts if isinstance(const, CodeType))
-  return merge_dicts(scope_values_by_path, *children)
-
-def get_fn_captured_vals(fn: Callable) -> list[Any]:
-  closure_scope = {k: get_cell_contents(v) for k, v in zip(fn.__code__.co_freevars, fn.__closure__ or [])}
-  global_vals = extract_scope_values(fn.__code__, AttrDict(fn.__globals__), closure=False)
-  closure_vals = extract_scope_values(fn.__code__, AttrDict(closure_scope), closure=True)
-  sorted_items = chain(sorted(global_vals.items()), sorted(closure_vals.items()))
-  return drop_none(map(itemgetter(1), sorted_items))
-
-def get_fn_body(fn: Callable) -> str:
-  source = "".join(inspect.getsourcelines(fn)[0])
-  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_fn_path(fn: Callable) -> Path:
-  return Path(inspect.getfile(fn)).resolve()
-
-def is_user_fn(candidate_fn) -> TypeGuard[Callable]:
-  return isinstance(candidate_fn, FunctionType) \
-    and cwd in get_fn_path(candidate_fn).parents
-
-def append_fn_depends(checkpoint_fns: set[CheckpointFn], captured_vals_by_fn: dict[Callable, list[Any]], fn: Callable, capture: bool) -> None:
-  from .checkpoint import CheckpointFn
-  captured_vals = get_fn_captured_vals(fn)
-  captured_vals_by_fn[fn] = [v for v in captured_vals if capture and not callable(v)]
-  callables = [unwrap_fn(val, checkpoint_fn=True) for val in captured_vals if callable(val)]
-  depends = {val for val in callables if is_user_fn(val)}
-  checkpoint_fns.update({val for val in callables if isinstance(val, CheckpointFn)})
-  not_appended = depends - captured_vals_by_fn.keys()
-  captured_vals_by_fn.update({fn: [] for fn in not_appended})
-  for child_fn in not_appended:
-    append_fn_depends(checkpoint_fns, captured_vals_by_fn, child_fn, capture)
-
-def get_depend_fns(fn: Callable, capture: bool) -> tuple[set[CheckpointFn], dict[Callable, list[Any]]]:
-  checkpoint_fns: set[CheckpointFn] = set()
-  captured_vals_by_fn: dict[Callable, list[Any]] = {}
-  append_fn_depends(checkpoint_fns, captured_vals_by_fn, fn, capture)
-  return checkpoint_fns, captured_vals_by_fn
-
-def get_function_hash(fn: Callable, capture: bool) -> tuple[str, list[Callable]]:
-  checkpoint_fns, captured_vals_by_fn = get_depend_fns(fn, capture)
-  checkpoint_fns = sorted(checkpoint_fns, key=lambda fn: unwrap_fn(fn).__qualname__)
-  checkpoint_hashes = [check.fn_hash for check in checkpoint_fns]
-  depend_fns, depend_captured_vals = transpose(sorted(captured_vals_by_fn.items(), key=lambda x: x[0].__qualname__), 2)
-  fn_bodies = list(map(get_fn_body, [fn] + depend_fns))
-  fn_hash = hashing.hash((fn_bodies, depend_captured_vals, checkpoint_hashes), "blake2b")
-  return fn_hash, checkpoint_fns + depend_fns

checkpointer-2.1.0.dist-info/RECORD

@@ -1,14 +0,0 @@
-checkpointer/__init__.py,sha256=t-dv0hIfgJHFx2M8tjCUMC9DlucPM8hvJOwGv86owUo,411
-checkpointer/checkpoint.py,sha256=NHY_63EzlY3X6eqbOBE-dIprMZ_-X_GRC-nhy6cI1QQ,4990
-checkpointer/function_body.py,sha256=DAq5fj1MMgb3az_Pfdxzqg6woJ6esFgvaqkkKqogJBY,4074
-checkpointer/print_checkpoint.py,sha256=21aeqgM9CMjNAJyScqFmXCWWfh3jBIn7o7i5zJkZGaA,1369
-checkpointer/types.py,sha256=SslunQTXxovFuGOR_VKfL7z5Vif9RD1PPx0J1FQdGLw,564
-checkpointer/utils.py,sha256=qT7pk3o6GkX-1ylfi6I-DJO5fdVOIHMuaEK8dTEAoVw,1465
-checkpointer/storages/__init__.py,sha256=G7JrOAyCGITd1wOz-u6_4RZVgxzxGLVLHPwBuW1sx1U,300
-checkpointer/storages/bcolz_storage.py,sha256=UoeREc3oS8skFClu9sULpgpqbIVcp3tVd8CeYfAe5yM,2220
-checkpointer/storages/memory_storage.py,sha256=RQ4WTVapxJGVPv1DNlb9VFTifxtyQy8YVo8fwaRLfdk,692
-checkpointer/storages/pickle_storage.py,sha256=nyrBWLXKnyzXgZIMwrpWUOAGRozpX3jL9pCyCV29e4E,787
-checkpointer-2.1.0.dist-info/METADATA,sha256=tdpLxisGi4Wx3gvs_W52t1SQqRwgDuke26pkfzHVuKY,9926
-checkpointer-2.1.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-checkpointer-2.1.0.dist-info/licenses/LICENSE,sha256=0cmUKqBotzbBcysIexd52AhjwbphhlGYiWbvg5l2QAU,1054
-checkpointer-2.1.0.dist-info/RECORD,,