checkpointer 2.14.2__tar.gz → 2.14.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpointer
-Version: 2.14.2
+Version: 2.14.4
 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.2 → checkpointer-2.14.4}/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-2.14.2 → checkpointer-2.14.4}/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
@@ -242,9 +244,9 @@ class CachedFunction(Generic[Fn]):
       raise CheckpointError("Could not load checkpoint") from ex
 
   @overload
-  def get_or(self: Callable[P, Coro[R]], default: R, *args: P.args, **kw: P.kwargs) -> R: ...
+  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: R, *args: P.args, **kw: P.kwargs) -> R: ...
+  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

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

@@ -1,3 +1,4 @@
+import gc
 from typing import Any
 from pathlib import Path
 from datetime import datetime
@@ -41,3 +42,12 @@ 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()
+  memory_stores = [obj for obj in gc.get_objects() if isinstance(obj, MemoryStorage)]
+  storage_keys = {store.fn_dir() for store in memory_stores}
+  for key in item_map.keys() - storage_keys:
+    del item_map[key]
+  for store in memory_stores:
+    store.cleanup(invalidated=False)

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

@@ -1,13 +1,27 @@
 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)
 
+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):
     return self.fn_dir() / f"{call_hash[:2]}/{call_hash[2:]}.pkl"
@@ -16,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):

{checkpointer-2.14.2 → checkpointer-2.14.4}/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)
 

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

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.14.2"
+version = "2.14.4"
 requires-python = ">=3.11"
 dependencies = []
 authors = [
@@ -38,6 +38,7 @@ dev = [
   "numpy>=2.2.1",
   "omg>=1.3.6",
   "poethepoet>=0.30.0",
+  "polars>=1.31.0",
   "pytest>=8.3.5",
   "pytest-asyncio>=0.26.0",
   "rich>=14.0.0",

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

@@ -8,7 +8,7 @@ resolution-markers = [
 
 [[package]]
 name = "checkpointer"
-version = "2.14.2"
+version = "2.14.4"
 source = { editable = "." }
 
 [package.dev-dependencies]
@@ -16,6 +16,7 @@ dev = [
     { name = "numpy" },
     { name = "omg" },
     { name = "poethepoet" },
+    { name = "polars" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
     { name = "rich" },
@@ -29,6 +30,7 @@ dev = [
     { name = "numpy", specifier = ">=2.2.1" },
     { name = "omg", specifier = ">=1.3.6" },
     { name = "poethepoet", specifier = ">=0.30.0" },
+    { name = "polars", specifier = ">=1.31.0" },
     { name = "pytest", specifier = ">=8.3.5" },
     { name = "pytest-asyncio", specifier = ">=0.26.0" },
     { name = "rich", specifier = ">=14.0.0" },
@@ -388,6 +390,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/06/e1/04f56c9d848d6135ca3328c5a2ca84d3303c358ad7828db290385e36a8cc/poethepoet-0.31.1-py3-none-any.whl", hash = "sha256:7fdfa0ac6074be9936723e7231b5bfaad2923e96c674a9857e81d326cf8ccdc2", size = 80238 },
 ]
 
+[[package]]
+name = "polars"
+version = "1.31.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/fd/f5/de1b5ecd7d0bd0dd87aa392937f759f9cc3997c5866a9a7f94eabf37cd48/polars-1.31.0.tar.gz", hash = "sha256:59a88054a5fc0135386268ceefdbb6a6cc012d21b5b44fed4f1d3faabbdcbf32", size = 4681224 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3d/6e/bdd0937653c1e7a564a09ae3bc7757ce83fedbf19da600c8b35d62c0182a/polars-1.31.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ccc68cd6877deecd46b13cbd2663ca89ab2a2cb1fe49d5cfc66a9cef166566d9", size = 34511354 },
+    { url = "https://files.pythonhosted.org/packages/77/fe/81aaca3540c1a5530b4bc4fd7f1b6f77100243d7bb9b7ad3478b770d8b3e/polars-1.31.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:a94c5550df397ad3c2d6adc212e59fd93d9b044ec974dd3653e121e6487a7d21", size = 31377712 },
+    { url = "https://files.pythonhosted.org/packages/b8/d9/5e2753784ea30d84b3e769a56f5e50ac5a89c129e87baa16ac0773eb4ef7/polars-1.31.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ada7940ed92bea65d5500ae7ac1f599798149df8faa5a6db150327c9ddbee4f1", size = 35050729 },
+    { url = "https://files.pythonhosted.org/packages/20/e8/a6bdfe7b687c1fe84bceb1f854c43415eaf0d2fdf3c679a9dc9c4776e462/polars-1.31.0-cp39-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:b324e6e3e8c6cc6593f9d72fe625f06af65e8d9d47c8686583585533a5e731e1", size = 32260836 },
+    { url = "https://files.pythonhosted.org/packages/6e/f6/9d9ad9dc4480d66502497e90ce29efc063373e1598f4bd9b6a38af3e08e7/polars-1.31.0-cp39-abi3-win_amd64.whl", hash = "sha256:3fd874d3432fc932863e8cceff2cff8a12a51976b053f2eb6326a0672134a632", size = 35156211 },
+    { url = "https://files.pythonhosted.org/packages/40/4b/0673a68ac4d6527fac951970e929c3b4440c654f994f0c957bd5556deb38/polars-1.31.0-cp39-abi3-win_arm64.whl", hash = "sha256:62ef23bb9d10dca4c2b945979f9a50812ac4ace4ed9e158a6b5d32a7322e6f75", size = 31469078 },
+]
+
 [[package]]
 name = "pygments"
 version = "2.19.1"