checkpointer 2.7.0__tar.gz → 2.8.0__tar.gz

{checkpointer-2.7.0 → checkpointer-2.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.7.0
+Version: 2.8.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
@@ -255,7 +255,7 @@ async def main():
     result2 = await async_compute_sum(3, 7)
     print(result2)  # Outputs 10
 
-    result3 = await async_compute_sum.get(3, 7)
+    result3 = async_compute_sum.get(3, 7)
     print(result3)  # Outputs 10
 
 asyncio.run(main())

{checkpointer-2.7.0 → checkpointer-2.8.0}/README.md

@@ -235,7 +235,7 @@ async def main():
     result2 = await async_compute_sum(3, 7)
     print(result2)  # Outputs 10
 
-    result3 = await async_compute_sum.get(3, 7)
+    result3 = async_compute_sum.get(3, 7)
     print(result3)  # Outputs 10
 
 asyncio.run(main())

{checkpointer-2.7.0 → checkpointer-2.8.0}/checkpointer/__init__.py

@@ -9,7 +9,8 @@ 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")
+tmp_checkpoint = Checkpointer(root_path=f"{tempfile.gettempdir()}/checkpoints")
+static_checkpoint = Checkpointer(fn_hash=ObjectHash())
 
 def cleanup_all(invalidated=True, expired=True):
   for obj in gc.get_objects():

{checkpointer-2.7.0 → checkpointer-2.8.0}/checkpointer/checkpoint.py

@@ -4,14 +4,16 @@ import re
 from datetime import datetime
 from functools import update_wrapper
 from pathlib import Path
-from typing import Any, Callable, Generic, Iterable, Literal, Type, TypedDict, TypeVar, Unpack, cast, overload
+from typing import Any, Awaitable, Callable, Generic, Iterable, Literal, ParamSpec, 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
+from .utils import AwaitableValue, unwrap_fn
 
 Fn = TypeVar("Fn", bound=Callable)
+P = ParamSpec("P")
+R = TypeVar("R")
 
 DEFAULT_DIR = Path.home() / ".cache/checkpoints"
 
@@ -70,7 +72,6 @@ class CheckpointFn(Generic[Fn]):
     deep_hashes = [child._set_ident().fn_hash_raw for child in iterate_checkpoint_fns(self)]
     self.fn_hash = str(params.fn_hash or ObjectHash().write_text(self.fn_hash_raw, *deep_hashes))
     self.fn_subdir = f"{fn_file}/{fn_name}/{self.fn_hash[:16]}"
-    self.is_async: bool = self.fn.is_async if isinstance(self.fn, CheckpointFn) else inspect.iscoroutinefunction(self.fn)
     self.storage = Storage(self)
     self.cleanup = self.storage.cleanup
 
@@ -96,7 +97,12 @@ class CheckpointFn(Generic[Fn]):
     call_hash = ObjectHash(hash_params, digest_size=16)
     return f"{self.fn_subdir}/{call_hash}"
 
-  async def _store_on_demand(self, args: tuple, kw: dict, rerun: bool):
+  async def _resolve_awaitable(self, checkpoint_path: Path, awaitable: Awaitable):
+    data = await awaitable
+    self.storage.store(checkpoint_path, AwaitableValue(data))
+    return data
+
+  def _store_on_demand(self, args: tuple, kw: dict, rerun: bool):
     params = self.checkpointer
     checkpoint_id = self.get_checkpoint_id(args, kw)
     checkpoint_path = params.root_path / checkpoint_id
@@ -107,10 +113,11 @@ class CheckpointFn(Generic[Fn]):
     if refresh:
       print_checkpoint(params.verbosity >= 1, "MEMORIZING", checkpoint_id, "blue")
       data = self.fn(*args, **kw)
-      if inspect.iscoroutine(data):
-        data = await data
-      self.storage.store(checkpoint_path, data)
-      return data
+      if inspect.isawaitable(data):
+        return self._resolve_awaitable(checkpoint_path, data)
+      else:
+        self.storage.store(checkpoint_path, data)
+        return data
 
     try:
       data = self.storage.load(checkpoint_path)
@@ -119,29 +126,33 @@ class CheckpointFn(Generic[Fn]):
     except (EOFError, FileNotFoundError):
       pass
     print_checkpoint(params.verbosity >= 1, "CORRUPTED", checkpoint_id, "yellow")
-    return await self._store_on_demand(args, kw, True)
+    return self._store_on_demand(args, kw, True)
 
   def _call(self, args: tuple, kw: dict, rerun=False):
     if not self.checkpointer.when:
       return self.fn(*args, **kw)
-    coroutine = self._store_on_demand(args, kw, rerun)
-    return coroutine if self.is_async else sync_resolve_coroutine(coroutine)
+    return self._store_on_demand(args, kw, rerun)
+
+  __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))
+
+  @overload
+  def get(self: Callable[P, Awaitable[R]], *args: P.args, **kw: P.kwargs) -> R: ...
+  @overload
+  def get(self: Callable[P, R], *args: P.args, **kw: P.kwargs) -> R: ...
 
-  def _get(self, args, kw) -> Any:
+  def get(self, *args, **kw):
     checkpoint_path = self.checkpointer.root_path / self.get_checkpoint_id(args, kw)
     try:
-      val = self.storage.load(checkpoint_path)
-      return resolved_awaitable(val) if self.is_async else val
+      data = self.storage.load(checkpoint_path)
+      return data.value if isinstance(data, AwaitableValue) else data
     except Exception as ex:
       raise CheckpointError("Could not load checkpoint") from ex
 
-  def exists(self, *args: tuple, **kw: dict) -> bool:
+  def exists(self: Callable[P, R], *args: P.args, **kw: P.kwargs) -> bool: # type: ignore
+    self = cast(CheckpointFn, self)
     return self.storage.exists(self.checkpointer.root_path / self.get_checkpoint_id(args, kw))
 
-  __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]}>"
 

{checkpointer-2.7.0 → checkpointer-2.8.0}/checkpointer/object_hash.py

@@ -56,6 +56,9 @@ class ObjectHash:
 
   __str__ = hexdigest
 
+  def __eq__(self, value: object) -> bool:
+    return isinstance(value, ObjectHash) and str(self) == str(value)
+
   def nested_hash(self, *objs: Any) -> str:
     return ObjectHash(iter=objs, tolerate_errors=self.tolerate_errors.value).hexdigest()
 

{checkpointer-2.7.0 → checkpointer-2.8.0}/checkpointer/test_checkpointer.py

@@ -6,7 +6,7 @@ from . import checkpoint
 from .checkpoint import CheckpointError
 from .utils import AttrDict
 
-def global_multiply(a, b):
+def global_multiply(a: int, b: int) -> int:
   return a * b
 
 @pytest.fixture(autouse=True)
@@ -27,15 +27,15 @@ def test_basic_caching():
 
 def test_cache_invalidation():
   @checkpoint
-  def multiply(a, b):
+  def multiply(a: int, b: int):
     return a * b
 
   @checkpoint
-  def helper(x):
+  def helper(x: int):
     return multiply(x + 1, 2)
 
   @checkpoint
-  def compute(a, b):
+  def compute(a: int, b: int):
     return helper(a) + helper(b)
 
   result1 = compute(3, 4)
@@ -46,7 +46,7 @@ def test_layered_caching():
 
   @checkpoint(format="memory")
   @dev_checkpoint
-  def expensive_function(x):
+  def expensive_function(x: int):
     return x ** 2
 
   assert expensive_function(4) == 16
@@ -79,9 +79,10 @@ async def test_async_caching():
     return x ** 2
 
   result1 = await async_square(3)
-  result2 = await async_square.get(3)
+  result2 = await async_square(3)
+  result3 = async_square.get(3)
 
-  assert result1 == result2 == 9
+  assert result1 == result2 == result3 == 9
 
 def test_force_recalculation():
   @checkpoint
@@ -95,7 +96,7 @@ def test_force_recalculation():
 def test_multi_layer_decorator():
   @checkpoint(format="memory")
   @checkpoint(format="pickle")
-  def add(a, b):
+  def add(a: int, b: int) -> int:
     return a + b
 
   assert add(2, 3) == 5
@@ -124,18 +125,18 @@ def test_capture():
   assert test_a.fn_hash != init_hash_a
 
 def test_depends():
-  def multiply_wrapper(a, b):
+  def multiply_wrapper(a: int, b: int) -> int:
     return global_multiply(a, b)
 
-  def helper(a, b):
+  def helper(a: int, b: int) -> int:
     return multiply_wrapper(a + 1, b + 1)
 
   @checkpoint
-  def test_a(a, b):
+  def test_a(a: int, b: int) -> int:
     return helper(a, b)
 
   @checkpoint
-  def test_b(a, b):
+  def test_b(a: int, b: int) -> int:
     return test_a(a, b) + multiply_wrapper(a, b)
 
   assert set(test_a.depends) == {test_a.fn, helper, multiply_wrapper, global_multiply}
@@ -143,17 +144,17 @@ def test_depends():
 
 def test_lazy_init():
   @checkpoint
-  def fn1(x):
+  def fn1(x: object) -> object:
     return fn2(x)
 
   @checkpoint
-  def fn2(x):
+  def fn2(x: object) -> object:
     return fn1(x)
 
-  assert type(object.__getattribute__(fn1, "_getattribute")) == MethodType
+  assert type(object.__getattribute__(fn1, "_getattribute")) is 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 type(object.__getattribute__(fn1, "_getattribute")) is MethodWrapperType
   assert set(fn1.depends) == {fn1.fn, fn2}
   assert set(fn2.depends) == {fn1, fn2.fn}

{checkpointer-2.7.0 → checkpointer-2.8.0}/checkpointer/utils.py

@@ -2,11 +2,10 @@ import inspect
 import tokenize
 from contextlib import contextmanager
 from io import StringIO
-from types import coroutine
-from typing import Any, Callable, Coroutine, Generator, Generic, Iterable, TypeVar, cast
+from typing import Any, Callable, Generic, Iterable, TypeVar, cast
 
 T = TypeVar("T")
-T_Callable = TypeVar("T_Callable", bound=Callable)
+Fn = TypeVar("Fn", bound=Callable)
 
 def distinct(seq: Iterable[T]) -> list[T]:
   return list(dict.fromkeys(seq))
@@ -33,28 +32,20 @@ def get_cell_contents(fn: Callable) -> Iterable[tuple[str, Any]]:
     except ValueError:
       pass
 
-def unwrap_fn(fn: T_Callable, checkpoint_fn=False) -> T_Callable:
+def unwrap_fn(fn: Fn, checkpoint_fn=False) -> Fn:
   from .checkpoint import CheckpointFn
   while True:
     if (checkpoint_fn and isinstance(fn, CheckpointFn)) or not hasattr(fn, "__wrapped__"):
-      return cast(T_Callable, fn)
+      return cast(Fn, fn)
     fn = getattr(fn, "__wrapped__")
 
-async def resolved_awaitable(value: T) -> T:
-  return value
-
-@coroutine
-def coroutine_as_generator(coroutine: Coroutine[None, None, T]) -> Generator[None, None, T]:
-  val = yield from coroutine
-  return val
+class AwaitableValue:
+  def __init__(self, value):
+    self.value = value
 
-def sync_resolve_coroutine(coroutine: Coroutine[None, None, T]) -> T:
-  gen = cast(Generator, coroutine_as_generator(coroutine))
-  try:
-    while True:
-      next(gen)
-  except StopIteration as ex:
-    return ex.value
+  def __await__(self):
+    yield
+    return self.value
 
 class AttrDict(dict):
   def __init__(self, *args, **kwargs):

{checkpointer-2.7.0 → checkpointer-2.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.7.0"
+version = "2.8.0"
 requires-python = ">=3.11"
 dependencies = []
 authors = [
@@ -23,8 +23,8 @@ dev = [
   "numpy>=2.2.1",
   "omg>=1.2.3",
   "poethepoet>=0.30.0",
-  "pytest>=8.3.3",
-  "pytest-asyncio>=0.24.0",
+  "pytest>=8.3.5",
+  "pytest-asyncio>=0.26.0",
   "relib>=1.0.8",
   "torch>=2.5.1",
 ]
@@ -39,3 +39,6 @@ build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
 packages = ["checkpointer", "checkpointer.storages"]
+
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope = "session"

{checkpointer-2.7.0 → checkpointer-2.8.0}/uv.lock

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.7.0"
+version = "2.8.0"
 source = { editable = "." }
 
 [package.dev-dependencies]
@@ -30,8 +30,8 @@ dev = [
     { name = "numpy", specifier = ">=2.2.1" },
     { name = "omg", specifier = ">=1.2.3" },
     { name = "poethepoet", specifier = ">=0.30.0" },
-    { name = "pytest", specifier = ">=8.3.3" },
-    { name = "pytest-asyncio", specifier = ">=0.24.0" },
+    { name = "pytest", specifier = ">=8.3.5" },
+    { name = "pytest-asyncio", specifier = ">=0.26.0" },
     { name = "relib", specifier = ">=1.0.8" },
     { name = "torch", specifier = ">=2.5.1" },
 ]
@@ -364,7 +364,7 @@ wheels = [
 
 [[package]]
 name = "pytest"
-version = "8.3.4"
+version = "8.3.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "colorama", marker = "sys_platform == 'win32'" },
@@ -372,21 +372,21 @@ dependencies = [
     { name = "packaging" },
     { name = "pluggy" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/05/35/30e0d83068951d90a01852cb1cef56e5d8a09d20c7f511634cc2f7e0372a/pytest-8.3.4.tar.gz", hash = "sha256:965370d062bce11e73868e0335abac31b4d3de0e82f4007408d242b4f8610761", size = 1445919 }
+sdist = { url = "https://files.pythonhosted.org/packages/ae/3c/c9d525a414d506893f0cd8a8d0de7706446213181570cdbd766691164e40/pytest-8.3.5.tar.gz", hash = "sha256:f4efe70cc14e511565ac476b57c279e12a855b11f48f212af1080ef2263d3845", size = 1450891 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/11/92/76a1c94d3afee238333bc0a42b82935dd8f9cf8ce9e336ff87ee14d9e1cf/pytest-8.3.4-py3-none-any.whl", hash = "sha256:50e16d954148559c9a74109af1eaf0c945ba2d8f30f0a3d3335edde19788b6f6", size = 343083 },
+    { url = "https://files.pythonhosted.org/packages/30/3d/64ad57c803f1fa1e963a7946b6e0fea4a70df53c1a7fed304586539c2bac/pytest-8.3.5-py3-none-any.whl", hash = "sha256:c69214aa47deac29fad6c2a4f590b9c4a9fdb16a403176fe154b79c0b4d4d820", size = 343634 },
 ]
 
 [[package]]
 name = "pytest-asyncio"
-version = "0.25.0"
+version = "0.26.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "pytest" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/94/18/82fcb4ee47d66d99f6cd1efc0b11b2a25029f303c599a5afda7c1bca4254/pytest_asyncio-0.25.0.tar.gz", hash = "sha256:8c0610303c9e0442a5db8604505fc0f545456ba1528824842b37b4a626cbf609", size = 53298 }
+sdist = { url = "https://files.pythonhosted.org/packages/8e/c4/453c52c659521066969523e87d85d54139bbd17b78f09532fb8eb8cdb58e/pytest_asyncio-0.26.0.tar.gz", hash = "sha256:c4df2a697648241ff39e7f0e4a73050b03f123f760673956cf0d72a4990e312f", size = 54156 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/88/56/2ee0cab25c11d4e38738a2a98c645a8f002e2ecf7b5ed774c70d53b92bb1/pytest_asyncio-0.25.0-py3-none-any.whl", hash = "sha256:db5432d18eac6b7e28b46dcd9b69921b55c3b1086e85febfe04e70b18d9e81b3", size = 19245 },
+    { url = "https://files.pythonhosted.org/packages/20/7f/338843f449ace853647ace35870874f69a764d251872ed1b4de9f234822c/pytest_asyncio-0.26.0-py3-none-any.whl", hash = "sha256:7b51ed894f4fbea1340262bdae5135797ebbe21d8638978e35d31c6d19f72fb0", size = 19694 },
 ]
 
 [[package]]