diff options
| author | Mikka Petterson <m.petterson@student.maastrichtuniversity.nl> | 2026-10-14 07:45:20 +0200 |
|---|---|---|
| committer | Mikka Petterson <m.petterson@student.maastrichtuniversity.nl> | 2026-10-14 07:45:20 +0200 |
| commit | 0bcc5be35b6a99109bd04925c22890347c9d6e88 (patch) | |
| tree | fdd88728ad1747268b0a02aba0de08a3af6a52b2 /src/teradata/cache/file.py | |
Diffstat (limited to 'src/teradata/cache/file.py')
| -rw-r--r-- | src/teradata/cache/file.py | 86 |
1 files changed, 86 insertions, 0 deletions
diff --git a/src/teradata/cache/file.py b/src/teradata/cache/file.py new file mode 100644 index 0000000..ebb9bda --- /dev/null +++ b/src/teradata/cache/file.py @@ -0,0 +1,86 @@ +lazy import ast +lazy from collections.abc import Callable +lazy from functools import wraps +lazy from inspect import getsource, signature, unwrap +lazy from pathlib import Path +lazy from sys import version_info +lazy from textwrap import dedent +lazy from types import FunctionType +lazy from uuid import uuid4 + +lazy import dill +lazy from blake3 import blake3 +lazy from loguru import logger + +BASE_FOLDER = Path( + f"./.cache/{version_info.major}.{version_info.minor}.{version_info.micro}/" +) + + +def _function_hash[**P, R](fn: Callable[P, R]) -> str: + fn = unwrap(fn) + + if not isinstance(fn, FunctionType): + raise ValueError(f"fn {fn} is not an instance of types.FunctionType") + + source = getsource(fn) + parsed = ast.parse(dedent(source)) + code = ast.dump(parsed) + sig = signature(fn) + name: str = fn.__qualname__ + + hasher = blake3() + hasher.update(name.encode("utf-8")) + hasher.update(sig.format().encode("utf-8")) + hasher.update(code.encode("utf-8")) + + return hasher.hexdigest() + + +def _parameter_hash[**P, R]( + fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs +) -> str: + fn = unwrap(fn) + sig = signature(fn) + + bound = sig.bind(*args, **kwargs) + bound.apply_defaults() + + serialized = dill.dumps(bound.arguments) + hasher = blake3(serialized) + + return hasher.hexdigest() + + +def persistent[**P, R](func: Callable[P, R]) -> Callable[P, R]: + fn_hash = _function_hash(func) + + folder = BASE_FOLDER / fn_hash + folder.mkdir(parents=True, exist_ok=True) + + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + param_hash = _parameter_hash(func, *args, **kwargs) + + path = folder / f"{param_hash}.dill" + + if not path.exists(): + logger.opt(colors=True).debug( + f"<white>[<cyan>CACHE</cyan>] [<red>MISS</red>] {func.__qualname__}({args}, {kwargs})</white>" + ) + result = func(*args, **kwargs) + payload = dill.dumps(result) + + temp_path = path.with_suffix(f".tmp.{uuid4().hex}") + temp_path.write_bytes(payload) + temp_path.replace(path) + else: + logger.opt(colors=True).debug( + f"<white>[<cyan>CACHE</cyan>] [<green>HIT</green>] {func.__qualname__}({args}, {kwargs})</white>" + ) + payload = path.read_bytes() + result = dill.loads(payload) + + return result + + return wrapper |
