"""Content-addressed ledger, so a heavy call runs once and is remembered.
The storage half of wallet-helper. A call is identified by a key built from a
namespace plus a payload (the arguments, file, or bytes that determine the
result). The ledger stores the result under that key, so an identical call is
served from disk instead of running again, across process restarts.
The default store is a directory with one JSON file per entry: local, no server,
easy to inspect and to delete. It is content-addressed, so a renamed input file
still hits and two different inputs never collide.
Author
------
Warith HARCHAOUI, https://linkedin.com/in/warith-harchaoui
"""
from __future__ import annotations
import functools
import glob as glob_module
import json
import os
import tempfile
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Iterator, Protocol, runtime_checkable
import os_helper as osh
try:
import fcntl # POSIX advisory file locks (macOS, Linux)
except ImportError: # pragma: no cover - Windows has no fcntl
fcntl = None # type: ignore[assignment]
# Default store location, outside any repo so cached results are not committed by
# accident. Overridable per instance or through the environment variable.
_DEFAULT_DIR = Path(os.environ.get("WALLET_HELPER_DIR", str(Path.home() / ".cache" / "wallet-helper")))
# One in-process lock per store directory, shared across Ledger instances, so
# threads in the same process never interleave a read-modify-write of an entry.
_DIR_LOCKS: dict[str, threading.Lock] = {}
_DIR_LOCKS_GUARD = threading.Lock()
def _atomic_write(path: Path, text: str) -> None:
"""Write ``text`` to ``path`` atomically, so a reader never sees a half file.
The data goes to a temporary file in the same directory, then ``os.replace``
swaps it into place in one step (atomic on POSIX and Windows). A concurrent
writer or a crash mid-write can never leave a truncated or corrupt entry.
"""
fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix=path.name + ".", suffix=".tmp")
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
handle.write(text)
os.replace(tmp, path)
except BaseException:
# Never leave the temporary file behind if the swap did not happen.
if os.path.exists(tmp):
os.unlink(tmp)
raise
@contextmanager
def _entry_lock(directory: Path) -> Iterator[None]:
"""Serialize read-modify-write on a store directory, within and across processes.
Holds a per-directory :class:`threading.Lock` (so threads in this process
take turns) and, where available, an exclusive advisory lock on a small lock
file (so separate processes take turns too). On platforms without ``fcntl``
the cross-process guard is skipped and only the in-process lock applies.
"""
key = str(directory)
with _DIR_LOCKS_GUARD:
lock = _DIR_LOCKS.setdefault(key, threading.Lock())
with lock:
if fcntl is None:
yield # Windows: in-process lock only (documented on register_hit)
return
directory.mkdir(parents=True, exist_ok=True)
handle = open(directory / ".wallet-helper.lock", "a+")
try:
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
yield
finally:
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
handle.close()
[docs]
@runtime_checkable
class LedgerLike(Protocol):
"""The storage contract every ledger backend fulfils.
The default :class:`Ledger` keeps one JSON file per entry, which is ideal for
a single process. :class:`wallet_helper.sqlite_ledger.SqliteLedger` keeps
everything in one SQLite file for a shared, concurrency-safe store. A
:class:`wallet_helper.guard.Wallet` and the command-line tools accept either.
"""
@property
def location(self) -> str:
"""A readable pointer to where entries live (a directory or a file)."""
...
[docs]
def has(self, key: str) -> bool:
"""Return ``True`` if a result is already stored for ``key``."""
...
[docs]
def get(self, key: str) -> Any | None:
"""Return the stored result for ``key``, or ``None`` if absent."""
...
[docs]
def get_record(self, key: str) -> dict | None:
"""Return the full stored record, or ``None`` if absent."""
...
[docs]
def put(self, key: str, result: Any, *, ttl: float | None = None) -> None:
"""Store ``result`` for ``key`` (overwrites), expiring after ``ttl`` seconds."""
...
[docs]
def register_hit(self, key: str) -> None:
"""Count one reuse of the cached result for ``key`` (no-op if absent)."""
...
[docs]
def stats(self, namespace: str | None = None) -> dict:
"""Return ``{entries, hits}`` for the whole store or one namespace."""
...
[docs]
def clear(self, namespace: str | None = None) -> None:
"""Remove entries, all of them or just one namespace (irreversible)."""
...
[docs]
def evict(self, *, max_entries: int | None = None, older_than: float | None = None) -> int:
"""Prune entries and return how many were removed (see :meth:`Ledger.evict`)."""
...
# Markers standing in for a content hash inside a canonicalised payload. The
# leading NUL means a real path or argument string never begins this way, and any
# user string that somehow does is escaped (see _canonical) so it can never be
# mistaken for a marker. Order matters: _ESC is checked first when escaping.
_ESC = "\x00wallet_helper.esc:"
_FILE_MARK = "\x00wallet_helper.file:"
_BYTES_MARK = "\x00wallet_helper.bytes:"
_SET_MARK = "\x00wallet_helper.set:"
_KEY_MARK = "\x00wallet_helper.key:"
_OBJ_MARK = "\x00wallet_helper.obj:"
_MARKS = (_ESC, _FILE_MARK, _BYTES_MARK, _SET_MARK, _KEY_MARK, _OBJ_MARK)
def _has_content_repr(value: Any) -> bool:
"""Return ``True`` if ``str(value)`` reflects content, not just identity.
An object that overrides ``__str__`` or ``__repr__`` (``enum``, ``Decimal``,
``datetime``, ``uuid``, numpy scalars, dataclasses, ...) has a meaningful text
form we can key on. One that inherits both defaults would only stringify to
``<Class at 0x...>``, whose heap address changes every run.
"""
cls = type(value)
return cls.__str__ is not object.__str__ or cls.__repr__ is not object.__repr__
_MISSING = object()
def _safe_getattr(value: Any, name: str) -> Any:
"""Return ``value.name`` or :data:`_MISSING`, swallowing any error.
A custom ``__getattr__`` may raise something other than ``AttributeError``;
probing an attribute during key construction must never propagate that.
"""
try:
return getattr(value, name)
except Exception:
return _MISSING
def _object_state(value: Any) -> dict | None:
"""Return an object's own attributes (``__dict__`` and ``__slots__``), or ``None``.
This is the content we key an opaque object on, instead of its address, so the
same logical object hashes the same across processes and two objects with
different state never collide. Both attribute mechanisms are merged, so a class
that uses one, the other, or both is covered. ``__slots__`` is collected from
every class in the MRO, not just ``type(value)``: a subclass's ``__slots__``
only lists the names it adds, so reading just that tuple would silently drop
any slot inherited from a base class, and two instances differing only in that
inherited slot would then key alike (a real collision, not just a miss).
``__weakref__`` is a runtime plumbing slot, not user state, so it is skipped
like ``__dict__``. Attribute reads are guarded, so a misbehaving
``__getattr__`` cannot crash key construction.
"""
state: dict = {}
own = _safe_getattr(value, "__dict__")
if isinstance(own, dict) and own:
state.update(own)
for klass in type(value).__mro__:
slots = klass.__dict__.get("__slots__", ())
if isinstance(slots, str):
slots = (slots,)
for name in slots:
if name in ("__dict__", "__weakref__"):
continue
attr = _safe_getattr(value, name)
if attr is not _MISSING:
state[name] = attr
return state or None
def _file_path(value: Any) -> str | None:
"""Return the string path if ``value`` names an existing file, else ``None``.
A path reaches us in several forms: a plain ``str``, a :class:`pathlib.Path`,
or any other :class:`os.PathLike` (``__fspath__``). All are accepted and
resolved with :func:`os.fspath`. ``bytes`` are excluded on purpose: they are
hashed as raw content, not treated as a filesystem path. Never raises: a value
that cannot be resolved to an existing file (a bad ``__fspath__``, an embedded
NUL, an over-long string) is simply reported as not a file.
"""
if isinstance(value, (bytes, bytearray)) or not isinstance(value, (str, os.PathLike)):
return None
try:
path = os.fspath(value)
if isinstance(path, bytes):
path = path.decode()
return path if osh.file_exists(path) else None
except Exception:
# Any failure resolving or stat-ing the value means it is not a file leaf;
# key construction must never crash before the wrapped call even runs.
return None
def _canonical(value: Any, seen: dict[int, int] | None = None) -> Any:
"""Replace file-path, ``bytes``, and opaque-object leaves with stable markers.
A path argument is usually one item inside a call's
``{"args": ..., "kwargs": ...}`` payload, not the whole payload. Walking the
structure lets us key such a path by the file's *bytes* rather than its text,
so two identical files at different paths (or one file later renamed) share a
single cache entry, and two different files never collide even when their
paths look alike. Leaves that are neither paths nor ``bytes`` pass through
untouched, so a payload with no file leaves serialises exactly as before.
A user string that itself begins with a marker prefix is escaped, so a crafted
argument can never forge a marker and collide with a real one. Dict *keys* are
normalised through :func:`_canonical_key`: a plain ``str`` key keeps its text
(kwarg names, the common case), while any other key type is encoded
structurally. ``seen`` maps the id of each container and object on the current
walk path to its depth, so a self-referential graph resolves to a cycle marker
instead of recursing without end. The depth is recorded in the marker so a
back-edge to one ancestor is told apart from a back-edge to another.
"""
if isinstance(value, (bytes, bytearray, memoryview)):
# Every byte-like type is keyed by its content, and hashes the same
# whatever wrapper it arrived in. latin-1 maps bytes to text losslessly.
return _BYTES_MARK + osh.hash_string(bytes(value).decode("latin-1"))
path = _file_path(value)
if path is not None:
return _FILE_MARK + osh.hashfile(path)
if isinstance(value, str):
# Keep user data and synthesised markers in disjoint spaces.
return _ESC + value if value.startswith(_MARKS) else value
if value is None or isinstance(value, (int, float, bool)):
return value # immutable scalar: no file, no cycle, hashes as itself
# From here every value is a container or an object that could form a cycle.
marker = id(value)
if seen is None:
seen = {}
if marker in seen:
# Back-edge: encode which ancestor (by its depth on this path) it targets,
# so two graphs that differ only in the back-edge destination stay distinct.
return [_OBJ_MARK, "cycle", seen[marker]]
seen[marker] = len(seen) # depth along the current walk path
try:
if isinstance(value, dict):
return {_canonical_key(k, seen): _canonical(v, seen) for k, v in value.items()}
if isinstance(value, (list, tuple)):
return [_canonical(v, seen) for v in value]
if isinstance(value, (set, frozenset)):
# Unordered: canonicalise members then sort by their JSON form so the
# same set always hashes the same, and a file inside a set is content-
# addressed. _SET_MARK keeps a set distinct from a list of equal members
# (a user list starting with _SET_MARK is escaped, so it cannot forge one).
members = sorted((_canonical(v, seen) for v in value), key=_dumps)
return [_SET_MARK, members]
if isinstance(value, functools.partial):
# A partial has an address-bearing repr but a structural identity: the
# wrapped callable plus its bound args and keywords, each canonicalised.
return [
_OBJ_MARK,
"partial",
_canonical(value.func, seen),
_canonical(list(value.args), seen),
_canonical(dict(value.keywords), seen),
]
qualname = _safe_getattr(value, "__qualname__")
if qualname is not _MISSING and (callable(value) or isinstance(value, type)):
# A function, method, class, or builtin: its repr embeds an address, but
# its identity is a stable (module, qualified name). Key on that so a
# callback passed as an argument dedups across processes.
module = _safe_getattr(value, "__module__")
return [_OBJ_MARK, "def", module if isinstance(module, str) else "", qualname]
if not _has_content_repr(value):
# An opaque object (identity-only repr) would otherwise leak its heap
# address into the key. Key it by its own state instead, canonicalised,
# so it is deterministic across processes and distinguishes instances.
# The type name keeps unrelated types apart; _OBJ_MARK blocks forgery.
# Volatile handles held in that state should still be dropped with
# ignore=/key=.
state = _object_state(value)
type_name = f"{type(value).__module__}.{type(value).__qualname__}"
return [_OBJ_MARK, type_name, _canonical(state, seen) if state is not None else None]
# A value with a content-bearing str (enum, Decimal, datetime, numpy, ...):
# leave it for json.dumps (its default=str renders it).
return value
finally:
del seen[marker]
def _canonical_key(key: Any, seen: dict[int, int] | None = None) -> str:
"""Return a stable string form of a dict key so the dict always serialises.
A plain string key is kept as is (and escaped only if it resembles a marker),
so the common payload hashes exactly as a plain ``json.dumps`` would. A
non-string key (a tuple, ``bytes``, an enum, or one of several mixed with
strings) would otherwise make ``json.dumps`` raise or fail to sort; it is
encoded as the JSON of its canonical form under a distinct ``_KEY_MARK`` so it
never collides with a string key of the same text (``{1: v}`` vs ``{"1": v}``).
"""
if isinstance(key, str):
return _ESC + key if key.startswith(_MARKS) else key
return _KEY_MARK + _dumps(_canonical(key, seen))
def _dumps(obj: Any) -> str:
"""Serialise an already-canonical object to a stable, key-sorted JSON string.
By this point :func:`_canonical` has turned every opaque object into its
state, so ``default=str`` only ever renders values with a content-bearing
``str`` (``enum``, ``Decimal``, ``datetime``, and the like).
"""
return json.dumps(obj, sort_keys=True, default=str, ensure_ascii=False)
def _digest(payload: Any) -> str:
"""Return a stable content hash for a key payload.
Delegates to os_helper's hashing, so wallet-helper reuses the suite's tested
content-addressing instead of rolling its own. The payload is first
canonicalised (:func:`_canonical`), which resolves every file-path and
``bytes`` leaf to its content hash *wherever it sits* (top level or nested in
the arguments), then the whole canonical form is hashed as key-sorted JSON. So
a file reached by a different path still hits the cache, and one rule governs
top-level and nested values alike.
Parameters
----------
payload : Any
A path to an existing file as a ``str`` or any :class:`os.PathLike`, raw
``bytes``, or any JSON-serialisable value, at any depth. Files and
``bytes`` are keyed by content (in disjoint spaces: a path and equal raw
bytes do not alias); everything else by its canonical JSON.
Returns
-------
str
A fixed-length hex digest.
"""
# One rule for every depth: canonicalise (files and bytes -> content markers),
# then hash the key-sorted JSON. Order-independent and enum-tolerant.
return osh.hash_string(_dumps(_canonical(payload)))
[docs]
def is_fresh(record: dict, *, now: float | None = None) -> bool:
"""Return ``True`` if a record has not expired.
Parameters
----------
record : dict
A stored record, which may carry an ``expires_at`` timestamp.
now : float, optional
The current time; defaults to :func:`time.time`. Passing it lets a caller
judge many records against one instant.
Returns
-------
bool
``True`` when there is no expiry, or the expiry is still in the future.
"""
expires_at = record.get("expires_at")
if expires_at is None:
return True
return (now if now is not None else time.time()) < expires_at
[docs]
def make_key(namespace: str, payload: Any) -> str:
"""Build a ledger key from a ``namespace`` and a content ``payload``.
Parameters
----------
namespace : str
A scope for the call, for example ``"transcribe"`` or ``"openai.chat"``,
so unrelated calls do not collide even if their payloads hash alike.
payload : Any
See :func:`_digest`.
Returns
-------
str
``"<namespace>_<hash>"``, safe to use as a filename.
Notes
-----
A ``str`` or :class:`os.PathLike` argument that names an existing file is
keyed by the file's content, by design, so the same file reached by a
different path still hits (this applies to a path used as an argument or as a
value; a plain ``str`` used as a dict *key* keeps its text). Two consequences
worth knowing:
- A string you meant as plain data (a title, an id) that *happens* to match an
existing file in the working directory is content-addressed too. If that is
not what you want, wrap the value so it is not a bare path, or exclude the
argument with ``ignore=`` / a ``key=`` builder.
- If a file is deleted in the brief moment between detection and hashing, the
key falls back to the path text for that call. The wrapped function would
fail to read the missing file anyway, so no stale result is served.
Examples
--------
>>> make_key("demo", {"b": 2, "a": 1}) == make_key("demo", {"a": 1, "b": 2})
True
"""
return f"{namespace}_{_digest(payload)}"
# Characters that are legal inside a ledger key (they come from a namespace or
# a qualname such as "outer.<locals>.inner") but illegal in a Windows filename.
# The JSON store turns a key into a filename, so on Windows an un-encoded key
# with any of these raises OSError. Percent-encoding them (and control chars,
# and "%" itself) keeps the store working on Windows too.
_FORBIDDEN_IN_FILENAME = '<>:"/\\|?*'
def _safe_filename(name: str) -> str:
"""Return ``name`` with Windows-illegal filename characters percent-encoded.
A ledger key is ``"<namespace>_<hash>"``. The hash is hex (always safe), but
the namespace defaults to ``module.qualname`` and, for a nested function or a
lambda, contains ``<locals>`` / ``<lambda>``; ``<`` and ``>`` (and ``: " / \\
| ? *`` and control characters) are illegal in a Windows filename, so the
JSON store crashed there. This encoding is a no-op for a key that has none of
them, so the common filename is byte-identical to before; and it is injective
(``%`` is itself encoded), so two distinct keys never collide on one file.
"""
out: list[str] = []
for ch in name:
if ch == "%" or ch in _FORBIDDEN_IN_FILENAME or ord(ch) < 0x20:
out.append(f"%{ord(ch):02X}")
else:
out.append(ch)
return "".join(out)
def _namespace_glob(namespace: str) -> str:
"""Return a glob pattern matching only entries under ``namespace``.
A namespace is normally a ``module.qualname`` and never contains a glob
metacharacter, but nothing stops a custom one from including ``*``, ``?``,
or ``[``. Left unescaped, ``stats("a*b")`` or ``clear("a*b")`` would match
entries under an unrelated namespace whose name merely fits the pattern. The
namespace is first mapped through :func:`_safe_filename` (the same transform
the on-disk names use), then :func:`glob.escape` neutralises any remaining
glob metacharacter; the trailing ``_*.json`` stays a real wildcard. Because
the transform is per-character and leaves ``_`` and hex alone, this matches
exactly the files whose key begins with ``namespace_``.
"""
return f"{glob_module.escape(_safe_filename(namespace))}_*.json"
[docs]
class Ledger:
"""A directory-backed store of results, one JSON file per entry.
Parameters
----------
cache_dir : str or pathlib.Path, optional
Where entries live. Defaults to ``$WALLET_HELPER_DIR`` then
``~/.cache/wallet-helper``.
max_entries : int, optional
A size cap. When set, each :meth:`put` also evicts down to the newest
``max_entries`` entries, so the store cannot grow without bound. ``None``
(default) means no automatic bound; call :meth:`evict` yourself.
Examples
--------
>>> import os_helper as osh
>>> with osh.temporary_folder() as tmp:
... Ledger(tmp).has("demo_x")
False
"""
def __init__(self, cache_dir: str | Path | None = None, max_entries: int | None = None) -> None:
self.dir = Path(cache_dir) if cache_dir is not None else _DEFAULT_DIR
self.max_entries = max_entries
@property
def location(self) -> str:
"""The directory that holds the JSON entries (for display)."""
return str(self.dir)
def _path(self, key: str) -> Path:
"""Absolute path of the JSON entry for ``key`` (filename made Windows-safe)."""
return self.dir / f"{_safe_filename(key)}.json"
[docs]
def has(self, key: str) -> bool:
"""Return ``True`` if a result is already stored for ``key``."""
return osh.file_exists(str(self._path(key)))
[docs]
def get_record(self, key: str) -> dict | None:
"""Return the full stored record, or ``None`` if absent."""
path = self._path(key)
if not path.exists():
return None
return json.loads(path.read_text(encoding="utf-8"))
[docs]
def get(self, key: str) -> Any | None:
"""Return just the stored result for ``key``, or ``None`` if absent."""
record = self.get_record(key)
return None if record is None else record["result"]
[docs]
def put(self, key: str, result: Any, *, ttl: float | None = None) -> None:
"""Store ``result`` for ``key`` (overwrites any previous entry).
Parameters
----------
key : str
The ledger key (see :func:`make_key`).
result : Any
The JSON-serialisable result to store.
ttl : float, optional
Seconds until the entry is considered stale. ``None`` (default) means
it never expires. A stale entry is treated as a miss on the next call
and is removed by :meth:`evict`.
"""
osh.make_directory(str(self.dir))
now = time.time()
record = {
"key": key,
"result": result,
"created_at": now,
"expires_at": now + ttl if ttl is not None else None,
"hits": 0, # incremented every time the cached result is reused
}
# Atomic swap: overwriting an entry never exposes a partial file, so a
# concurrent writer or a crash cannot corrupt it.
_atomic_write(self._path(key), json.dumps(record, ensure_ascii=False, indent=2))
if self.max_entries is not None:
self.evict(max_entries=self.max_entries)
[docs]
def register_hit(self, key: str) -> None:
"""Count one reuse of the cached result for ``key`` (no-op if absent).
The read-modify-write runs under a lock so concurrent hits do not lose a
count: threads in this process take turns, and separate processes take
turns too where advisory file locks are available (POSIX). On Windows the
cross-process count is best-effort; use :class:`SqliteLedger` for an exact
count under heavy multi-process concurrency.
"""
with _entry_lock(self.dir):
record = self.get_record(key)
if record is None:
return
record["hits"] = int(record.get("hits", 0)) + 1
_atomic_write(self._path(key), json.dumps(record, ensure_ascii=False, indent=2))
[docs]
def clear(self, namespace: str | None = None) -> None:
"""Delete entries, all or just one ``namespace`` (irreversible).
With a namespace, only its entries are removed (this is what a memoized
function's ``cache_clear()`` calls). Without one, the whole directory is
removed and recreated lazily on the next :meth:`put`.
"""
if namespace is None:
osh.remove_directory(str(self.dir))
return
# Keys are "<namespace>_<hash>", one file each; unlink the matches.
osh.remove_files([str(p) for p in self.dir.glob(_namespace_glob(namespace))])
[docs]
def stats(self, namespace: str | None = None) -> dict:
"""Count stored entries and their reuses, for the store or one namespace.
Parameters
----------
namespace : str, optional
Restrict the count to entries under this namespace. ``None`` (default)
counts the whole ledger.
Returns
-------
dict
``{"entries": int, "hits": int}`` where ``hits`` is how many times a
cached result was reused, that is, how many real calls were saved.
"""
entries = hits = 0
pattern = "*.json" if namespace is None else _namespace_glob(namespace)
for path in self.dir.glob(pattern):
record = json.loads(path.read_text(encoding="utf-8"))
entries += 1
hits += int(record.get("hits", 0))
return {"entries": entries, "hits": hits}
[docs]
def evict(self, *, max_entries: int | None = None, older_than: float | None = None) -> int:
"""Prune entries and return how many were removed.
Expired entries (past their ``ttl``) are always removed. In addition:
Parameters
----------
max_entries : int, optional
Keep only the newest ``max_entries`` by creation time; remove the
rest. This is a simple size cap.
older_than : float, optional
Remove entries created more than this many seconds ago.
Returns
-------
int
The number of entries removed.
"""
now = time.time()
# Read every entry once, with its creation time, so we can rank and prune.
items = []
for path in self.dir.glob("*.json"):
record = json.loads(path.read_text(encoding="utf-8"))
items.append((path, float(record.get("created_at", 0.0)), record))
doomed: set[Path] = set()
for path, created_at, record in items:
if not is_fresh(record, now=now):
doomed.add(path) # expired entries always go
elif older_than is not None and (now - created_at) > older_than:
doomed.add(path)
if max_entries is not None:
survivors = sorted((it for it in items if it[0] not in doomed), key=lambda it: it[1], reverse=True)
for path, _created_at, _record in survivors[max_entries:]:
doomed.add(path) # keep the newest max_entries, drop the older tail
osh.remove_files([str(p) for p in doomed])
return len(doomed)