Skip to content

Instantly share code, notes, and snippets.

@sonhanguyen

sonhanguyen/cache.py

Last active Oct 17, 2020
Embed
What would you like to do?
A filesystem-based memoization decorator
import os
import collections
import pickle
import json
import urllib.parse
import re
from functools import wraps
from typing import Callable, Any
from dataclasses import dataclass, fields
@dataclass
class Key:
function: Callable
args: tuple
kwarg: dict
def __str__(self):
stringify = lambda key: re.sub(
r'[{}"[\]]',
'',
json.dumps(
key,
default=lambda unserializable: '',
separators=('_', '-')
)
)
args = stringify(self.args)
kwarg = stringify(self.kwarg)
params = '_'.join(filter(lambda s: s, [args, kwarg]))
func = self.function.__name__
return os.path.join(
func,
urllib.parse.quote(params)) if params\
else func
@dataclass
class CacheOptions:
read: Callable[[Key], Any]
write: Callable[[Key], None]
invalidate: Callable[[Any, dict, tuple], bool]
key: Callable[[Callable, tuple, dict], Key]
def __post_init__(self):
self.invalidate = self.invalidate or (lambda *_: False)
self.key = self.key or Key
def cache(options: CacheOptions):
def cached(func):
@wraps(func)
def with_cache(*args, **kwarg):
key = options.key(func, args, kwarg)
try:
result = options.read(key)
except KeyError:
result = func(*args, **kwarg)
if not options.invalidate(result, kwarg, args): options.write(key, result)
return result
return with_cache
return cached
def file_cache(**options):
def dumps(fname, serialized):
os.makedirs(os.path.dirname(fname), exist_ok=True)
with open(fname, 'wb') as file:
file.write(serialized)
def loads(fname):
with open(fname, 'rb') as file:
return file.read()
@dataclass
class Options(CacheOptions):
dir: str
marshal: Callable
unmarshal: Callable
def __init__(self, **kwarg):
attrs = self.__dict__
attrs.update(kwarg)
path = lambda key: os.path.join(attrs.get('dir', '.'), str(key))
unmarshal = attrs.get('unmarshal',
lambda serialized, key: pickle.loads(serialized))
marshal = attrs.get('marshal', lambda val, key: pickle.dumps(val))
def get(key):
try:
return unmarshal(loads(path(key)), key)
except (FileNotFoundError, NotADirectoryError):
raise KeyError('Cache misses for ' + str(key))
props = map(lambda field: field.name, fields(super()))
super().__init__(**{
**{ key: kwarg.get(key) for key in props },
**{
'write': lambda key, val: dumps(path(key), marshal(val, key)),
'read': get
}
})
return cache(Options(**options))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment