viat.providers.storage ¶

class JsonAttributeReader(ViatAttributeReader, Mapping[str, JsonT]):
    """A storage reader for JSON objects.

    Args:
        path: The path to which the table corresponds.
        json_object: The object to proxy.
    """

    path: pathlib.Path
    """The path used to initialize the reader."""

    json_object: JsonObjectT
    """The object used to initialize the reader."""

    def __init__(self, path: pathlib.Path, json_object: JsonObjectT) -> None:
        self.path = path
        self.json_object = json_object

    def __getitem__(self, key: str) -> JsonT:
        try:
            return self.json_object[key]
        except KeyError:
            raise ViatMissingAttributeError(self.path, key) from None

    def __iter__(self) -> Iterator[str]:
        return iter(self.json_object)

    def __len__(self) -> int:
        return len(self.json_object)

class JsonAttributeStorage(JsonAttributeStorageMixin, ViatAttributeStorage):
    """The JSON file storage class.

    Args:
        config: All configuration required for the storage.
        resolver: A path resolver used by connections.
    """

    config: JsonAttributeStorageConfig
    """The configuration used to initialize the storage."""

    resolver: ViatPathResolver | None
    """The resolver used to initialize the storage."""

    def __init__(
        self,
        config: JsonAttributeStorageConfig,
        resolver: ViatPathResolver | None = None,
    ) -> None:
        self.config = config
        self.resolver = resolver

    @override
    def _get_json_schema_path(self) -> pathlib.Path | None:
        return self.config.json_schema_path

    @override
    def _load_storage_data(self) -> MutableMapping[str, JsonT]:
        try:
            with self.config.storage_path.open() as file:
                return json.load(file)
        except FileNotFoundError:
            return {}
        except OSError as err:
            raise ViatAttributeStorageError('Could not read the storage file') from err
        except json.JSONDecodeError as err:
            raise ViatMalformedStoredDataError(self.config.storage_path) from err

    @override
    def _dump_storage_data(self, data: JsonObjectT) -> None:
        try:
            json_string = json.dumps(data, indent=self.config.indent)
        except json.JSONDecodeError as err:
            raise ViatAttributeStorageError('Could not serialize the stored attribute contents') from err
        finally:
            self._active_conn = None

        self.config.storage_path.write_text(json_string, encoding='utf-8')

@dataclass
class JsonAttributeStorageConfig:
    """Configuration for the JSON file storage."""

    storage_path: pathlib.Path
    """The path to a file for storing attributes.

    The configuration loader sets this to the default value `storage.json`.
    """

    json_schema_path: pathlib.Path | None = None
    """Path to a JSON schema file.

    If empty, we assume no schema.

    The configuration loader (but not the storage itself) uses the default value `schema.json` if it exists.
    """

    indent: int | str | None = None
    """An indent parameter for Python's JSON writer."""

@dataclass
class JsonAttributeStorageConnection(ViatAttributeStorageConnection):
    """A connection class for the JSON and TOML storages.

    Args:
        payload: The initial state of the storage.
        resolver: A path resolver used when creating readers and mutators.
        validator: A validator to be used on the contents of the storage.
    """

    payload: MutableMapping[str, JsonT]
    """The payload used to initialize the connection."""

    validator: Callable[[JsonT], None] | None
    """The validator used to initialize the connection."""

    resolver: ViatPathResolver | None
    """The resolver used to initialize the connection."""

    has_mutations: bool
    """An indicator whether the connection payload has been mutated."""

    _locked: MutableSet[pathlib.Path]

    def __init__(
        self,
        payload: MutableMapping[str, JsonT],
        resolver: ViatPathResolver | None = None,
        validator: Callable[[JsonT], None] | None = None,
    ) -> None:
        self.payload = payload
        self.resolver = resolver
        self.validator = validator
        self.has_mutations = False
        self._locked = set[pathlib.Path]()

        if validator:
            for key, value in self.payload.items():
                if not isinstance(value, JsonObject):
                    raise ViatMalformedStoredDataError(pathlib.Path(key))

                if self.validator:
                    try:
                        self.validator(value)
                    except fastjsonschema.JsonSchemaValueException as err:
                        emit_warning(
                            ViatStoredDataValidationWarning(pathlib.Path(key), err),
                            stacklevel=2,
                        )

    def _resolve_path(self, path: pathlib.Path | str) -> pathlib.Path:
        return self.resolver.relativize(path) if self.resolver else pathlib.Path(path)

    @contextlib.contextmanager
    def get_reader(self, path: pathlib.Path | str) -> Generator[JsonAttributeReader]:
        rel_path = self._resolve_path(path)

        if rel_path in self._locked:
            raise ViatAttributeStorageError(f'There is already an active reader or mutator for {rel_path}')

        stored_data = self.payload.get(rel_path.as_posix())

        if stored_data is not None and not isinstance(stored_data, JsonObject):
            raise ViatMalformedStoredDataError(rel_path)

        payload: JsonObjectT = stored_data or {}
        self._locked.add(rel_path)

        try:
            yield JsonAttributeReader(rel_path, payload)
        finally:
            self._locked.remove(rel_path)

    @contextlib.contextmanager
    def get_mutator(self, path: pathlib.Path | str) -> Generator[JsonAttributeMutator]:
        rel_path = self._resolve_path(path)

        if rel_path in self._locked:
            raise ViatAttributeStorageError(f'There is already an active reader or mutator for {rel_path}')

        stored_data = self.payload.get(rel_path.as_posix())

        if stored_data is not None and not isinstance(stored_data, MutableMapping):
            raise ViatMalformedStoredDataError(rel_path)

        payload: MutableMapping[str, JsonT] = stored_data or {}
        self._locked.add(rel_path)
        self.has_mutations = True

        try:
            yield JsonAttributeMutator(rel_path, payload)
        finally:
            self._locked.remove(rel_path)

        if len(payload) == 0:
            if stored_data is not None:
                del self.payload[rel_path.as_posix()]

            return

        if self.validator:
            try:
                self.validator(payload)
            except fastjsonschema.JsonSchemaValueException as err:
                raise ViatValidationError(rel_path) from err

        self.payload[rel_path.as_posix()] = payload

    @override
    def iter_known_paths(self) -> Iterable[pathlib.Path]:
        for key in self.payload:
            yield pathlib.Path(key)

class TomlAttributeStorage(JsonAttributeStorageMixin, ViatAttributeStorage):
    """The TOML file storage class.

    Args:
        config: All configuration required for the storage.
        resolver: A path resolver used by connections.
    """

    config: TomlAttributeStorageConfig
    """The configuration used to initialize the storage."""

    resolver: ViatPathResolver | None
    """The resolver used to initialize the storage."""

    def __init__(
        self,
        config: TomlAttributeStorageConfig,
        resolver: ViatPathResolver | None = None,
    ) -> None:
        self.config = config
        self.resolver = resolver

    @override
    def _get_json_schema_path(self) -> pathlib.Path | None:
        return self.config.json_schema_path

    @override
    def _load_storage_data(self) -> MutableMapping[str, JsonT]:
        try:
            with self.config.storage_path.open('rb') as file:
                return tomllib.load(file)
        except FileNotFoundError:
            return {}
        except OSError as err:
            raise ViatAttributeStorageError('Could not read the storage file') from err
        except tomllib.TOMLDecodeError as err:
            raise ViatMalformedStoredDataError(self.config.storage_path) from err

    @override
    def _dump_storage_data(self, data: JsonObjectT) -> None:
        try:
            toml_string = tomli_w.dumps(data)
        except (ValueError, TypeError) as err:
            raise ViatAttributeStorageError('Could not serialize the stored attribute contents') from err
        finally:
            self._active_conn = None

        self.config.storage_path.write_text(toml_string, encoding='utf-8')

@dataclass
class TomlAttributeStorageConfig:
    """Configuration for the TOML file storage."""

    storage_path: pathlib.Path
    """The path to a file for storing attributes.

    The configuration loader sets this to the default value `storage.toml`.
    """

    json_schema_path: pathlib.Path | None = None
    """Path to a JSON schema file.

    We utilize the fact that TOML maps to JSON to enable validation via JSON schemas.

    If empty, we assume no schema.

    The configuration loader (but not the storage itself) uses the default value `schema.json` if it exists.
    """