from io import BytesIO import os class DiskCache: def __init__(self, path=None, enabled=True): if path is None: self.path = os.path.join(os.path.dirname(__file__), 'cache') else: self.path = path self.enabled = enabled if self.enabled is False: return os.makedirs(self.path, exist_ok=True) def set(self, identifier, buf): if self.enabled is False: return with open(self._build_path(identifier), 'wb') as f: f.write(buf.read()) def get(self, identifier): if self.enabled is False: raise FileNotFoundError with open(self._build_path(identifier), 'rb') as f: return BytesIO(f.read()) def _build_path(self, identifier): return os.path.join(self.path, identifier)