checkpointer 2.14.3__py3-none-any.whl → 2.14.5__py3-none-any.whl

checkpointer/__init__.py

@@ -5,6 +5,7 @@ from typing import Callable
 from .checkpoint import CachedFunction, Checkpointer, CheckpointError, FunctionIdent
 from .object_hash import ObjectHash
 from .storages import MemoryStorage, PickleStorage, Storage
+from .storages.memory_storage import cleanup_memory_storage
 from .types import AwaitableValue, Captured, CapturedOnce, CaptureMe, CaptureMeOnce, HashBy, NoHash
 
 checkpoint = Checkpointer()

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

checkpointer/storages/memory_storage.py

@@ -1,11 +1,23 @@
-from typing import Any
-from pathlib import Path
+from __future__ import annotations
+import gc
 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(), {})
 
@@ -41,3 +53,11 @@ class MemoryStorage(Storage):
     for key in list(item_map.keys()):
       if key.parent == fn_path:
         del item_map[key]
+
+def cleanup_memory_storage():
+  gc.collect()
+  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 mem_stores:
+    store.cleanup()

checkpointer/storages/pickle_storage.py

@@ -54,14 +54,16 @@ class PickleStorage(Storage):
       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__}")
+      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

{checkpointer-2.14.3.dist-info → checkpointer-2.14.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.14.3
+Version: 2.14.5
 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.3.dist-info → checkpointer-2.14.5.dist-info}/RECORD

@@ -1,5 +1,5 @@
-checkpointer/__init__.py,sha256=mT706hd-BoTpQXYH_MiqreplONVDvjCWjwwGQ7oRe_8,921
-checkpointer/checkpoint.py,sha256=N3MtNF4Oq3FV69LCfIuyTfNxrcnw9-dGzX312eYigVQ,10454
+checkpointer/__init__.py,sha256=lprhgo4BcVbr4RczWdkNIWroVX_9M-akZcMbxOAWHf0,981
+checkpointer/checkpoint.py,sha256=Vw0uKPBtp01Rit-ydE28TrMO3t3HQ2dkT0uJT2gtyms,10509
 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
@@ -8,11 +8,11 @@ checkpointer/print_checkpoint.py,sha256=uUQ493fJCaB4nhp4Ox60govSCiBTIPbBX15zt2Qi
 checkpointer/types.py,sha256=GFqbGACdDxzQX3bb2LmF9UxQVWOEisGvdtobnqCBAOA,1129
 checkpointer/utils.py,sha256=BAv_qvcanbYqgr7cCOCKNiG0_rGbNbLomBbpvoYlltc,3553
 checkpointer/storages/__init__.py,sha256=p-r4YrPXn505_S3qLrSXHSlsEtb13w_DFnCt9IiUomk,296
-checkpointer/storages/memory_storage.py,sha256=1YHZU-WSqvLA_0og53UZY141ziDgnaBub2mDjDaeAj8,1261
-checkpointer/storages/pickle_storage.py,sha256=WCbWGdDH046_GrlnttcqqGZ55wGOz2Rb7Gu0slI448g,2203
-checkpointer/storages/storage.py,sha256=JVxdq1DMhbq83bvflvIuW7okE1CQCa7poQPsj3x0ACg,1366
-checkpointer-2.14.3.dist-info/METADATA,sha256=YehES3ivV7IqJfs3jKVNCfrgRZwdEySFJrwqem8MVPo,11215
-checkpointer-2.14.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-checkpointer-2.14.3.dist-info/licenses/ATTRIBUTION.md,sha256=WF6L7-sD4s9t9ytVJOhjhpDoZ6TrWpqE3_bMdDIeJxI,1078
-checkpointer-2.14.3.dist-info/licenses/LICENSE,sha256=drXs6vIb7uW49r70UuMz2A1VtOCl626kiTbcmrar1Xo,1072
-checkpointer-2.14.3.dist-info/RECORD,,
+checkpointer/storages/memory_storage.py,sha256=M5O3gRMHBiG3CgMcFdA6-iDGvfafKqTi9nN3u8sETCg,1801
+checkpointer/storages/pickle_storage.py,sha256=3_it1YQL6y2WQwnOHT_UPZ1BQkVJZvuR9txcuiuPwYQ,2242
+checkpointer/storages/storage.py,sha256=0dwxQvg0HXlgU93FRXrpZnF_Qb6-KE0WyWhS_GJmmm0,1366
+checkpointer-2.14.5.dist-info/METADATA,sha256=vNYVwNi98WgyKujBVSLhnkWqsx3VMo7KHhgOiWe-i9w,11215
+checkpointer-2.14.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+checkpointer-2.14.5.dist-info/licenses/ATTRIBUTION.md,sha256=WF6L7-sD4s9t9ytVJOhjhpDoZ6TrWpqE3_bMdDIeJxI,1078
+checkpointer-2.14.5.dist-info/licenses/LICENSE,sha256=drXs6vIb7uW49r70UuMz2A1VtOCl626kiTbcmrar1Xo,1072
+checkpointer-2.14.5.dist-info/RECORD,,