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)