checkpointer 2.7.1__py3-none-any.whl → 2.8.1__py3-none-any.whl

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

@@ -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.dist-info → checkpointer-2.8.1.dist-info}/RECORD

@@ -1,16 +1,16 @@
 checkpointer/__init__.py,sha256=HRLsQ24ZhxgmDcHchZ-hX6wA0NMCSedGA0NmCnUdS_c,832
-checkpointer/checkpoint.py,sha256=H1RQJIIuEmRcXd42Y-qUfyU3z0OGp3smY0LoKQN0IYU,6425
+checkpointer/checkpoint.py,sha256=FZSKQVIXj_Enja0253WDUauO8siUHwlzcr4frHMJzB0,6538
 checkpointer/fn_ident.py,sha256=SWaksNCTlskMom0ztqjECSRjZYPWXUA1p1ZCb-9tWo0,4297
 checkpointer/object_hash.py,sha256=cxuWRDrg4F9wC18aC12zOZYOPv3bk2Qf6tZ0_WgAb6Y,7484
 checkpointer/print_checkpoint.py,sha256=aJCeWMRJiIR3KpyPk_UOKTaD906kArGrmLGQ3LqcVgo,1369
-checkpointer/test_checkpointer.py,sha256=VdINXiaA_BoDdVYEB73ctfQ42fw3EDoYa9vYacoB13A,3768
-checkpointer/utils.py,sha256=E1AV96NTh3tuVmgPrr0JSKZaokw-Jely5Y6-NjlMCp8,3141
+checkpointer/test_checkpointer.py,sha256=DM5f1Ci4z7MhDnxK-2Z-V3a3ntzBJqORQvCFROAf2SI,3807
+checkpointer/utils.py,sha256=2yk1ksKszXofwnSqBrNCFRSR4C3YLPEAdHUFf-cviRU,2755
 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=idh9sBMdWuyvS220oa_7bAUpc9Xo9v6Ud9aYKGWasUs,1593
 checkpointer/storages/storage.py,sha256=_m18Z8TKrdAbi6YYYQmuNOnhna4RB2sJDn1v3liaU3U,721
-checkpointer-2.7.1.dist-info/METADATA,sha256=wEdB7ZEnYVW3-bcwFBpZAQXbv76MNU1PvUpKMhW1Ids,10606
-checkpointer-2.7.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-checkpointer-2.7.1.dist-info/licenses/LICENSE,sha256=9xVsdtv_-uSyY9Xl9yujwAPm4-mjcCLeVy-ljwXEWbo,1059
-checkpointer-2.7.1.dist-info/RECORD,,
+checkpointer-2.8.1.dist-info/METADATA,sha256=z0SbrcyKvHTOgJ6oq0dshV25fBAsqGjMCl_uWmxSFRM,10600
+checkpointer-2.8.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+checkpointer-2.8.1.dist-info/licenses/LICENSE,sha256=9xVsdtv_-uSyY9Xl9yujwAPm4-mjcCLeVy-ljwXEWbo,1059
+checkpointer-2.8.1.dist-info/RECORD,,