checkpointer 2.14.4__py3-none-any.whl → 2.14.6__py3-none-any.whl

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

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

@@ -1,18 +1,18 @@
 checkpointer/__init__.py,sha256=lprhgo4BcVbr4RczWdkNIWroVX_9M-akZcMbxOAWHf0,981
-checkpointer/checkpoint.py,sha256=N3MtNF4Oq3FV69LCfIuyTfNxrcnw9-dGzX312eYigVQ,10454
+checkpointer/checkpoint.py,sha256=BG2LkBbnCxWAzSTCNdYFJElef7LIFu29MvZqRp0RFrc,10665
 checkpointer/fn_ident.py,sha256=Z2HtkoG8n0ddqNzjejnKW6Lo-ID2unKQZa1tOpRGujs,6911
 checkpointer/fn_string.py,sha256=YldjAU91cwiZNwuE_AQN_DMiQCPAhGiiC3lPi0uvM0g,2579
 checkpointer/import_mappings.py,sha256=ESqWvZTzYAmaVnJ6NulUvn3_8CInOOPmEKUXO2WD_WA,1794
 checkpointer/object_hash.py,sha256=_VgyTEoe3H6268KBAOFZxSgWHb6otgG-3NzhHyehxe4,8595
 checkpointer/print_checkpoint.py,sha256=uUQ493fJCaB4nhp4Ox60govSCiBTIPbBX15zt2QiRGo,1356
 checkpointer/types.py,sha256=GFqbGACdDxzQX3bb2LmF9UxQVWOEisGvdtobnqCBAOA,1129
-checkpointer/utils.py,sha256=BAv_qvcanbYqgr7cCOCKNiG0_rGbNbLomBbpvoYlltc,3553
+checkpointer/utils.py,sha256=WGfVeHeZua010YXMuDPwAYfyQyB6k2bm0n0nbZdJ28Q,3565
 checkpointer/storages/__init__.py,sha256=p-r4YrPXn505_S3qLrSXHSlsEtb13w_DFnCt9IiUomk,296
-checkpointer/storages/memory_storage.py,sha256=VmU7uDscuibqpxIb0F3FRUXncVEbAOrLv-VaquP8Bsk,1598
-checkpointer/storages/pickle_storage.py,sha256=WCbWGdDH046_GrlnttcqqGZ55wGOz2Rb7Gu0slI448g,2203
-checkpointer/storages/storage.py,sha256=JVxdq1DMhbq83bvflvIuW7okE1CQCa7poQPsj3x0ACg,1366
-checkpointer-2.14.4.dist-info/METADATA,sha256=AoUxO7FgKWaEIeQD7sOIeFN5d1FkxzdMMSvMn8j23WU,11215
-checkpointer-2.14.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-checkpointer-2.14.4.dist-info/licenses/ATTRIBUTION.md,sha256=WF6L7-sD4s9t9ytVJOhjhpDoZ6TrWpqE3_bMdDIeJxI,1078
-checkpointer-2.14.4.dist-info/licenses/LICENSE,sha256=drXs6vIb7uW49r70UuMz2A1VtOCl626kiTbcmrar1Xo,1072
-checkpointer-2.14.4.dist-info/RECORD,,
+checkpointer/storages/memory_storage.py,sha256=M5O3gRMHBiG3CgMcFdA6-iDGvfafKqTi9nN3u8sETCg,1801
+checkpointer/storages/pickle_storage.py,sha256=L6E7cugbn-XGMMlcHCFBJ-rBLbmtzbiPwRrlgfRd7yc,2304
+checkpointer/storages/storage.py,sha256=5v_9Mp8__oqPnyrGsUEiT7jU4eMk9L4ORHDZeQw-TU4,1370
+checkpointer-2.14.6.dist-info/METADATA,sha256=q-QYZcNIAkQ-3AcnIuKaXtifxwlIqYPU3vy86gTFGBQ,11215
+checkpointer-2.14.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+checkpointer-2.14.6.dist-info/licenses/ATTRIBUTION.md,sha256=WF6L7-sD4s9t9ytVJOhjhpDoZ6TrWpqE3_bMdDIeJxI,1078
+checkpointer-2.14.6.dist-info/licenses/LICENSE,sha256=drXs6vIb7uW49r70UuMz2A1VtOCl626kiTbcmrar1Xo,1072
+checkpointer-2.14.6.dist-info/RECORD,,