checkpointer 2.7.1__tar.gz → 2.8.1__tar.gz

{checkpointer-2.7.1 → checkpointer-2.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.7.1
+Version: 2.8.1
 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.1 → checkpointer-2.8.1}/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.1 → checkpointer-2.8.1}/checkpointer/checkpoint.py

@@ -4,12 +4,12 @@ import re
 from datetime import datetime
 from functools import update_wrapper
 from pathlib import Path
-from typing import Any, Callable, Generic, Iterable, Literal, ParamSpec, 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")
@@ -70,9 +70,8 @@ class CheckpointFn(Generic[Fn]):
     update_wrapper(cast(Callable, self), wrapped)
     Storage = STORAGE_MAP[params.format] if isinstance(params.format, str) else params.format
     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.fn_hash = str(params.fn_hash or ObjectHash(digest_size=16).write_text(self.fn_hash_raw, *deep_hashes))
+    self.fn_subdir = f"{fn_file}/{fn_name}/{self.fn_hash[:32]}"
     self.storage = Storage(self)
     self.cleanup = self.storage.cleanup
 
@@ -98,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
@@ -109,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)
@@ -121,20 +126,26 @@ 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)
 
-  def get(self: Callable[P, R], *args: P.args, **kw: P.kwargs) -> R: # type: ignore
-    self = cast(CheckpointFn, self)
+  __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):
     checkpoint_path = self.checkpointer.root_path / self.get_checkpoint_id(args, kw)
     try:
-      val = self.storage.load(checkpoint_path)
-      return cast(R, 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
 
@@ -142,9 +153,6 @@ class CheckpointFn(Generic[Fn]):
     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))
-
   def __repr__(self) -> str:
     return f"<CheckpointFn {self.fn.__name__} {self.fn_hash[:6]}>"
 

{checkpointer-2.7.1 → checkpointer-2.8.1}/checkpointer/test_checkpointer.py

@@ -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

{checkpointer-2.7.1 → checkpointer-2.8.1}/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.1 → checkpointer-2.8.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.7.1"
+version = "2.8.1"
 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.1 → checkpointer-2.8.1}/uv.lock

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.7.1"
+version = "2.8.1"
 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]]