checkpointer 2.11.2__py3-none-any.whl → 2.13.0__py3-none-any.whl

checkpointer/__init__.py

@@ -1,15 +1,15 @@
 import gc
 import tempfile
 from typing import Callable
-from .checkpoint import CachedFunction, Checkpointer, CheckpointError
+from .checkpoint import CachedFunction, Checkpointer, CheckpointError, FunctionIdent
 from .object_hash import ObjectHash
 from .storages import MemoryStorage, PickleStorage, Storage
-from .types import AwaitableValue, HashBy, NoHash
+from .types import AwaitableValue, Captured, CapturedOnce, CaptureMe, CaptureMeOnce, 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")
+memory_checkpoint = Checkpointer(storage="memory", verbosity=0)
+tmp_checkpoint = Checkpointer(directory=f"{tempfile.gettempdir()}/checkpoints")
 static_checkpoint = Checkpointer(fn_hash_from=())
 
 def cleanup_all(invalidated=True, expired=True):

checkpointer/checkpoint.py

@@ -2,30 +2,27 @@ 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 inspect import Parameter, iscoroutine, signature, unwrap
+from itertools import chain
 from pathlib import Path
 from typing import (
-  Annotated, Callable, Concatenate, Coroutine, Generic,
-  Iterable, Literal, Self, Type, TypedDict,
-  Unpack, cast, get_args, get_origin, overload,
+  Callable, Concatenate, Coroutine, Generic, Iterable,
+  Literal, Self, Type, TypedDict, Unpack, cast, overload,
 )
-from .fn_ident import RawFunctionIdent, get_fn_ident
+from .fn_ident import Capturable, 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, C, Coro, Fn, HashBy, P, R
-from .utils import unwrap_fn
+from .storages import STORAGE_MAP, Storage, StorageType
+from .types import AwaitableValue, C, Coro, Fn, P, R, hash_by_from_annotation
 
 DEFAULT_DIR = Path.home() / ".cache/checkpoints"
 
-empty_set = cast(set, frozenset())
-
 class CheckpointError(Exception):
   pass
 
 class CheckpointerOpts(TypedDict, total=False):
-  format: Type[Storage] | Literal["pickle", "memory", "bcolz"]
-  root_path: Path | str | None
+  storage: Type[Storage] | StorageType
+  directory: Path | str | None
   when: bool
   verbosity: Literal[0, 1, 2]
   should_expire: Callable[[datetime], bool] | None
@@ -34,8 +31,8 @@ class CheckpointerOpts(TypedDict, total=False):
 
 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.storage = opts.get("storage", "pickle")
+    self.directory = Path(opts.get("directory", DEFAULT_DIR) or ".")
     self.when = opts.get("when", True)
     self.verbosity = opts.get("verbosity", 1)
     self.should_expire = opts.get("should_expire")
@@ -59,119 +56,164 @@ class FunctionIdent:
   Separated from CachedFunction to prevent hash desynchronization
   among bound instances when `.reinit()` is called.
   """
-  def __init__(self, cached_fn: CachedFunction):
-    self.__dict__.clear()
+  __slots__ = (
+    "checkpointer", "cached_fn", "fn", "fn_dir", "pos_names",
+    "arg_names", "default_args", "hash_by_map", "__dict__",
+  )
+
+  def __init__(self, cached_fn: CachedFunction, checkpointer: Checkpointer, fn: Callable):
+    wrapped = unwrap(fn)
+    fn_file = Path(wrapped.__code__.co_filename).name
+    fn_name = re.sub(r"[^\w.]", "", wrapped.__qualname__)
+    params = list(signature(wrapped).parameters.values())
+    pos_param_types = (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD)
+    named_param_types = (Parameter.KEYWORD_ONLY,) + pos_param_types
+    name_by_kind = {Parameter.VAR_POSITIONAL: b"*", Parameter.VAR_KEYWORD: b"**"}
+    self.checkpointer = checkpointer
     self.cached_fn = cached_fn
+    self.fn = fn
+    self.fn_dir = f"{fn_file}/{fn_name}"
+    self.pos_names = [param.name for param in params if param.kind in pos_param_types]
+    self.arg_names = {param.name for param in params if param.kind in named_param_types}
+    self.default_args = {param.name: param.default for param in params if param.default is not Parameter.empty}
+    self.hash_by_map = {
+      name_by_kind.get(param.kind, param.name): hash_by
+      for param in params
+      if (hash_by := hash_by_from_annotation(param.annotation))
+    }
+
+  def reset(self):
+    self.__dict__.clear()
+
+  def is_static(self) -> bool:
+    return self.checkpointer.fn_hash_from is not None
 
   @cached_property
   def raw_ident(self) -> RawFunctionIdent:
-    return get_fn_ident(unwrap_fn(self.cached_fn.fn), self.cached_fn.checkpointer.capture)
+    return get_fn_ident(unwrap(self.fn), self.checkpointer.capture)
 
   @cached_property
   def fn_hash(self) -> str:
-    if (hash_from := self.cached_fn.checkpointer.fn_hash_from) is not None:
-      return str(ObjectHash(hash_from, digest_size=16))
-    deep_hashes = [depend.ident.raw_ident.fn_hash for depend in self.cached_fn.deep_depends()]
+    if self.is_static():
+      return str(ObjectHash(self.checkpointer.fn_hash_from, digest_size=16))
+    depends = self.deep_idents(past_static=False)
+    deep_hashes = [d.fn_hash if d.is_static() else d.raw_ident.fn_hash for d in 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 capturables(self) -> list[Capturable]:
+    return sorted({
+      capturable.key: capturable
+      for depend in self.deep_idents()
+      for capturable in depend.raw_ident.capturables
+    }.values())
+
+  def deep_depends(self, past_static=True, visited: set[Callable] = set()) -> Iterable[Callable]:
+    if self.cached_fn not in visited:
+      yield self.cached_fn
+      visited = visited or set()
+      visited.add(self.cached_fn)
+      stop = not past_static and self.is_static()
+      depends = [] if stop else self.raw_ident.depends
+      for depend in depends:
+        if isinstance(depend, CachedFunction):
+          yield from depend.ident.deep_depends(past_static, visited)
+        elif depend not in visited:
+          yield depend
+          visited.add(depend)
 
-  def reset(self):
-    self.__init__(self.cached_fn)
+  def deep_idents(self, past_static=True) -> Iterable[FunctionIdent]:
+    return (fn.ident for fn in self.deep_depends(past_static) if isinstance(fn, CachedFunction))
 
 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}"
+    store_format = checkpointer.storage
+    Storage = STORAGE_MAP[store_format] if isinstance(store_format, str) else store_format
+    update_wrapper(cast(Callable, self), unwrap(fn))
+    self.ident = FunctionIdent(self, checkpointer, fn)
     self.storage = Storage(self)
-    self.cleanup = self.storage.cleanup
     self.bound = ()
-    self.attrname: str | None = None
-
-    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)
-
-  def __set_name__(self, _, name: str):
-    assert self.attrname is None
-    self.attrname = name
 
   @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: 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
-    assert self.attrname is not None
     bound_fn = object.__new__(CachedFunction)
     bound_fn.__dict__ |= self.__dict__
     bound_fn.bound = (instance,)
-    if hasattr(instance, "__dict__"):
-      setattr(instance, self.attrname, bound_fn)
     return bound_fn
 
   @property
-  def depends(self) -> list[Callable]:
-    return self.ident.raw_ident.depends
+  def fn(self) -> Fn:
+    return cast(Fn, self.ident.fn)
+
+  @property
+  def cleanup(self):
+    return self.storage.cleanup
 
   def reinit(self, recursive=False) -> CachedFunction[Fn]:
-    depend_idents = [depend.ident for depend in self.deep_depends()] if recursive else [self.ident]
+    depend_idents = list(self.ident.deep_idents()) if recursive else [self.ident]
     for ident in depend_idents: ident.reset()
     for ident in depend_idents: ident.fn_hash
     return self
 
-  def get_call_hash(self, args: tuple, kw: dict[str, object]) -> str:
+  def _get_call_hash(self, args: tuple, kw: dict[str, object]) -> str:
+    ident = self.ident
     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):
+    pos_args = args[len(ident.pos_names):]
+    named_pos_args = dict(zip(ident.pos_names, args))
+    named_args = {**ident.default_args, **named_pos_args, **kw}
+    for key, hash_by in ident.hash_by_map.items():
+      if isinstance(key, str):
+        named_args[key] = hash_by(named_args[key])
+      elif key == b"*":
+        pos_args = map(hash_by, pos_args)
+      elif key == b"**":
+        for key in kw.keys() - ident.arg_names:
+          named_args[key] = hash_by(named_args[key])
+    named_args_iter = chain.from_iterable(sorted(named_args.items()))
+    captured = chain.from_iterable(capturable.capture() for capturable in ident.capturables)
+    call_hash = ObjectHash(digest_size=16) \
+      .update(iter=named_args_iter, header="NAMED") \
+      .update(iter=pos_args, header="POS") \
+      .update(iter=captured, header="CAPTURED")
+    return str(call_hash)
+
+  def get_call_hash(self: CachedFunction[Callable[P, R]], *args: P.args, **kw: P.kwargs) -> str:
+    return self._get_call_hash(args, kw)
+
+  async def _store_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
+    params = self.ident.checkpointer
+    storage = self.storage
     if not params.when:
       return self.fn(*full_args, **kw)
 
-    call_hash = self.get_call_hash(args, kw)
-    call_id = f"{self.storage.fn_id()}/{call_hash}"
+    call_hash = self._get_call_hash(args, kw)
+    call_id = f"{storage.fn_id()}/{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)))
+      or not storage.exists(call_hash) \
+      or (params.should_expire and params.should_expire(storage.checkpoint_date(call_hash)))
 
     if refresh:
       print_checkpoint(params.verbosity >= 1, "MEMORIZING", call_id, "blue")
       data = self.fn(*full_args, **kw)
       if iscoroutine(data):
-        return self._resolve_coroutine(call_hash, data)
-      return self.storage.store(call_hash, data)
+        return self._store_coroutine(call_hash, data)
+      return storage.store(call_hash, data)
 
     try:
-      data = self.storage.load(call_hash)
+      data = storage.load(call_hash)
       print_checkpoint(params.verbosity >= 2, "REMEMBERED", call_id, "green")
       return data
     except (EOFError, FileNotFoundError):
@@ -186,17 +228,17 @@ class CachedFunction(Generic[Fn]):
     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))
+    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))
+    self.storage.delete(self._get_call_hash(args, kw))
 
   @overload
   def get(self: Callable[P, Coro[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)
+    call_hash = self._get_call_hash(args, kw)
     try:
       data = self.storage.load(call_hash)
       return data.value if isinstance(data, AwaitableValue) else data
@@ -208,33 +250,9 @@ class CachedFunction(Generic[Fn]):
   @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)
+    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] = empty_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 {}
+    initialized = "fn_hash" in self.ident.__dict__
+    fn_hash = self.ident.fn_hash[:6] if initialized else "- uninitialized"
+    return f"<CachedFunction {self.fn.__name__} {fn_hash}>"

checkpointer/fn_ident.py

@@ -1,103 +1,161 @@
 import dis
-import inspect
-from itertools import takewhile
-from pathlib import Path
-from types import CodeType, FunctionType, MethodType
-from typing import Callable, Iterable, NamedTuple, Type, TypeGuard
+from inspect import Parameter, getmodule, signature, unwrap
+from types import CodeType, MethodType, ModuleType
+from typing import Annotated, Callable, Iterable, NamedTuple, Type, get_args, get_origin
+from .fn_string import get_fn_aststr
+from .import_mappings import resolve_annotation
 from .object_hash import ObjectHash
-from .utils import AttrDict, distinct, get_cell_contents, iterate_and_upcoming, unwrap_fn
+from .types import hash_by_from_annotation, is_capture_me, is_capture_me_once, to_none
+from .utils import (
+  cwd, distinct, get_at, get_cell_contents,
+  get_file, is_class, is_user_fn, seekable, takewhile,
+)
 
-cwd = Path.cwd().resolve()
+AttrPath = tuple[str, ...]
+CapturableByFn = dict[Callable, list["Capturable"]]
 
 class RawFunctionIdent(NamedTuple):
   fn_hash: str
-  captured_hash: str
   depends: list[Callable]
+  capturables: set["Capturable"]
 
-def is_class(obj) -> TypeGuard[Type]:
-  # isinstance works too, but needlessly triggers _lazyinit()
-  return issubclass(type(obj), type)
+class Capturable(NamedTuple):
+  key: str
+  module: ModuleType
+  attr_path: AttrPath
+  hash_by: Callable | None
+  hash: str | None = None
 
-def extract_classvars(code: CodeType, scope_vars: AttrDict) -> dict[str, dict[str, Type]]:
-  attr_path: tuple[str, ...] = ()
+  def capture(self) -> tuple[str, object]:
+    if obj := self.hash:
+      return self.key, obj
+    obj = get_at(self.module, *self.attr_path)
+    obj = self.hash_by(obj) if self.hash_by else obj
+    return self.key, obj
+
+  @staticmethod
+  def new(module: ModuleType, attr_path: AttrPath, hash_by: Callable | None, capture_once: bool) -> "Capturable":
+    file = str(get_file(module).relative_to(cwd))
+    key = file + "/" + ".".join(attr_path)
+    cap = Capturable(key, module, attr_path, hash_by)
+    if not capture_once:
+      return cap
+    obj_hash = str(ObjectHash(cap.capture()[1]))
+    return Capturable(key, module, attr_path, None, obj_hash)
+
+def extract_classvars(code: CodeType, scope_vars: dict) -> dict[str, dict[str, Type]]:
+  attr_path = AttrPath(())
   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 = ()
+  instructs = seekable(dis.get_instructions(code))
+  for instruct in instructs:
+    if instruct.opname in scope_vars and not attr_path:
+      attrs = takewhile((x.opname == "LOAD_ATTR", x.argval) for x in instructs)
+      attr_path = AttrPath((instruct.opname, instruct.argval, *attrs))
+      instructs.step(-1)
+    elif instruct.opname == "CALL":
+      obj = get_at(scope_vars, *attr_path)
+      attr_path = AttrPath(())
       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
+    elif instruct.opname in ("STORE_FAST", "STORE_DEREF") and scope_obj:
+      load_key = instruct.opname.replace("STORE", "LOAD")
+      classvars.setdefault(load_key, {})[instruct.argval] = scope_obj
       scope_obj = None
   return classvars
 
-def extract_scope_values(code: CodeType, scope_vars: AttrDict) -> Iterable[tuple[tuple[str, ...], object]]:
+def extract_scope_values(code: CodeType, scope_vars: dict) -> Iterable[tuple[AttrPath, 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)):
-    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
+  scope_vars = {**scope_vars, **{k: {**scope_vars[k], **v} for k, v in classvars.items()}}
+  instructs = seekable(dis.get_instructions(code))
+  for instruct in instructs:
+    if instruct.opname in scope_vars:
+      attrs = takewhile((x.opname in ("LOAD_ATTR", "LOAD_METHOD"), x.argval) for x in instructs)
+      attr_path = AttrPath((instruct.opname, instruct.argval, *attrs))
+      parent_path = attr_path[:-1]
+      instructs.step(-1)
+      obj = get_at(scope_vars, *attr_path)
+      if obj is not None:
+        yield attr_path, obj
+      if callable(obj) and parent_path[1:]:
+        parent_obj = get_at(scope_vars, *parent_path)
+        yield parent_path, parent_obj
   for const in code.co_consts:
     if isinstance(const, CodeType):
-      yield from extract_scope_values(const, scope_vars)
+      next_deref = {**scope_vars["LOAD_DEREF"], **scope_vars["LOAD_FAST"]}
+      next_scope_vars = {**scope_vars, "LOAD_FAST": {}, "LOAD_DEREF": next_deref}
+      yield from extract_scope_values(const, next_scope_vars)
+
+def resolve_class_annotations(anno: object) -> Type | None:
+  if anno in (None, Annotated):
+    return None
+  elif is_class(anno):
+    return anno
+  elif get_origin(anno) is Annotated:
+    return resolve_class_annotations(next(iter(get_args(anno)), None))
+  return resolve_class_annotations(get_origin(anno))
 
-def get_self_value(fn: Callable) -> type | object | None:
+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)
+  parts = fn.__qualname__.split(".")[:-1]
+  cls = parts and get_at(fn.__globals__, *parts)
   if is_class(cls):
     return cls
 
-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 {}),
-    "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 get_capturables(fn: Callable, capture: bool, captured_vars: dict[AttrPath, object]) -> Iterable[Capturable]:
+  module = getmodule(fn)
+  if not module or not is_user_fn(fn):
+    return
+  for (instruct, *attr_path), obj in captured_vars.items():
+    attr_path = AttrPath(attr_path)
+    if instruct == "LOAD_GLOBAL" and not callable(obj) and not isinstance(obj, ModuleType):
+      anno = resolve_annotation(module, ".".join(attr_path))
+      if capture or is_capture_me(anno) or is_capture_me_once(anno):
+        hash_by = hash_by_from_annotation(anno)
+        if hash_by is not to_none:
+          yield Capturable.new(module, attr_path, hash_by, is_capture_me_once(anno))
 
-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_fn_captures(fn: Callable, capture: bool) -> tuple[list[Callable], list[Capturable]]:
+  scope_vars_signature: dict[str, Type | object] = {
+    param.name: class_anno
+    for param in signature(fn).parameters.values()
+    if param.annotation is not Parameter.empty
+    if param.kind not in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD)
+    if (class_anno := resolve_class_annotations(param.annotation))
+  }
+  if self_obj := get_self_value(fn):
+    scope_vars_signature["self"] = self_obj
+  scope_vars = {
+    "LOAD_FAST": scope_vars_signature,
+    "LOAD_DEREF": dict(get_cell_contents(fn)),
+    "LOAD_GLOBAL": fn.__globals__,
+  }
+  captured_vars = dict(extract_scope_values(fn.__code__, scope_vars))
+  captured_callables = [obj for obj in captured_vars.values() if callable(obj)]
+  capturables = list(get_capturables(fn, capture, captured_vars))
+  return captured_callables, capturables
 
-def get_depend_fns(fn: Callable, captured_vals_by_fn: dict[Callable, list[object]] = {}) -> dict[Callable, list[object]]:
+def get_depend_fns(fn: Callable, capture: bool, capturable_by_fn: CapturableByFn = {}) -> CapturableByFn:
   from .checkpoint import CachedFunction
-  captured_vals = get_fn_captured_vals(fn)
-  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, captured_vals_by_fn)
-  return captured_vals_by_fn
+  captured_callables, capturables = get_fn_captures(fn, capture)
+  capturable_by_fn = capturable_by_fn or {}
+  capturable_by_fn[fn] = capturables
+  for depend_fn in captured_callables:
+    depend_fn = unwrap(depend_fn, stop=lambda f: isinstance(f, CachedFunction))
+    if isinstance(depend_fn, CachedFunction):
+      capturable_by_fn[depend_fn.ident.cached_fn] = []
+    elif depend_fn not in capturable_by_fn and is_user_fn(depend_fn):
+      get_depend_fns(depend_fn, capture, capturable_by_fn)
+  return capturable_by_fn
 
 def get_fn_ident(fn: Callable, capture: bool) -> RawFunctionIdent:
   from .checkpoint import CachedFunction
-  captured_vals_by_fn = get_depend_fns(fn)
-  depend_captured_vals = list(captured_vals_by_fn.values()) * capture
-  depends = captured_vals_by_fn.keys()
+  capturable_by_fn = get_depend_fns(fn, capture)
+  capturables = {capt for capts in capturable_by_fn.values() for capt in capts}
+  depends = capturable_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)]
-  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)
+  depend_callables = [fn for fn in depends if not isinstance(fn, CachedFunction)]
+  assert fn == depend_callables[0]
+  fn_hash = str(ObjectHash(iter=map(get_fn_aststr, depend_callables)))
+  return RawFunctionIdent(fn_hash, depends, capturables)