checkpointer 2.0.1__tar.gz → 2.0.2__tar.gz

{checkpointer-2.0.1 → checkpointer-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: checkpointer
-Version: 2.0.1
+Version: 2.0.2
 Summary: A Python library for memoizing function results with support for multiple storage backends, async runtimes, and automatic cache invalidation
 Project-URL: Repository, https://github.com/Reddan/checkpointer.git
 Author: Hampus Hallman
@@ -17,13 +17,13 @@ Description-Content-Type: text/markdown
 
 # checkpointer · [![License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/Reddan/checkpointer/blob/master/LICENSE) [![pypi](https://img.shields.io/pypi/v/checkpointer)](https://pypi.org/project/checkpointer/) [![Python 3.12](https://img.shields.io/badge/python-3.12-blue)](https://pypi.org/project/checkpointer/)
 
-`checkpointer` is a Python library for memoizing function results. It simplifies caching by providing a decorator-based API and supports various storage backends. It's designed for computationally expensive operations where caching can save time, or during development to avoid waiting for redundant computations. ⚡️
+`checkpointer` is a Python library for memoizing function results. It provides a decorator-based API with support for multiple storage backends. Use it for computationally expensive operations where caching can save time, or during development to avoid waiting for redundant computations.
 
 Adding or removing `@checkpoint` doesn't change how your code works, and it can be applied to any function, including ones you've already written, without altering their behavior or introducing side effects. The original function remains unchanged and can still be called directly when needed.
 
 ### Key Features:
-- 🗂️ **Multiple Storage Backends**: Supports in-memory, pickle, or your own custom storage.
-- 🎯 **Simple Decorator API**: Apply `@checkpoint` to functions.
+- 🗂️ **Multiple Storage Backends**: Built-in support for in-memory and pickle-based storage, or create your own.
+- 🎯 **Simple Decorator API**: Apply `@checkpoint` to functions without boilerplate.
 - 🔄 **Async and Sync Compatibility**: Works with synchronous functions and any Python async runtime (e.g., `asyncio`, `Trio`, `Curio`).
 - ⏲️ **Custom Expiration Logic**: Automatically invalidate old checkpoints.
 - 📂 **Flexible Path Configuration**: Control where checkpoints are stored.
@@ -48,21 +48,21 @@ def expensive_function(x: int) -> int:
     print("Computing...")
     return x ** 2
 
-result = expensive_function(4)  # Computes and stores result
-result = expensive_function(4)  # Loads from checkpoint
+result = expensive_function(4)  # Computes and stores the result
+result = expensive_function(4)  # Loads from the cache
 ```
 
 ---
 
 ## How It Works
 
-When you use `@checkpoint`, the function's **arguments** (`args`, `kwargs`) are hashed to create a unique identifier for each call. This identifier is used to store and retrieve cached results. If the same arguments are passed again, `checkpointer` will return the cached result instead of recomputing.
+When you use `@checkpoint`, the function's **arguments** (`args`, `kwargs`) are hashed to create a unique identifier for each call. This identifier is used to store and retrieve cached results. If the same arguments are passed again, `checkpointer` loads the cached result instead of recomputing.
 
-Additionally, `checkpointer` ensures that caches are invalidated when a function’s implementation or any of its dependencies change. Each function is assigned a hash based on:
-1. **Its source code**: Changes to the function’s code update its hash.
-2. **Dependent functions**: If a function calls others, changes to those will also update the hash.
+Additionally, `checkpointer` ensures that caches are invalidated when a function's implementation or any of its dependencies change. Each function is assigned a hash based on:
+1. **Its source code**: Changes to the function's code update its hash.
+2. **Dependent functions**: If a function calls others, changes in those dependencies will also update the hash.
 
-### Example: Cache Invalidation by Function Dependencies
+### Example: Cache Invalidation
 
 ```python
 def multiply(a, b):
@@ -77,95 +77,62 @@ def compute(a, b):
     return helper(a) + helper(b)
 ```
 
-If you change `multiply`, the checkpoints for both `helper` and `compute` will be invalidated and recomputed.
+If you modify `multiply`, caches for both `helper` and `compute` are invalidated and recomputed.
 
 ---
 
 ## Parameterization
 
-### Global Configuration
+### Custom Configuration
 
-You can configure a custom `Checkpointer`:
+Set up a `Checkpointer` instance with custom settings, and extend it by calling itself with overrides:
 
 ```python
 from checkpointer import checkpoint
 
-checkpoint = checkpoint(format="memory", root_path="/tmp/checkpoints")
-```
-
-Extend this configuration by calling itself again:
+IS_DEVELOPMENT = True  # Toggle based on your environment
 
-```python
-extended_checkpoint = checkpoint(format="pickle", verbosity=0)
+tmp_checkpoint = checkpoint(root_path="/tmp/checkpoints")
+dev_checkpoint = tmp_checkpoint(when=IS_DEVELOPMENT)  # Adds development-specific behavior
 ```
 
-### Per-Function Customization
-
-```python
-@checkpoint(format="pickle", verbosity=0)
-def my_function(x, y):
-    return x + y
-```
-
-### Combining Configurations
-
-```python
-checkpoint = checkpoint(format="memory", verbosity=1)
-quiet_checkpoint = checkpoint(verbosity=0)
-pickle_checkpoint = checkpoint(format="pickle", root_path="/tmp/pickle_checkpoints")
-
-@checkpoint
-def compute_square(n: int) -> int:
-    return n ** 2
-
-@quiet_checkpoint
-def compute_quietly(n: int) -> int:
-    return n ** 3
-
-@pickle_checkpoint
-def compute_sum(a: int, b: int) -> int:
-    return a + b
-```
+### Per-Function Customization & Layered Caching
 
-### Layered Caching
+Layer caches by stacking checkpoints:
 
 ```python
-IS_DEVELOPMENT = True  # Toggle based on environment
-
-dev_checkpoint = checkpoint(when=IS_DEVELOPMENT)
-
-@checkpoint(format="memory")
-@dev_checkpoint
+@checkpoint(format="memory")  # Always use memory storage
+@dev_checkpoint  # Adds caching during development
 def some_expensive_function():
     print("Performing a time-consuming operation...")
     return sum(i * i for i in range(10**6))
 ```
 
-- In development: Both `dev_checkpoint` and `memory` caches are active.
-- In production: Only the `memory` cache is active.
+- **In development**: Both `dev_checkpoint` and `memory` caches are active.
+- **In production**: Only the `memory` cache is active.
 
 ---
 
 ## Usage
 
 ### Force Recalculation
-Use `rerun` to force a recalculation and overwrite the stored checkpoint:
+Force a recalculation and overwrite the stored checkpoint:
 
 ```python
 result = expensive_function.rerun(4)
 ```
 
-### Bypass Checkpointer
+### Call the Original Function
 Use `fn` to directly call the original, undecorated function:
 
 ```python
 result = expensive_function.fn(4)
 ```
 
-This is especially useful **inside recursive functions**. By using `.fn` within the function itself, you avoid redundant caching of intermediate recursive calls while still caching the final result at the top level.
+This is especially useful **inside recursive functions** to avoid redundant caching of intermediate steps while still caching the final result.
 
 ### Retrieve Stored Checkpoints
-Access stored results without recalculating:
+Access cached results without recalculating:
 
 ```python
 stored_result = expensive_function.get(4)
@@ -175,14 +142,15 @@ stored_result = expensive_function.get(4)
 
 ## Storage Backends
 
-`checkpointer` supports flexible storage backends, including built-in options and custom implementations.
+`checkpointer` works with both built-in and custom storage backends, so you can use what's provided or roll your own as needed.
 
 ### Built-In Backends
 
-1. **PickleStorage**: Saves checkpoints to disk using Python's `pickle` module.
-2. **MemoryStorage**: Caches checkpoints in memory for fast, non-persistent use.
+1. **PickleStorage**: Stores checkpoints on disk using Python's `pickle`.
+2. **MemoryStorage**: Keeps checkpoints in memory for non-persistent, fast caching.
+
+You can specify a storage backend using either its name (`"pickle"` or `"memory"`) or its corresponding class (`PickleStorage` or `MemoryStorage`) in the `format` parameter:
 
-To use these backends, pass either `"pickle"` or `PickleStorage` (and similarly for `"memory"` or `MemoryStorage`) to the `format` parameter:
 ```python
 from checkpointer import checkpoint, PickleStorage, MemoryStorage
 
@@ -197,33 +165,27 @@ def memory_cached(x: int) -> int:
 
 ### Custom Storage Backends
 
-Create custom storage backends by implementing methods for storing, loading, and managing checkpoints. For example, a custom storage backend might use a database, cloud storage, or a specialized format.
+Create a custom storage backend by inheriting from the `Storage` class and implementing its methods. Access configuration options through the `self.checkpointer` attribute, an instance of `Checkpointer`.
+
+#### Example: Custom Storage Backend
 
-Example usage:
 ```python
 from checkpointer import checkpoint, Storage
-from typing import Any
-from pathlib import Path
 from datetime import datetime
 
-class CustomStorage(Storage):  # Optional for type hinting
-    @staticmethod
-    def exists(path: Path) -> bool: ...
-    @staticmethod
-    def checkpoint_date(path: Path) -> datetime: ...
-    @staticmethod
-    def store(path: Path, data: Any) -> None: ...
-    @staticmethod
-    def load(path: Path) -> Any: ...
-    @staticmethod
-    def delete(path: Path) -> None: ...
+class CustomStorage(Storage):
+    def exists(self, path) -> bool: ...  # Check if a checkpoint exists at the given path
+    def checkpoint_date(self, path) -> datetime: ...  # Return the date the checkpoint was created
+    def store(self, path, data): ...  # Save the checkpoint data
+    def load(self, path): ...  # Return the checkpoint data
+    def delete(self, path): ...  # Delete the checkpoint
 
 @checkpoint(format=CustomStorage)
 def custom_cached(x: int):
     return x ** 2
 ```
 
-This flexibility allows you to adapt `checkpointer` to meet any storage requirement, whether persistent or in-memory.
+Using a custom backend lets you tailor storage to your application, whether it involves databases, cloud storage, or custom file formats.
 
 ---
 

{checkpointer-2.0.1 → checkpointer-2.0.2}/README.md

@@ -1,12 +1,12 @@
 # checkpointer · [![License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/Reddan/checkpointer/blob/master/LICENSE) [![pypi](https://img.shields.io/pypi/v/checkpointer)](https://pypi.org/project/checkpointer/) [![Python 3.12](https://img.shields.io/badge/python-3.12-blue)](https://pypi.org/project/checkpointer/)
 
-`checkpointer` is a Python library for memoizing function results. It simplifies caching by providing a decorator-based API and supports various storage backends. It's designed for computationally expensive operations where caching can save time, or during development to avoid waiting for redundant computations. ⚡️
+`checkpointer` is a Python library for memoizing function results. It provides a decorator-based API with support for multiple storage backends. Use it for computationally expensive operations where caching can save time, or during development to avoid waiting for redundant computations.
 
 Adding or removing `@checkpoint` doesn't change how your code works, and it can be applied to any function, including ones you've already written, without altering their behavior or introducing side effects. The original function remains unchanged and can still be called directly when needed.
 
 ### Key Features:
-- 🗂️ **Multiple Storage Backends**: Supports in-memory, pickle, or your own custom storage.
-- 🎯 **Simple Decorator API**: Apply `@checkpoint` to functions.
+- 🗂️ **Multiple Storage Backends**: Built-in support for in-memory and pickle-based storage, or create your own.
+- 🎯 **Simple Decorator API**: Apply `@checkpoint` to functions without boilerplate.
 - 🔄 **Async and Sync Compatibility**: Works with synchronous functions and any Python async runtime (e.g., `asyncio`, `Trio`, `Curio`).
 - ⏲️ **Custom Expiration Logic**: Automatically invalidate old checkpoints.
 - 📂 **Flexible Path Configuration**: Control where checkpoints are stored.
@@ -31,21 +31,21 @@ def expensive_function(x: int) -> int:
     print("Computing...")
     return x ** 2
 
-result = expensive_function(4)  # Computes and stores result
-result = expensive_function(4)  # Loads from checkpoint
+result = expensive_function(4)  # Computes and stores the result
+result = expensive_function(4)  # Loads from the cache
 ```
 
 ---
 
 ## How It Works
 
-When you use `@checkpoint`, the function's **arguments** (`args`, `kwargs`) are hashed to create a unique identifier for each call. This identifier is used to store and retrieve cached results. If the same arguments are passed again, `checkpointer` will return the cached result instead of recomputing.
+When you use `@checkpoint`, the function's **arguments** (`args`, `kwargs`) are hashed to create a unique identifier for each call. This identifier is used to store and retrieve cached results. If the same arguments are passed again, `checkpointer` loads the cached result instead of recomputing.
 
-Additionally, `checkpointer` ensures that caches are invalidated when a function’s implementation or any of its dependencies change. Each function is assigned a hash based on:
-1. **Its source code**: Changes to the function’s code update its hash.
-2. **Dependent functions**: If a function calls others, changes to those will also update the hash.
+Additionally, `checkpointer` ensures that caches are invalidated when a function's implementation or any of its dependencies change. Each function is assigned a hash based on:
+1. **Its source code**: Changes to the function's code update its hash.
+2. **Dependent functions**: If a function calls others, changes in those dependencies will also update the hash.
 
-### Example: Cache Invalidation by Function Dependencies
+### Example: Cache Invalidation
 
 ```python
 def multiply(a, b):
@@ -60,95 +60,62 @@ def compute(a, b):
     return helper(a) + helper(b)
 ```
 
-If you change `multiply`, the checkpoints for both `helper` and `compute` will be invalidated and recomputed.
+If you modify `multiply`, caches for both `helper` and `compute` are invalidated and recomputed.
 
 ---
 
 ## Parameterization
 
-### Global Configuration
+### Custom Configuration
 
-You can configure a custom `Checkpointer`:
+Set up a `Checkpointer` instance with custom settings, and extend it by calling itself with overrides:
 
 ```python
 from checkpointer import checkpoint
 
-checkpoint = checkpoint(format="memory", root_path="/tmp/checkpoints")
-```
-
-Extend this configuration by calling itself again:
+IS_DEVELOPMENT = True  # Toggle based on your environment
 
-```python
-extended_checkpoint = checkpoint(format="pickle", verbosity=0)
+tmp_checkpoint = checkpoint(root_path="/tmp/checkpoints")
+dev_checkpoint = tmp_checkpoint(when=IS_DEVELOPMENT)  # Adds development-specific behavior
 ```
 
-### Per-Function Customization
-
-```python
-@checkpoint(format="pickle", verbosity=0)
-def my_function(x, y):
-    return x + y
-```
-
-### Combining Configurations
-
-```python
-checkpoint = checkpoint(format="memory", verbosity=1)
-quiet_checkpoint = checkpoint(verbosity=0)
-pickle_checkpoint = checkpoint(format="pickle", root_path="/tmp/pickle_checkpoints")
-
-@checkpoint
-def compute_square(n: int) -> int:
-    return n ** 2
-
-@quiet_checkpoint
-def compute_quietly(n: int) -> int:
-    return n ** 3
-
-@pickle_checkpoint
-def compute_sum(a: int, b: int) -> int:
-    return a + b
-```
+### Per-Function Customization & Layered Caching
 
-### Layered Caching
+Layer caches by stacking checkpoints:
 
 ```python
-IS_DEVELOPMENT = True  # Toggle based on environment
-
-dev_checkpoint = checkpoint(when=IS_DEVELOPMENT)
-
-@checkpoint(format="memory")
-@dev_checkpoint
+@checkpoint(format="memory")  # Always use memory storage
+@dev_checkpoint  # Adds caching during development
 def some_expensive_function():
     print("Performing a time-consuming operation...")
     return sum(i * i for i in range(10**6))
 ```
 
-- In development: Both `dev_checkpoint` and `memory` caches are active.
-- In production: Only the `memory` cache is active.
+- **In development**: Both `dev_checkpoint` and `memory` caches are active.
+- **In production**: Only the `memory` cache is active.
 
 ---
 
 ## Usage
 
 ### Force Recalculation
-Use `rerun` to force a recalculation and overwrite the stored checkpoint:
+Force a recalculation and overwrite the stored checkpoint:
 
 ```python
 result = expensive_function.rerun(4)
 ```
 
-### Bypass Checkpointer
+### Call the Original Function
 Use `fn` to directly call the original, undecorated function:
 
 ```python
 result = expensive_function.fn(4)
 ```
 
-This is especially useful **inside recursive functions**. By using `.fn` within the function itself, you avoid redundant caching of intermediate recursive calls while still caching the final result at the top level.
+This is especially useful **inside recursive functions** to avoid redundant caching of intermediate steps while still caching the final result.
 
 ### Retrieve Stored Checkpoints
-Access stored results without recalculating:
+Access cached results without recalculating:
 
 ```python
 stored_result = expensive_function.get(4)
@@ -158,14 +125,15 @@ stored_result = expensive_function.get(4)
 
 ## Storage Backends
 
-`checkpointer` supports flexible storage backends, including built-in options and custom implementations.
+`checkpointer` works with both built-in and custom storage backends, so you can use what's provided or roll your own as needed.
 
 ### Built-In Backends
 
-1. **PickleStorage**: Saves checkpoints to disk using Python's `pickle` module.
-2. **MemoryStorage**: Caches checkpoints in memory for fast, non-persistent use.
+1. **PickleStorage**: Stores checkpoints on disk using Python's `pickle`.
+2. **MemoryStorage**: Keeps checkpoints in memory for non-persistent, fast caching.
+
+You can specify a storage backend using either its name (`"pickle"` or `"memory"`) or its corresponding class (`PickleStorage` or `MemoryStorage`) in the `format` parameter:
 
-To use these backends, pass either `"pickle"` or `PickleStorage` (and similarly for `"memory"` or `MemoryStorage`) to the `format` parameter:
 ```python
 from checkpointer import checkpoint, PickleStorage, MemoryStorage
 
@@ -180,33 +148,27 @@ def memory_cached(x: int) -> int:
 
 ### Custom Storage Backends
 
-Create custom storage backends by implementing methods for storing, loading, and managing checkpoints. For example, a custom storage backend might use a database, cloud storage, or a specialized format.
+Create a custom storage backend by inheriting from the `Storage` class and implementing its methods. Access configuration options through the `self.checkpointer` attribute, an instance of `Checkpointer`.
+
+#### Example: Custom Storage Backend
 
-Example usage:
 ```python
 from checkpointer import checkpoint, Storage
-from typing import Any
-from pathlib import Path
 from datetime import datetime
 
-class CustomStorage(Storage):  # Optional for type hinting
-    @staticmethod
-    def exists(path: Path) -> bool: ...
-    @staticmethod
-    def checkpoint_date(path: Path) -> datetime: ...
-    @staticmethod
-    def store(path: Path, data: Any) -> None: ...
-    @staticmethod
-    def load(path: Path) -> Any: ...
-    @staticmethod
-    def delete(path: Path) -> None: ...
+class CustomStorage(Storage):
+    def exists(self, path) -> bool: ...  # Check if a checkpoint exists at the given path
+    def checkpoint_date(self, path) -> datetime: ...  # Return the date the checkpoint was created
+    def store(self, path, data): ...  # Save the checkpoint data
+    def load(self, path): ...  # Return the checkpoint data
+    def delete(self, path): ...  # Delete the checkpoint
 
 @checkpoint(format=CustomStorage)
 def custom_cached(x: int):
     return x ** 2
 ```
 
-This flexibility allows you to adapt `checkpointer` to meet any storage requirement, whether persistent or in-memory.
+Using a custom backend lets you tailor storage to your application, whether it involves databases, cloud storage, or custom file formats.
 
 ---
 

{checkpointer-2.0.1 → checkpointer-2.0.2}/checkpointer/__init__.py

@@ -5,5 +5,5 @@ import tempfile
 
 create_checkpointer = Checkpointer
 checkpoint = Checkpointer()
-memory_checkpoint = Checkpointer(format="memory")
+memory_checkpoint = Checkpointer(format="memory", verbosity=0)
 tmp_checkpoint = Checkpointer(root_path=tempfile.gettempdir() + "/checkpoints")

{checkpointer-2.0.1 → checkpointer-2.0.2}/checkpointer/checkpoint.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
 import inspect
 import relib.hashing as hashing
-from typing import Generic, TypeVar, TypedDict, Callable, Unpack, Literal, Union, Any, cast, overload
-from datetime import datetime
+from typing import Generic, TypeVar, Type, TypedDict, Callable, Unpack, Literal, Any, cast, overload
 from pathlib import Path
+from datetime import datetime
 from functools import update_wrapper
 from .types import Storage
 from .function_body import get_function_hash
@@ -15,13 +16,13 @@ from .print_checkpoint import print_checkpoint
 Fn = TypeVar("Fn", bound=Callable)
 
 DEFAULT_DIR = Path.home() / ".cache/checkpoints"
-STORAGE_MAP = {"memory": MemoryStorage, "pickle": PickleStorage, "bcolz": BcolzStorage}
+STORAGE_MAP: dict[str, Type[Storage]] = {"memory": MemoryStorage, "pickle": PickleStorage, "bcolz": BcolzStorage}
 
 class CheckpointError(Exception):
   pass
 
 class CheckpointerOpts(TypedDict, total=False):
-  format: Storage | Literal["pickle", "memory", "bcolz"]
+  format: Type[Storage] | Literal["pickle", "memory", "bcolz"]
   root_path: Path | str | None
   when: bool
   verbosity: Literal[0, 1]
@@ -37,14 +38,11 @@ class Checkpointer:
     self.path = opts.get("path")
     self.should_expire = opts.get("should_expire")
 
-  def get_storage(self) -> Storage:
-    return STORAGE_MAP[self.format] if isinstance(self.format, str) else self.format
-
   @overload
-  def __call__(self, fn: Fn, **override_opts: Unpack[CheckpointerOpts]) -> "CheckpointFn[Fn]": ...
+  def __call__(self, fn: Fn, **override_opts: Unpack[CheckpointerOpts]) -> CheckpointFn[Fn]: ...
   @overload
-  def __call__(self, fn: None=None, **override_opts: Unpack[CheckpointerOpts]) -> "Checkpointer": ...
-  def __call__(self, fn: Fn | None=None, **override_opts: Unpack[CheckpointerOpts]) -> Union["Checkpointer", "CheckpointFn[Fn]"]:
+  def __call__(self, fn: None=None, **override_opts: Unpack[CheckpointerOpts]) -> Checkpointer: ...
+  def __call__(self, fn: Fn | None=None, **override_opts: Unpack[CheckpointerOpts]) -> Checkpointer | CheckpointFn[Fn]:
     if override_opts:
       opts = CheckpointerOpts(**{**self.__dict__, **override_opts})
       return Checkpointer(**opts)(fn)
@@ -56,11 +54,13 @@ class CheckpointFn(Generic[Fn]):
     wrapped = unwrap_fn(fn)
     file_name = Path(wrapped.__code__.co_filename).name
     update_wrapper(cast(Callable, self), wrapped)
+    storage = STORAGE_MAP[checkpointer.format] if isinstance(checkpointer.format, str) else checkpointer.format
     self.checkpointer = checkpointer
     self.fn = fn
     self.fn_hash = get_function_hash(wrapped)
     self.fn_id = f"{file_name}/{wrapped.__name__}"
     self.is_async = inspect.iscoroutinefunction(wrapped)
+    self.storage = storage(checkpointer)
 
   def get_checkpoint_id(self, args: tuple, kw: dict) -> str:
     if not callable(self.checkpointer.path):
@@ -73,27 +73,26 @@ class CheckpointFn(Generic[Fn]):
   async def _store_on_demand(self, args: tuple, kw: dict, rerun: bool):
     checkpoint_id = self.get_checkpoint_id(args, kw)
     checkpoint_path = self.checkpointer.root_path / checkpoint_id
-    storage = self.checkpointer.get_storage()
-    should_log = storage is not MemoryStorage and self.checkpointer.verbosity > 0
+    should_log = self.checkpointer.verbosity > 0
     refresh = rerun \
-      or not storage.exists(checkpoint_path) \
-      or (self.checkpointer.should_expire and self.checkpointer.should_expire(storage.checkpoint_date(checkpoint_path)))
+      or not self.storage.exists(checkpoint_path) \
+      or (self.checkpointer.should_expire and self.checkpointer.should_expire(self.storage.checkpoint_date(checkpoint_path)))
 
     if refresh:
       print_checkpoint(should_log, "MEMORIZING", checkpoint_id, "blue")
       data = self.fn(*args, **kw)
       if inspect.iscoroutine(data):
         data = await data
-      storage.store(checkpoint_path, data)
+      self.storage.store(checkpoint_path, data)
       return data
 
     try:
-      data = storage.load(checkpoint_path)
+      data = self.storage.load(checkpoint_path)
       print_checkpoint(should_log, "REMEMBERED", checkpoint_id, "green")
       return data
     except (EOFError, FileNotFoundError):
       print_checkpoint(should_log, "CORRUPTED", checkpoint_id, "yellow")
-      storage.delete(checkpoint_path)
+      self.storage.delete(checkpoint_path)
       return await self._store_on_demand(args, kw, rerun)
 
   def _call(self, args: tuple, kw: dict, rerun=False):
@@ -107,8 +106,7 @@ class CheckpointFn(Generic[Fn]):
 
   def get(self, *args, **kw) -> Any:
     checkpoint_path = self.checkpointer.root_path / self.get_checkpoint_id(args, kw)
-    storage = self.checkpointer.get_storage()
     try:
-      return storage.load(checkpoint_path)
+      return self.storage.load(checkpoint_path)
     except:
       raise CheckpointError("Could not load checkpoint")

{checkpointer-2.0.1 → checkpointer-2.0.2}/checkpointer/storages/bcolz_storage.py

@@ -24,16 +24,13 @@ def insert_data(path: Path, data):
   c.flush()
 
 class BcolzStorage(Storage):
-  @staticmethod
-  def exists(path):
+  def exists(self, path):
     return path.exists()
 
-  @staticmethod
-  def checkpoint_date(path):
+  def checkpoint_date(self, path):
     return datetime.fromtimestamp(path.stat().st_mtime)
 
-  @staticmethod
-  def store(path, data):
+  def store(self, path, data):
     metapath = get_metapath(path)
     path.parent.mkdir(parents=True, exist_ok=True)
     data_type_str = get_data_type_str(data)
@@ -48,12 +45,11 @@ class BcolzStorage(Storage):
     if data_type_str in ["tuple", "dict"]:
       for i in range(len(fields)):
         child_path = Path(f"{path} ({i})")
-        BcolzStorage.store(child_path, data[fields[i]])
+        self.store(child_path, data[fields[i]])
     else:
       insert_data(path, data)
 
-  @staticmethod
-  def load(path):
+  def load(self, path):
     import bcolz
     metapath = get_metapath(path)
     meta_data = bcolz.open(metapath)[:][0]
@@ -61,7 +57,7 @@ class BcolzStorage(Storage):
     if data_type_str in ["tuple", "dict"]:
       fields = meta_data["fields"]
       partitions = range(len(fields))
-      data = [BcolzStorage.load(Path(f"{path} ({i})")) for i in partitions]
+      data = [self.load(Path(f"{path} ({i})")) for i in partitions]
       if data_type_str == "tuple":
         return tuple(data)
       else:
@@ -75,8 +71,7 @@ class BcolzStorage(Storage):
       else:
         return data[:]
 
-  @staticmethod
-  def delete(path):
+  def delete(self, path):
     # NOTE: Not recursive
     metapath = get_metapath(path)
     try:

checkpointer-2.0.2/checkpointer/storages/memory_storage.py

@@ -0,0 +1,25 @@
+from typing import Any
+from pathlib import Path
+from datetime import datetime
+from ..types import Storage
+
+item_map: dict[str, tuple[datetime, Any]] = {}
+
+class MemoryStorage(Storage):
+  def get_short_path(self, path: Path):
+    return str(path.relative_to(self.checkpointer.root_path))
+
+  def exists(self, path):
+    return self.get_short_path(path) in item_map
+
+  def checkpoint_date(self, path):
+    return item_map[self.get_short_path(path)][0]
+
+  def store(self, path, data):
+    item_map[self.get_short_path(path)] = (datetime.now(), data)
+
+  def load(self, path):
+    return item_map[self.get_short_path(path)][1]
+
+  def delete(self, path):
+    del item_map[self.get_short_path(path)]

{checkpointer-2.0.1 → checkpointer-2.0.2}/checkpointer/storages/pickle_storage.py

@@ -7,29 +7,24 @@ def get_path(path: Path):
   return path.with_name(f"{path.name}.pkl")
 
 class PickleStorage(Storage):
-  @staticmethod
-  def exists(path):
+  def exists(self, path):
     return get_path(path).exists()
 
-  @staticmethod
-  def checkpoint_date(path):
+  def checkpoint_date(self, path):
     return datetime.fromtimestamp(get_path(path).stat().st_mtime)
 
-  @staticmethod
-  def store(path, data):
+  def store(self, path, data):
     full_path = get_path(path)
     full_path.parent.mkdir(parents=True, exist_ok=True)
     with full_path.open("wb") as file:
       pickle.dump(data, file, -1)
 
-  @staticmethod
-  def load(path):
+  def load(self, path):
     full_path = get_path(path)
     with full_path.open("rb") as file:
       return pickle.load(file)
 
-  @staticmethod
-  def delete(path):
+  def delete(self, path):
     try:
       get_path(path).unlink()
     except FileNotFoundError:

checkpointer-2.0.2/checkpointer/types.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+from typing import Any, TYPE_CHECKING
+from pathlib import Path
+from datetime import datetime
+
+if TYPE_CHECKING:
+  from .checkpoint import Checkpointer
+
+class Storage:
+  checkpointer: Checkpointer
+
+  def __init__(self, checkpointer: Checkpointer):
+    self.checkpointer = checkpointer
+
+  def exists(self, path: Path) -> bool: ...
+
+  def checkpoint_date(self, path: Path) -> datetime: ...
+
+  def store(self, path: Path, data: Any) -> None: ...
+
+  def load(self, path: Path) -> Any: ...
+
+  def delete(self, path: Path) -> None: ...

{checkpointer-2.0.1 → checkpointer-2.0.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "checkpointer"
-version = "2.0.1"
+version = "2.0.2"
 requires-python = ">=3.12"
 dependencies = [
   "relib",

{checkpointer-2.0.1 → checkpointer-2.0.2}/uv.lock

@@ -3,7 +3,7 @@ requires-python = ">=3.12"
 
 [[package]]
 name = "checkpointer"
-version = "2.0.1"
+version = "2.0.2"
 source = { editable = "." }
 dependencies = [
     { name = "relib" },

checkpointer-2.0.1/checkpointer/storages/memory_storage.py

@@ -1,28 +0,0 @@
-from datetime import datetime
-from ..types import Storage
-
-store = {}
-date_stored = {}
-
-class MemoryStorage(Storage):
-  @staticmethod
-  def exists(path):
-    return str(path) in store
-
-  @staticmethod
-  def checkpoint_date(path):
-    return date_stored[str(path)]
-
-  @staticmethod
-  def store(path, data):
-    store[str(path)] = data
-    date_stored[str(path)] = datetime.now()
-
-  @staticmethod
-  def load(path):
-    return store[str(path)]
-
-  @staticmethod
-  def delete(path):
-    del store[str(path)]
-    del date_stored[str(path)]

checkpointer-2.0.1/checkpointer/types.py

@@ -1,19 +0,0 @@
-from typing import Protocol, Any
-from pathlib import Path
-from datetime import datetime
-
-class Storage(Protocol):
-  @staticmethod
-  def exists(path: Path) -> bool: ...
-
-  @staticmethod
-  def checkpoint_date(path: Path) -> datetime: ...
-
-  @staticmethod
-  def store(path: Path, data: Any) -> None: ...
-
-  @staticmethod
-  def load(path: Path) -> Any: ...
-
-  @staticmethod
-  def delete(path: Path) -> None: ...