1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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
|
