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

checkpointer/checkpoint.py

@@ -12,8 +12,8 @@ from .fn_ident import Capturable, RawFunctionIdent, get_fn_ident
 from .object_hash import ObjectHash
 from .print_checkpoint import print_checkpoint
 from .storages import STORAGE_MAP, Storage, StorageType
-from .types import AwaitableValue, C, Coro, Fn, P, R, hash_by_from_annotation
-from .utils import flatten
+from .types import AwaitableValue, C, Coro, Fn, P, R, T, hash_by_from_annotation
+from .utils import flatten, to_coroutine
 
 DEFAULT_DIR = Path.home() / ".cache/checkpoints"
 
@@ -211,6 +211,8 @@ class CachedFunction(Generic[Fn]):
     try:
       data = storage.load(call_hash)
       print_checkpoint(params.verbosity >= 2, "REMEMBERED", call_id, "green")
+      if isinstance(data, AwaitableValue):
+        return to_coroutine(data.value)  # type: ignore
       return data
     except (EOFError, FileNotFoundError):
       pass
@@ -241,6 +243,16 @@ class CachedFunction(Generic[Fn]):
     except Exception as ex:
       raise CheckpointError("Could not load checkpoint") from ex
 
+  @overload
+  def get_or(self: Callable[P, Coro[R]], default: T, *args: P.args, **kw: P.kwargs) -> R | T: ...
+  @overload
+  def get_or(self: Callable[P, R], default: T, *args: P.args, **kw: P.kwargs) -> R | T: ...
+  def get_or(self, default, *args, **kw):
+    try:
+      return self.get(*args, **kw)  # type: ignore
+    except CheckpointError:
+      return default
+
   @overload
   def set(self: Callable[P, Coro[R]], value: AwaitableValue[R], *args: P.args, **kw: P.kwargs): ...
   @overload

checkpointer/storages/pickle_storage.py

@@ -1,22 +1,26 @@
 import pickle
 import shutil
 from datetime import datetime
+from io import BytesIO
 from pathlib import Path
+from ..utils import clear_directory
 from .storage import Storage
 
+try:
+  import polars as pl
+except:
+  pl = None
+
 def filedate(path: Path) -> datetime:
   return datetime.fromtimestamp(path.stat().st_mtime)
 
-def empty_dirs(path: Path):
-  nonempty_count = 0
-  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
-  if nonempty_count == 0:
-    yield path
+class ExtendedPickler(pickle.Pickler):
+  def reducer_override(self, obj): # type: ignore
+    if pl and isinstance(obj, pl.DataFrame):
+      buffer = BytesIO()
+      obj.rechunk().write_parquet(buffer)
+      return pl.read_parquet, (buffer.getvalue(),)
+    return NotImplemented
 
 class PickleStorage(Storage):
   def get_path(self, call_hash: str):
@@ -26,7 +30,7 @@ class PickleStorage(Storage):
     path = self.get_path(call_hash)
     path.parent.mkdir(parents=True, exist_ok=True)
     with path.open("wb") as file:
-      pickle.dump(data, file, -1)
+      ExtendedPickler(file, -1).dump(data)
     return data
 
   def exists(self, call_hash):
@@ -58,11 +62,7 @@ class PickleStorage(Storage):
           count += 1
           pkl_path.unlink(missing_ok=True)
       print(f"Removed {count} expired checkpoints for {self.cached_fn.__qualname__}")
-    if fn_path.exists():
-      for file in fn_path.glob("**/.DS_Store"):
-        file.unlink()
-      for directory in empty_dirs(fn_path):
-        directory.rmdir()
+    clear_directory(fn_path)
 
   def clear(self):
     fn_path = self.fn_dir().parent

checkpointer/utils.py

@@ -10,6 +10,9 @@ from .types import T
 cwd = Path.cwd().resolve()
 flatten = chain.from_iterable
 
+async def to_coroutine(value: T) -> T:
+  return value
+
 def is_class(obj) -> TypeGuard[Type]:
   return isinstance(obj, type)
 
@@ -112,3 +115,21 @@ class ContextVar(Generic[T]):
       yield
     finally:
       self.value = old
+
+def empty_dirs(path: Path) -> Iterable[Path]:
+  nonempty_count = 0
+  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
+  if nonempty_count == 0:
+    yield path
+
+def clear_directory(path: Path):
+  if path.is_dir():
+    for file in path.glob("**/.DS_Store"):
+      file.unlink()
+    for directory in empty_dirs(path):
+      directory.rmdir()

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

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

@@ -1,18 +1,18 @@
 checkpointer/__init__.py,sha256=mT706hd-BoTpQXYH_MiqreplONVDvjCWjwwGQ7oRe_8,921
-checkpointer/checkpoint.py,sha256=SUHu5ADRtSrapspk2DpBLaYnAQ7yuTc05lUELig_Tpc,9972
+checkpointer/checkpoint.py,sha256=N3MtNF4Oq3FV69LCfIuyTfNxrcnw9-dGzX312eYigVQ,10454
 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=CF28pGytCWV4URhRDFaIDawX8qmCAAByQp0ZkEwrA6c,3014
+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=PwFbCqYHjR3jiWuTnwxCU9cui7CINvq0wXIOKtf8dOk,2230
+checkpointer/storages/pickle_storage.py,sha256=WCbWGdDH046_GrlnttcqqGZ55wGOz2Rb7Gu0slI448g,2203
 checkpointer/storages/storage.py,sha256=JVxdq1DMhbq83bvflvIuW7okE1CQCa7poQPsj3x0ACg,1366
-checkpointer-2.14.1.dist-info/METADATA,sha256=KupseJCcWRNym_Ckit_F-oGLPIgqUsyKJvjl6ekEDZU,11215
-checkpointer-2.14.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-checkpointer-2.14.1.dist-info/licenses/ATTRIBUTION.md,sha256=WF6L7-sD4s9t9ytVJOhjhpDoZ6TrWpqE3_bMdDIeJxI,1078
-checkpointer-2.14.1.dist-info/licenses/LICENSE,sha256=drXs6vIb7uW49r70UuMz2A1VtOCl626kiTbcmrar1Xo,1072
-checkpointer-2.14.1.dist-info/RECORD,,
+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,,