checkpointer 2.14.4__tar.gz → 2.14.6__tar.gz

{checkpointer-2.14.4 → checkpointer-2.14.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.14.4
+Version: 2.14.6
 Summary: checkpointer adds code-aware caching to Python functions, maintaining correctness and speeding up execution as your code changes.
 Project-URL: Repository, https://github.com/Reddan/checkpointer.git
 Author: Hampus Hallman

{checkpointer-2.14.4 → checkpointer-2.14.6}/checkpointer/checkpoint.py

@@ -25,7 +25,7 @@ class CheckpointerOpts(TypedDict, total=False):
   directory: Path | str | None
   when: bool
   verbosity: Literal[0, 1, 2]
-  expiry: Callable[[datetime], bool] | timedelta | None
+  expiry: timedelta | Callable[[datetime], bool] | None
   capture: bool
   fn_hash_from: object
 
@@ -127,11 +127,8 @@ class FunctionIdent:
 
 class CachedFunction(Generic[Fn]):
   def __init__(self, checkpointer: Checkpointer, fn: Fn):
-    store_format = checkpointer.storage
-    Storage = STORAGE_MAP[store_format] if isinstance(store_format, str) else store_format
     update_wrapper(self, unwrap(fn))  # type: ignore
     self.ident = FunctionIdent(self, checkpointer, fn)
-    self.storage = Storage(self)
     self.bound = ()
 
   @overload
@@ -154,6 +151,12 @@ class CachedFunction(Generic[Fn]):
   def fn(self) -> Fn:
     return self.ident.fn  # type: ignore
 
+  @cached_property
+  def storage(self) -> Storage:
+    store_format = self.ident.checkpointer.storage
+    Storage = STORAGE_MAP[store_format] if isinstance(store_format, str) else store_format
+    return Storage(self)
+
   @property
   def cleanup(self):
     return self.storage.cleanup
@@ -260,6 +263,9 @@ class CachedFunction(Generic[Fn]):
   def set(self, value, *args, **kw):
     self.storage.store(self._get_call_hash(args, kw), value)
 
+  def set_awaitable(self: CachedFunction[Callable[P, Coro[R]]], value: R, *args: P.args, **kw: P.kwargs):
+    self.set(AwaitableValue(value), *args, **kw)
+
   def __repr__(self) -> str:
     initialized = "fn_hash" in self.ident.__dict__
     fn_hash = self.ident.fn_hash[:6] if initialized else "- uninitialized"

{checkpointer-2.14.4 → checkpointer-2.14.6}/checkpointer/storages/memory_storage.py

@@ -1,12 +1,23 @@
+from __future__ import annotations
 import gc
-from typing import Any
-from pathlib import Path
 from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from weakref import WeakSet
 from .storage import Storage
 
+if TYPE_CHECKING:
+  from ..checkpoint import CachedFunction
+
 item_map: dict[Path, dict[str, tuple[datetime, Any]]] = {}
+mem_stores: WeakSet[MemoryStorage] = WeakSet()
 
 class MemoryStorage(Storage):
+  def __init__(self, cached_fn: CachedFunction):
+    super().__init__(cached_fn)
+    self.cleanup()
+    mem_stores.add(self)
+
   def get_dict(self):
     return item_map.setdefault(self.fn_dir(), {})
 
@@ -45,9 +56,8 @@ class MemoryStorage(Storage):
 
 def cleanup_memory_storage():
   gc.collect()
-  memory_stores = [obj for obj in gc.get_objects() if isinstance(obj, MemoryStorage)]
-  storage_keys = {store.fn_dir() for store in memory_stores}
+  storage_keys = {store.fn_dir() for store in mem_stores}
   for key in item_map.keys() - storage_keys:
     del item_map[key]
-  for store in memory_stores:
-    store.cleanup(invalidated=False)
+  for store in mem_stores:
+    store.cleanup()

{checkpointer-2.14.4 → checkpointer-2.14.6}/checkpointer/storages/pickle_storage.py

@@ -53,15 +53,18 @@ class PickleStorage(Storage):
     if invalidated and fn_path.exists():
       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.cached_fn.__qualname__}")
+        for pkl_path in path.glob("**/*.pkl"):
+          pkl_path.unlink(missing_ok=True)
+      if old_dirs:
+        print(f"Removed {len(old_dirs)} invalidated directories for {self.cached_fn.__qualname__}")
     if expired and self.checkpointer.expiry:
       count = 0
       for pkl_path in fn_path.glob("**/*.pkl"):
         if self.expired_dt(filedate(pkl_path)):
           count += 1
           pkl_path.unlink(missing_ok=True)
-      print(f"Removed {count} expired checkpoints for {self.cached_fn.__qualname__}")
+      if count:
+        print(f"Removed {count} expired checkpoints for {self.cached_fn.__qualname__}")
     clear_directory(fn_path)
 
   def clear(self):

{checkpointer-2.14.4 → checkpointer-2.14.6}/checkpointer/storages/storage.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
-from typing import Any, TYPE_CHECKING
-from pathlib import Path
 from datetime import datetime, timedelta
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-  from ..checkpoint import Checkpointer, CachedFunction
+  from ..checkpoint import CachedFunction, Checkpointer
 
 class Storage:
   checkpointer: Checkpointer
@@ -22,7 +22,7 @@ class Storage:
     return self.checkpointer.directory / self.fn_id()
 
   def expired(self, call_hash: str) -> bool:
-    if not self.checkpointer.expiry:
+    if self.checkpointer.expiry is None:
       return False
     return self.expired_dt(self.checkpoint_date(call_hash))
 

{checkpointer-2.14.4 → checkpointer-2.14.6}/checkpointer/utils.py

@@ -121,15 +121,15 @@ def empty_dirs(path: Path) -> Iterable[Path]:
   for child in path.iterdir():
     nonempty_count += 1
     if child.is_dir():
-      for grand_child in empty_dirs(child):
-        yield grand_child
-        nonempty_count -= child == grand_child
+      for descendant in empty_dirs(child):
+        yield descendant
+        nonempty_count -= child == descendant
   if nonempty_count == 0:
     yield path
 
 def clear_directory(path: Path):
   if path.is_dir():
     for file in path.glob("**/.DS_Store"):
-      file.unlink()
+      file.unlink(missing_ok=True)
     for directory in empty_dirs(path):
       directory.rmdir()

{checkpointer-2.14.4 → checkpointer-2.14.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.14.4"
+version = "2.14.6"
 requires-python = ">=3.11"
 dependencies = []
 authors = [
@@ -36,7 +36,7 @@ Repository = "https://github.com/Reddan/checkpointer.git"
 [dependency-groups]
 dev = [
   "numpy>=2.2.1",
-  "omg>=1.3.6",
+  "omg>=1.3.9",
   "poethepoet>=0.30.0",
   "polars>=1.31.0",
   "pytest>=8.3.5",

{checkpointer-2.14.4 → checkpointer-2.14.6}/uv.lock

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.14.4"
+version = "2.14.6"
 source = { editable = "." }
 
 [package.dev-dependencies]
@@ -28,7 +28,7 @@ dev = [
 [package.metadata.requires-dev]
 dev = [
     { name = "numpy", specifier = ">=2.2.1" },
-    { name = "omg", specifier = ">=1.3.6" },
+    { name = "omg", specifier = ">=1.3.9" },
     { name = "poethepoet", specifier = ">=0.30.0" },
     { name = "polars", specifier = ">=1.31.0" },
     { name = "pytest", specifier = ">=8.3.5" },
@@ -340,14 +340,14 @@ wheels = [
 
 [[package]]
 name = "omg"
-version = "1.3.6"
+version = "1.3.9"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "watchdog" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/65/06/da0a3778b7ff8f1333ed7ddc0931ffff3c86ab5cb8bc4a96a1d0edb8671b/omg-1.3.6.tar.gz", hash = "sha256:465a51b7576fa31ef313e2b9a77d57f5d4816fb0a14dca0fc5c09ff471074fe6", size = 14268 }
+sdist = { url = "https://files.pythonhosted.org/packages/2a/a3/4a0fb81ac8265c688dbd8eddc4ac61d0621a850a1037090d4af88ae38f01/omg-1.3.9.tar.gz", hash = "sha256:4d2de31323a8b2bfd80e04e1670b8d3a5912cec041854bcae79a97fb10f0b7ea", size = 14683 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/dd/d2/87346e94dbecd3a65a09e2156c1adf30c162f31e69d0936343c3eff53e7a/omg-1.3.6-py3-none-any.whl", hash = "sha256:8e3ac99a18d5284ceef2ed98492d288d5f22ee2bb417591654a7d2433e196607", size = 7988 },
+    { url = "https://files.pythonhosted.org/packages/f8/44/0f768cb8c82a5ed4ae5a71246668be55351795cc4438045425bd9709099a/omg-1.3.9-py3-none-any.whl", hash = "sha256:585abf178a28749c1b80bbba154914a9cfeb688b3955883d72068e2726716c8a", size = 8121 },
 ]
 
 [[package]]