checkpointer 1.1.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

checkpointer/__init__.py

@@ -1,10 +1,9 @@
-import os
-
-from .checkpoint import create_checkpointer, read_only, default_dir
-from .storage import store_on_demand, read_from_store
+from .checkpoint import Checkpointer, CheckpointFn
+from .checkpoint import CheckpointError, CheckpointReadFail
+from .types import Storage
 from .function_body import get_function_hash
 
-storage_dir = os.environ.get('CHECKPOINTS_DIR', default_dir)
-verbosity = int(os.environ.get('CHECKPOINTS_VERBOSITY', '1'))
-
-checkpoint = create_checkpointer(root_path=storage_dir, verbosity=verbosity)
+create_checkpointer = Checkpointer
+checkpoint = Checkpointer()
+memory_checkpoint = Checkpointer(format="memory")
+tmp_checkpoint = Checkpointer(root_path="/tmp/checkpoints")

checkpointer/checkpoint.py

@@ -1,68 +1,123 @@
 import inspect
-from collections import namedtuple
-from pathlib import Path
-from functools import wraps
 import relib.hashing as hashing
-from . import storage
+from typing import Generic, TypeVar, TypedDict, Unpack, Literal, Union, Any, cast, overload
+from collections.abc import Callable
+from datetime import datetime
+from pathlib import Path
+from functools import update_wrapper
+from .types import Storage
 from .function_body import get_function_hash
-from .utils import unwrap_func, sync_resolve_coroutine
-
-default_dir = Path.home() / '.checkpoints'
-
-def get_invoke_path(func, function_hash, args, kwargs, path):
-  if type(path) == str:
-    return path
-  elif callable(path):
-    return path(*args, **kwargs)
-  else:
-    hash = hashing.hash([function_hash, args, kwargs or 0])
-    file_name = Path(func.__code__.co_filename).name
-    name = func.__name__
-    return file_name + '/' + name + '/' + hash
-
-def create_checkpointer_from_config(config):
-  def checkpoint(opt_func=None, format=config.format, path=None, should_expire=None, when=True):
-    def receive_func(func):
-      if not (config.when and when):
-        return func
-
-      config_ = config._replace(format=format)
-      is_async = inspect.iscoroutinefunction(func)
-      unwrapped_func = unwrap_func(func)
-      function_hash = get_function_hash(unwrapped_func)
-
-      @wraps(unwrapped_func)
-      def wrapper(*args, **kwargs):
-        compute = lambda: func(*args, **kwargs)
-        recheck = kwargs.pop('recheck', False)
-        invoke_path = get_invoke_path(unwrapped_func, function_hash, args, kwargs, path)
-        coroutine = storage.store_on_demand(compute, invoke_path, config_, recheck, should_expire)
-        if is_async:
-          return coroutine
-        else:
-          return sync_resolve_coroutine(coroutine)
-
-      wrapper.checkpoint_config = config_
-
-      return wrapper
-
-    return receive_func(opt_func) if callable(opt_func) else receive_func
-
-  return checkpoint
-
-def create_checkpointer(format='pickle', root_path=default_dir, when=True, verbosity=1):
-  root_path = None if root_path is None else Path(root_path)
-  opts = locals()
-  CheckpointerConfig = namedtuple('CheckpointerConfig', sorted(opts))
-  config = CheckpointerConfig(**opts)
-  return create_checkpointer_from_config(config)
-
-def read_only(wrapper_func, config, format='pickle', path=None):
-  func = unwrap_func(wrapper_func)
-  function_hash = get_function_hash(func)
-
-  def wrapper(*args, **kwargs):
-    invoke_path = get_invoke_path(func, function_hash, args, kwargs, path)
-    return storage.read_from_store(invoke_path, config, storage=format)
-
-  return wrapper
+from .utils import unwrap_fn, sync_resolve_coroutine
+from .storages.pickle_storage import PickleStorage
+from .storages.memory_storage import MemoryStorage
+from .storages.bcolz_storage import BcolzStorage
+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}
+
+class CheckpointError(Exception):
+  pass
+
+class CheckpointReadFail(CheckpointError):
+  pass
+
+StorageType = Literal["pickle", "memory", "bcolz"] | Storage
+CheckpointPath = str | Callable[..., str] | None
+ShouldExpire = Callable[[datetime], bool]
+
+class CheckpointerOpts(TypedDict, total=False):
+  format: StorageType
+  root_path: Path | str | None
+  when: bool
+  verbosity: Literal[0, 1]
+  path: CheckpointPath
+  should_expire: ShouldExpire
+
+class Checkpointer:
+  def __init__(self, **opts: Unpack[CheckpointerOpts]):
+    self.format = opts.get("format", "pickle")
+    self.root_path = Path(opts.get("root_path", DEFAULT_DIR) or ".")
+    self.when = opts.get("when", True)
+    self.verbosity = opts.get("verbosity", 1)
+    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]": ...
+  @overload
+  def __call__(self, fn=None, **override_opts: Unpack[CheckpointerOpts]) -> "Checkpointer": ...
+  def __call__(self, fn: Fn | None=None, **override_opts: Unpack[CheckpointerOpts]) -> Union["Checkpointer", "CheckpointFn[Fn]"]:
+    if override_opts:
+      opts = CheckpointerOpts(**{**self.__dict__, **override_opts})
+      return Checkpointer(**opts)(fn)
+
+    return CheckpointFn(self, fn) if callable(fn) else self
+
+class CheckpointFn(Generic[Fn]):
+  def __init__(self, checkpointer: Checkpointer, fn: Fn):
+    wrapped = unwrap_fn(fn)
+    file_name = Path(wrapped.__code__.co_filename).name
+    update_wrapper(cast(Callable, self), wrapped)
+    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(fn)
+
+  def get_checkpoint_id(self, args: tuple, kw: dict) -> str:
+    match self.checkpointer.path:
+      case str() as path:
+        return path
+      case Callable() as path:
+        p = path(*args, **kw)
+        assert isinstance(p, str), "path function must return a string"
+        return p
+      case _:
+        return f"{self.fn_id}/{hashing.hash([self.fn_hash, args, kw or 0])}"
+
+  async def _store_on_demand(self, args: tuple, kw: dict, force: 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
+    refresh = force \
+      or storage.is_expired(checkpoint_path) \
+      or (self.checkpointer.should_expire and storage.should_expire(checkpoint_path, self.checkpointer.should_expire))
+
+    if refresh:
+      print_checkpoint(should_log, "MEMORIZING", checkpoint_id, "blue")
+      data = self.fn(*args, **kw)
+      if inspect.iscoroutine(data):
+        data = await data
+      return storage.store_data(checkpoint_path, data)
+
+    try:
+      data = storage.load_data(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_data(checkpoint_path)
+      return await self._store_on_demand(args, kw, force)
+
+  def _call(self, args: tuple, kw: dict, force=False):
+    if not self.checkpointer.when:
+      return self.fn(*args, **kw)
+    coroutine = self._store_on_demand(args, kw, force)
+    return coroutine if self.is_async else sync_resolve_coroutine(coroutine)
+
+  __call__: Fn = cast(Fn, lambda self, *args, **kw: self._call(args, kw))
+  rerun: Fn = cast(Fn, lambda self, *args, **kw: self._call(args, kw, True))
+
+  def get(self, *args, **kw) -> Any:
+    checkpoint_path = self.checkpointer.root_path / self.get_checkpoint_id(args, kw)
+    try:
+      return self.checkpointer.get_storage().load_data(checkpoint_path)
+    except:
+      raise CheckpointReadFail()

checkpointer/function_body.py

@@ -1,52 +1,46 @@
 import inspect
-from types import FunctionType, CodeType
 import relib.hashing as hashing
+from collections.abc import Callable
+from types import FunctionType, CodeType
 from pathlib import Path
-from .utils import unwrap_func
+from .utils import unwrap_fn
 
 cwd = Path.cwd()
 
-def get_func_path(func):
-  return Path(inspect.getfile(func)).absolute()
+def get_fn_path(fn: Callable) -> Path:
+  return Path(inspect.getfile(fn)).resolve()
 
-def get_function_body(func):
+def get_function_body(fn: Callable) -> str:
   # TODO: Strip comments
-  lines = inspect.getsourcelines(func)[0]
+  lines = inspect.getsourcelines(fn)[0]
   lines = [line.rstrip() for line in lines]
   lines = [line for line in lines if line]
-  return '\n'.join(lines)
+  return "\n".join(lines)
 
-def get_code_children(__code__):
-  consts = [const for const in __code__.co_consts if isinstance(const, CodeType)]
+def get_code_children(code: CodeType) -> list[str]:
+  consts = [const for const in code.co_consts if isinstance(const, CodeType)]
   children = [child for const in consts for child in get_code_children(const)]
-  return list(__code__.co_names) + children
-
-def get_func_children(func, neighbor_funcs=[]):
-  def get_candidate_func(co_name):
-    candidate_func = func.__globals__.get(co_name, None)
-    return unwrap_func(candidate_func)
-
-  def clear_candidate(candidate_func):
-    return isinstance(candidate_func, FunctionType) \
-      and candidate_func not in neighbor_funcs \
-      and cwd in get_func_path(candidate_func).parents
-
-  code_children = get_code_children(func.__code__)
-
-  func_children = [get_candidate_func(child) for child in code_children]
-  func_children = [child for child in func_children if clear_candidate(child)]
-
-  funcs = [func] + [
-    deep_child
-    for child_func in func_children
-    for deep_child in get_func_children(child_func, func_children)
-  ]
-  funcs = sorted(set(funcs), key=lambda func: func.__name__)
-
-  return funcs
-
-def get_function_hash(func):
-  funcs = [func] + get_func_children(func)
-  function_bodies = list(map(get_function_body, funcs))
-  function_bodies_hash = hashing.hash(function_bodies)
-  return function_bodies_hash
+  return list(code.co_names) + children
+
+def is_user_fn(candidate_fn, cleared_fns: set[Callable]) -> bool:
+  return isinstance(candidate_fn, FunctionType) \
+    and candidate_fn not in cleared_fns \
+    and cwd in get_fn_path(candidate_fn).parents
+
+def append_fn_children(cleared_fns: set[Callable], fn: Callable) -> None:
+  code_children = get_code_children(fn.__code__)
+  fn_children = [unwrap_fn(fn.__globals__.get(co_name, None)) for co_name in code_children]
+  fn_children = [child for child in fn_children if is_user_fn(child, cleared_fns)]
+  cleared_fns.update(fn_children)
+  for child_fn in fn_children:
+    append_fn_children(cleared_fns, child_fn)
+
+def get_fn_children(fn: Callable) -> list[Callable]:
+  cleared_fns: set[Callable] = set()
+  append_fn_children(cleared_fns, fn)
+  return sorted(cleared_fns, key=lambda fn: fn.__name__)
+
+def get_function_hash(fn: Callable) -> str:
+  fns = [fn] + get_fn_children(fn)
+  fn_bodies = list(map(get_function_body, fns))
+  return hashing.hash(fn_bodies)

checkpointer/print_checkpoint.py

@@ -0,0 +1,52 @@
+import io
+import os
+import sys
+from typing import Literal
+
+Color = Literal[
+  "black", "grey", "red", "green", "yellow", "blue", "magenta",
+  "cyan", "light_grey", "dark_grey", "light_red", "light_green",
+  "light_yellow", "light_blue", "light_magenta", "light_cyan", "white",
+]
+
+COLOR_MAP: dict[Color, int] = {
+  "black": 30,
+  "grey": 30,
+  "red": 31,
+  "green": 32,
+  "yellow": 33,
+  "blue": 34,
+  "magenta": 35,
+  "cyan": 36,
+  "light_grey": 37,
+  "dark_grey": 90,
+  "light_red": 91,
+  "light_green": 92,
+  "light_yellow": 93,
+  "light_blue": 94,
+  "light_magenta": 95,
+  "light_cyan": 96,
+  "white": 97,
+}
+
+def allow_color() -> bool:
+  if "NO_COLOR" in os.environ or os.environ.get("TERM") == "dumb" or not hasattr(sys.stdout, "fileno"):
+    return False
+  try:
+    return os.isatty(sys.stdout.fileno())
+  except io.UnsupportedOperation:
+    return sys.stdout.isatty()
+
+def colored_(text: str, color: Color | None = None, on_color: Color | None = None) -> str:
+  if color:
+    text = f"\033[{COLOR_MAP[color]}m{text}"
+  if on_color:
+    text = f"\033[{COLOR_MAP[on_color] + 10}m{text}"
+  return text + "\033[0m"
+
+noop = lambda *args, **_: args[0]
+colored = colored_ if allow_color() else noop
+
+def print_checkpoint(should_log: bool, title: str, text: str, color: Color):
+  if should_log:
+    print(f"{colored(f" {title} ", "grey", color)} {colored(text, color)}")

checkpointer/storages/bcolz_storage.py

@@ -1,92 +1,98 @@
 import shutil
 from pathlib import Path
 from datetime import datetime
+from ..types import Storage
 
 def get_data_type_str(x):
   if isinstance(x, tuple):
-    return 'tuple'
+    return "tuple"
   elif isinstance(x, dict):
-    return 'dict'
+    return "dict"
   elif isinstance(x, list):
-    return 'list'
-  elif isinstance(x, str) or not hasattr(x, '__len__'):
-    return 'other'
+    return "list"
+  elif isinstance(x, str) or not hasattr(x, "__len__"):
+    return "other"
   else:
-    return 'ndarray'
+    return "ndarray"
 
-def get_paths(root_path, invoke_path):
-  full_path = Path(invoke_path) if root_path is None else root_path / invoke_path
-  meta_full_path = full_path.with_name(full_path.name + '_meta')
-  return full_path, meta_full_path
+def get_metapath(path: Path):
+  return path.with_name(f"{path.name}_meta")
 
-def get_collection_timestamp(config, path):
+def get_collection_timestamp(path: Path):
   import bcolz
-  _, meta_full_path = get_paths(config.root_path, path)
-  meta_data = bcolz.open(meta_full_path)[:][0]
-  return meta_data['created']
+  metapath = get_metapath(path)
+  meta_data = bcolz.open(metapath)[:][0]
+  return meta_data["created"]
 
-def get_is_expired(config, path):
-  try:
-    get_collection_timestamp(config, path)
-    return False
-  except (FileNotFoundError, EOFError):
-    return True
-
-def should_expire(config, path, expire_fn):
-  return expire_fn(get_collection_timestamp(config, path))
-
-def insert_data(path, data):
+def insert_data(path: Path, data):
   import bcolz
-  c = bcolz.carray(data, rootdir=path, mode='w')
+  c = bcolz.carray(data, rootdir=path, mode="w")
   c.flush()
 
-def store_data(config, path, data, expire_in=None):
-  full_path, meta_full_path = get_paths(config.root_path, path)
-  full_path.parent.mkdir(parents=True, exist_ok=True)
-  created = datetime.now()
-  data_type_str = get_data_type_str(data)
-  if data_type_str == 'tuple':
-    fields = list(range(len(data)))
-  elif data_type_str == 'dict':
-    fields = sorted(data.keys())
-  else:
-    fields = []
-  meta_data = {'created': created, 'data_type_str': data_type_str, 'fields': fields}
-  insert_data(meta_full_path, meta_data)
-  if data_type_str in ['tuple', 'dict']:
-    for i in range(len(fields)):
-      sub_path = f"{path} ({i})"
-      store_data(config, sub_path, data[fields[i]])
-  else:
-    insert_data(full_path, data)
-  return data
+class BcolzStorage(Storage):
+  @staticmethod
+  def is_expired(path):
+    try:
+      get_collection_timestamp(path)
+      return False
+    except (FileNotFoundError, EOFError):
+      return True
 
-def load_data(config, path):
-  import bcolz
-  full_path, meta_full_path = get_paths(config.root_path, path)
-  meta_data = bcolz.open(meta_full_path)[:][0]
-  data_type_str = meta_data['data_type_str']
-  if data_type_str in ['tuple', 'dict']:
-    fields = meta_data['fields']
-    partitions = range(len(fields))
-    data = [load_data(config, f"{path} ({i})") for i in partitions]
-    if data_type_str == 'tuple':
-      return tuple(data)
+  @staticmethod
+  def should_expire(path, expire_fn):
+    return expire_fn(get_collection_timestamp(path))
+
+  @staticmethod
+  def store_data(path, data):
+    metapath = get_metapath(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    created = datetime.now()
+    data_type_str = get_data_type_str(data)
+    if data_type_str == "tuple":
+      fields = list(range(len(data)))
+    elif data_type_str == "dict":
+      fields = sorted(data.keys())
     else:
-      return dict(zip(fields, data))
-  else:
-    data = bcolz.open(full_path)
-    if data_type_str == 'list':
-      return list(data)
-    elif data_type_str == 'other':
-      return data[0]
+      fields = []
+    meta_data = {"created": created, "data_type_str": data_type_str, "fields": fields}
+    insert_data(metapath, meta_data)
+    if data_type_str in ["tuple", "dict"]:
+      for i in range(len(fields)):
+        child_path = Path(f"{path} ({i})")
+        BcolzStorage.store_data(child_path, data[fields[i]])
+    else:
+      insert_data(path, data)
+    return data
+
+  @staticmethod
+  def load_data(path):
+    import bcolz
+    metapath = get_metapath(path)
+    meta_data = bcolz.open(metapath)[:][0]
+    data_type_str = meta_data["data_type_str"]
+    if data_type_str in ["tuple", "dict"]:
+      fields = meta_data["fields"]
+      partitions = range(len(fields))
+      data = [BcolzStorage.load_data(Path(f"{path} ({i})")) for i in partitions]
+      if data_type_str == "tuple":
+        return tuple(data)
+      else:
+        return dict(zip(fields, data))
     else:
-      return data[:]
+      data = bcolz.open(path)
+      if data_type_str == "list":
+        return list(data)
+      elif data_type_str == "other":
+        return data[0]
+      else:
+        return data[:]
 
-def delete_data(config, path):
-  full_path, meta_full_path = get_paths(config.root_path, path)
-  try:
-    shutil.rmtree(meta_full_path)
-    shutil.rmtree(full_path)
-  except FileNotFoundError:
-    pass
+  @staticmethod
+  def delete_data(path):
+    # NOTE: Not recursive
+    metapath = get_metapath(path)
+    try:
+      shutil.rmtree(metapath)
+      shutil.rmtree(path)
+    except FileNotFoundError:
+      pass

checkpointer/storages/memory_storage.py

@@ -1,18 +1,29 @@
 from datetime import datetime
+from ..types import Storage
 
 store = {}
 date_stored = {}
 
-def get_is_expired(config, path):
-  return path not in store
+class MemoryStorage(Storage):
+  @staticmethod
+  def is_expired(path):
+    return path not in store
 
-def should_expire(config, path, expire_fn):
-  return expire_fn(date_stored[path])
+  @staticmethod
+  def should_expire(path, expire_fn):
+    return expire_fn(date_stored[path])
 
-def store_data(config, path, data):
-  store[path] = data
-  date_stored[path] = datetime.now()
-  return data
+  @staticmethod
+  def store_data(path, data):
+    store[path] = data
+    date_stored[path] = datetime.now()
+    return data
 
-def load_data(config, path):
-  return store[path]
+  @staticmethod
+  def load_data(path):
+    return store[path]
+
+  @staticmethod
+  def delete_data(path):
+    del store[path]
+    del date_stored[path]

checkpointer/storages/pickle_storage.py

@@ -1,49 +1,55 @@
 import pickle
 from pathlib import Path
 from datetime import datetime
+from ..types import Storage
 
-def get_paths(root_path, invoke_path):
-  p = Path(invoke_path) if root_path is None else root_path / invoke_path
-  meta_full_path = p.with_name(p.name + '_meta.pkl')
-  pkl_full_path = p.with_name(p.name + '.pkl')
+def get_paths(path: Path):
+  meta_full_path = path.with_name(f"{path.name}_meta.pkl")
+  pkl_full_path = path.with_name(f"{path.name}.pkl")
   return meta_full_path, pkl_full_path
 
-def get_collection_timestamp(config, path):
-  meta_full_path, pkl_full_path = get_paths(config.root_path, path)
-  with meta_full_path.open('rb') as file:
+def get_collection_timestamp(path: Path):
+  meta_full_path, _ = get_paths(path)
+  with meta_full_path.open("rb") as file:
     meta_data = pickle.load(file)
-    return meta_data['created']
+    return meta_data["created"]
 
-def get_is_expired(config, path):
-  try:
-    get_collection_timestamp(config, path)
-    return False
-  except (FileNotFoundError, EOFError):
-    return True
+class PickleStorage(Storage):
+  @staticmethod
+  def is_expired(path):
+    try:
+      get_collection_timestamp(path)
+      return False
+    except (FileNotFoundError, EOFError):
+      return True
 
-def should_expire(config, path, expire_fn):
-  return expire_fn(get_collection_timestamp(config, path))
+  @staticmethod
+  def should_expire(path, expire_fn):
+    return expire_fn(get_collection_timestamp(path))
 
-def store_data(config, path, data):
-  created = datetime.now()
-  meta_data = {'created': created}
-  meta_full_path, pkl_full_path = get_paths(config.root_path, path)
-  pkl_full_path.parent.mkdir(parents=True, exist_ok=True)
-  with pkl_full_path.open('wb') as file:
-    pickle.dump(data, file, -1)
-  with meta_full_path.open('wb') as file:
-    pickle.dump(meta_data, file, -1)
-  return data
+  @staticmethod
+  def store_data(path, data):
+    created = datetime.now()
+    meta_data = {"created": created} # TODO: this should just be a JSON or binary dump of the unix timestamp and other metadata - not pickle
+    meta_full_path, pkl_full_path = get_paths(path)
+    pkl_full_path.parent.mkdir(parents=True, exist_ok=True)
+    with pkl_full_path.open("wb") as file:
+      pickle.dump(data, file, -1)
+    with meta_full_path.open("wb") as file:
+      pickle.dump(meta_data, file, -1)
+    return data
 
-def load_data(config, path):
-  _, full_path = get_paths(config.root_path, path)
-  with full_path.open('rb') as file:
-    return pickle.load(file)
+  @staticmethod
+  def load_data(path):
+    _, full_path = get_paths(path)
+    with full_path.open("rb") as file:
+      return pickle.load(file)
 
-def delete_data(config, path):
-  meta_full_path, pkl_full_path = get_paths(config.root_path, path)
-  try:
-    meta_full_path.unlink()
-    pkl_full_path.unlink()
-  except FileNotFoundError:
-    pass
+  @staticmethod
+  def delete_data(path):
+    meta_full_path, pkl_full_path = get_paths(path)
+    try:
+      meta_full_path.unlink()
+      pkl_full_path.unlink()
+    except FileNotFoundError:
+      pass

checkpointer/types.py

@@ -0,0 +1,19 @@
+from typing import Callable, Protocol, Any
+from pathlib import Path
+from datetime import datetime
+
+class Storage(Protocol):
+  @staticmethod
+  def is_expired(path: Path) -> bool: ...
+
+  @staticmethod
+  def should_expire(path: Path, expire_fn: Callable[[datetime], bool]) -> bool: ...
+
+  @staticmethod
+  def store_data(path: Path, data: Any) -> Any: ...
+
+  @staticmethod
+  def load_data(path: Path) -> Any: ...
+
+  @staticmethod
+  def delete_data(path: Path) -> None: ...

checkpointer/utils.py

@@ -1,10 +1,9 @@
 import types
 
-def unwrap_func(func):
-  if hasattr(func, '__wrapped__'):
-    return unwrap_func(func.__wrapped__)
-  else:
-    return func
+def unwrap_fn[T](fn: T) -> T:
+  while hasattr(fn, "__wrapped__"):
+    fn = getattr(fn, "__wrapped__")
+  return fn
 
 @types.coroutine
 def coroutine_as_generator(coroutine):

checkpointer-2.0.0.dist-info/METADATA

@@ -0,0 +1,212 @@
+Metadata-Version: 2.3
+Name: checkpointer
+Version: 2.0.0
+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
+License: Copyright 2024 Hampus Hallman
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+Requires-Python: >=3.12
+Requires-Dist: relib
+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. 🚀
+
+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.
+- **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.
+
+### 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.
+
+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.
+
+### Example: Cache Invalidation by Function Dependencies
+
+```python
+def multiply(a, b):
+    return a * b
+
+@checkpoint
+def helper(x):
+    return multiply(x + 1, 2)
+
+@checkpoint
+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.
+
+---
+
+## Installation
+
+```bash
+pip install checkpointer
+```
+
+---
+
+## Quick Start
+
+```python
+from checkpointer import checkpoint
+
+@checkpoint
+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
+```
+
+---
+
+## Parameterization
+
+### Global Configuration
+
+You can configure a custom `Checkpointer`:
+
+```python
+from checkpointer import checkpoint
+
+checkpoint = checkpoint(format="memory", root_path="/tmp/checkpoints")
+```
+
+Extend this configuration by calling itself again:
+
+```python
+extended_checkpoint = checkpoint(format="pickle", verbosity=0)
+```
+
+### 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
+```
+
+### Layered Caching
+
+```python
+IS_DEVELOPMENT = True  # Toggle based on environment
+
+dev_checkpoint = checkpoint(when=IS_DEVELOPMENT)
+
+@checkpoint(format="memory")
+@dev_checkpoint
+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.
+
+---
+
+## Usage
+
+### Force Recalculation
+Use `rerun` to force a recalculation and overwrite the stored checkpoint:
+
+```python
+result = expensive_function.rerun(4)
+```
+
+### Bypass Checkpointer
+Use `fn` to directly call the original, undecorated function:
+
+```python
+result = expensive_function.fn(4)
+```
+
+### Retrieve Stored Checkpoints
+Access stored results without recalculating:
+
+```python
+stored_result = expensive_function.get(4)
+```
+
+---
+
+## Configuration Options
+
+| Option         | Type                                | Default     | Description                                 |
+|----------------|-------------------------------------|-------------|---------------------------------------------|
+| `format`       | `"pickle"`, `"memory"`, `Storage`   | `"pickle"`  | Storage backend format.                     |
+| `root_path`    | `Path`, `str`, or `None`            | User Cache  | Root directory for storing checkpoints.     |
+| `when`         | `bool`                              | `True`      | Enable or disable checkpointing.            |
+| `verbosity`    | `0` or `1`                          | `1`         | Logging verbosity.                          |
+| `path`         | `str` or `Callable[..., str]`       | `None`      | Custom path for checkpoint storage.         |
+| `should_expire`| `Callable[[datetime], bool]`        | `None`      | Custom expiration logic.                    |
+
+---
+
+## Full Example
+
+```python
+import asyncio
+from checkpointer import checkpoint
+
+@checkpoint
+def compute_square(n: int) -> int:
+    print(f"Computing {n}^2...")
+    return n ** 2
+
+@checkpoint(format="memory")
+async def async_compute_sum(a: int, b: int) -> int:
+    await asyncio.sleep(1)
+    return a + b
+
+async def main():
+    result1 = compute_square(5)
+    print(result1)
+
+    result2 = await async_compute_sum(3, 7)
+    print(result2)
+
+    result3 = async_compute_sum.get(3, 7)
+    print(result3)
+
+asyncio.run(main())
+```

checkpointer-2.0.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+checkpointer/__init__.py,sha256=TODGBGbZYBJ5LIpz5t6tnQNJ7ODPRUvXjv3Ooqb8-cc,357
+checkpointer/checkpoint.py,sha256=V8JL8ibmMeqZjLjaigeAWa-8c948VfIiYqO8t5OFk48,4812
+checkpointer/function_body.py,sha256=92mnTY9d_JhKnKugeySYRP6qhU4fH6F6zesb7h2pEi0,1720
+checkpointer/print_checkpoint.py,sha256=wHC2xWNwNfFhRHyhrmLkadYoyThRTJWiox3NjgE9Ubc,1369
+checkpointer/types.py,sha256=yoNPnN_QJHfyK_Gs8c0SoywHHDUlU7uhKqPPTTWjRTE,469
+checkpointer/utils.py,sha256=UrQt689UHUjl7kXpTbUCGkHUgQZllByX2rbuvZdt9vk,368
+checkpointer/storages/bcolz_storage.py,sha256=5hbJB0VJ2k-FHf7rItywMXP74WT-JTqeNK5N8yftcnw,2647
+checkpointer/storages/memory_storage.py,sha256=5ITKjh_bVNfj1C6pcyMgB4YU4sy6jOLlvH0_3Pl1Elo,558
+checkpointer/storages/pickle_storage.py,sha256=ipXG2dht8YQAXJFEK5-OwOb8xP8ij_v7K0Qu5Xz9aVE,1622
+checkpointer-2.0.0.dist-info/METADATA,sha256=5DZmJ0rMnPeRZ9b5REXv1h7z8tbj0llgTO9yp8xEbnQ,7561
+checkpointer-2.0.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+checkpointer-2.0.0.dist-info/licenses/LICENSE,sha256=0cmUKqBotzbBcysIexd52AhjwbphhlGYiWbvg5l2QAU,1054
+checkpointer-2.0.0.dist-info/RECORD,,

{checkpointer-1.1.0.dist-info → checkpointer-2.0.0.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.42.0)
+Generator: hatchling 1.26.3
 Root-Is-Purelib: true
 Tag: py3-none-any
-

checkpointer-1.1.0.dist-info/LICENSE.txt → checkpointer-2.0.0.dist-info/licenses/LICENSE

@@ -1,4 +1,4 @@
-Copyright 2018 Hampus Hallman
+Copyright 2024 Hampus Hallman
 
 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 

checkpointer/storage.py

@@ -1,60 +0,0 @@
-import inspect
-from termcolor import colored
-from .storages import memory_storage, pickle_storage, bcolz_storage
-
-storages = {
-  'memory': memory_storage,
-  'pickle': pickle_storage,
-  'bcolz': bcolz_storage,
-}
-
-initialized_storages = set()
-
-def create_logger(should_log):
-  def log(color, title, text):
-    if should_log:
-      title_log = colored(f' {title} ', 'grey', 'on_' + color)
-      rest_log = colored(text, color)
-      print(title_log + ' ' + rest_log)
-  return log
-
-def get_storage(storage):
-  if type(storage) == str:
-    storage = storages[storage]
-  if storage not in initialized_storages:
-    if hasattr(storage, 'initialize'):
-      storage.initialize()
-    initialized_storages.add(storage)
-  return storage
-
-async def store_on_demand(get_data, name, config, force=False, should_expire=None):
-  storage = get_storage(config.format)
-  should_log = storage != memory_storage and config.verbosity != 0
-  log = create_logger(should_log)
-  refresh = force \
-    or storage.get_is_expired(config, name) \
-    or (should_expire and storage.should_expire(config, name, should_expire))
-
-  if refresh:
-    log('blue', 'MEMORIZING', name)
-    data = get_data()
-    if inspect.iscoroutine(data):
-      data = await data
-    return storage.store_data(config, name, data)
-  else:
-    try:
-      data = storage.load_data(config, name)
-      log('green', 'REMEMBERED', name)
-      return data
-    except (EOFError, FileNotFoundError):
-      log('yellow', 'CORRUPTED', name)
-      storage.delete_data(config, name)
-      result = await store_on_demand(get_data, name, config, force, should_expire)
-      return result
-
-def read_from_store(name, config, storage='pickle'):
-  storage = get_storage(storage)
-  try:
-    return storage.load_data(config, name)
-  except:
-    return None

checkpointer-1.1.0.dist-info/METADATA

@@ -1,12 +0,0 @@
-Metadata-Version: 2.1
-Name: checkpointer
-Version: 1.1.0
-Home-page: https://github.com/Reddan/checkpointer
-Author: Hampus Hallman
-Author-email: me@hampushallman.com
-License: MIT
-Requires-Python: ~=3.5
-License-File: LICENSE.txt
-Requires-Dist: relib
-Requires-Dist: termcolor
-

checkpointer-1.1.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-checkpointer/__init__.py,sha256=_RYcKsZbeUf08KZ-DcXlNn4eAMWh2LN9o-KvchYVmmk,380
-checkpointer/checkpoint.py,sha256=FyL78HvAvPtgl-esiAkt-CdekT18J2Sh0SriMtX4QLc,2367
-checkpointer/function_body.py,sha256=aitzBFoEIo_E-aL2oe4Nr_T_Oy8ehKInGb1XYjeLzX8,1703
-checkpointer/storage.py,sha256=Ofuh0dKF5vk4_B4djt3Q6qyZhIO5f59uCNCZjMrto0U,1782
-checkpointer/utils.py,sha256=Kzh2qXxq2Lgjiqfv6KcOwXdw7sh70uCn4a_uXYwLhiM,377
-checkpointer/storages/bcolz_storage.py,sha256=Yk7FI75noe9hZBWVFIRetiFSR7tkzbryYlBmxX-lVlw,2728
-checkpointer/storages/memory_storage.py,sha256=S4SgKSApbQE-pxxKRWLNJqyZMRQwaw5-N0DOIsZM7mE,364
-checkpointer/storages/pickle_storage.py,sha256=zcnX1GG6XPHvVxi7gCab5oFxKoz5E7LZHYH74VL1hkY,1542
-checkpointer-1.1.0.dist-info/LICENSE.txt,sha256=2c7g4mni-RUemFGkk6GnoFwknh-leF04BF_J_3gp4sg,1054
-checkpointer-1.1.0.dist-info/METADATA,sha256=_isCKu9zdEI1Z-umSB6mr-X7lEVBVjlU_QvRNYolUNQ,273
-checkpointer-1.1.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-checkpointer-1.1.0.dist-info/top_level.txt,sha256=uF0eyHpShnsHI3sobErnhQ8LWCT8DPViqAznKeTaZlw,13
-checkpointer-1.1.0.dist-info/RECORD,,

checkpointer-1.1.0.dist-info/top_level.txt

@@ -1 +0,0 @@
-checkpointer