# # This file is distributed under the MIT License. See LICENSE.md for details. # from __future__ import annotations import json from collections.abc import Buffer, Mapping from dataclasses import dataclass from io import StringIO from pathlib import Path from typing import Annotated, Dict, Generator, TextIO, Tuple, Type, final from .object import Kind, ObjectID, ObjectSet from .utils import bytes_to_string, is_mime_type_text, string_to_bytes from .utils.cabc import ABC, abstractmethod from .utils.registry import get_singleton ContainerID = Annotated[ str, """ This is used to identify a container in a savepoint. This is the name of the container. """, ] ConfigurationId = Annotated[ str, """ From a pipeline node, this represent the static and runtime configuration of this node, and all its dependencies. This is used to index the storage provider in SavePoints. """, ] Configuration = Annotated[ str, "The runtime-configuration of a pipe, most commonly it will be a json / yaml string that" "the pipe will parse and use to configure itself." "THIS IS DIFFERENT FROM ConfigurationId.", ] @dataclass(slots=True, frozen=True, order=True) class ContainerDeclaration: """ A ContainerDeclaration represents a container when describing a pipeline. It has a name and a type. Not to be confused with a Container, which is an instance of the container_type. """ name: str container_type: Type[Container] def instance(self) -> Container: return self.container_type() InvalidationList = list[Tuple[ConfigurationId, ContainerDeclaration, ObjectID]] def group_by_container( invalidation_list: InvalidationList, ) -> Generator[Tuple[ConfigurationId, ContainerDeclaration, ObjectSet], None, None]: if len(invalidation_list) == 0: return first_configuration_id, first_container, _ = invalidation_list[0] last = (first_configuration_id, first_container) def new_object_list(): return ObjectSet(first_container.container_type.kind, set()) object_list: ObjectSet = new_object_list() for configuration_id, container, obj in invalidation_list: if last != (configuration_id, container): # We changed container. Yield what we accumulated so far and prepare a new list yield (last[0], last[1], object_list) last = (configuration_id, container) object_list = new_object_list() # Record the current object object_list.add(obj) # Yield the last group yield (last[0], last[1], object_list) class Container(ABC): """ A Container contains objects of a certain kind. """ kind: Kind def __init__(self): """ This constructor just makes it explicit that a container should be able to be initialized without any arguments. """ @abstractmethod def objects(self) -> ObjectSet: pass @abstractmethod def deserialize(self, data: Mapping[ObjectID, Buffer]) -> None: """ Ingest data from a serialized format into this container. This is used to **add** cached objects to this container. """ @abstractmethod def serialize(self, objects: ObjectSet | None = None) -> Mapping[ObjectID, Buffer]: """ Dump objects from this container into a serialized format, dump all objects if none are specified. """ @classmethod @abstractmethod def mime_type(cls) -> str: """ The mime type of the serialized format of this container. This is used to inform the storage provider about the type of data it will be storing. """ @classmethod def is_text(cls) -> bool: """ Returns if the serialized format of this container is just text (e.g. JSON, YAML, etc.) or binary. This is used to improve transmission performance by avoiding unnecessary encoding/decoding steps. """ return is_mime_type_text(cls.mime_type()) @abstractmethod def verify(self) -> bool: pass def contains_all(self, obj: ObjectSet) -> bool: """ Check if an object set is fully contained in this container. This is a default implementation, the container can probably do it much more efficiently, so you are supposed to override this method. """ assert ( self.kind == obj.kind ), f"Container {self} has kind {self.kind}, but the object set has kind {obj.kind}." return self.objects().issubset(obj) @classmethod @final def from_dict(cls, data: dict[str, str]) -> Container: # Feed the data in the container container = cls() obj_id_type = get_singleton(ObjectID) # type: ignore[type-abstract] container.deserialize( { obj_id_type.deserialize(obj_id): string_to_bytes(content, container.is_text()) for obj_id, content in data.items() } ) return container @classmethod @final def from_string( cls, data: str, ) -> Container: """ Load a container into a serialized format. """ return cls.from_file(StringIO(data)) @classmethod @final def from_file( cls, path_or_file: str | Path | TextIO, ) -> Container: """ Load a container from a serialized format. This is used to **load** cached objects into this container. """ if isinstance(path_or_file, (str, Path)): path: str | Path = path_or_file with open(path, "r", encoding="utf-8") as f: return cls.from_file(f) # Read the file file: TextIO = path_or_file return cls.from_dict(json.load(file)) @final def to_dict(self) -> dict[str, str]: return { k.serialize(): bytes_to_string(v, self.is_text()) for k, v in self.serialize(self.objects()).items() } @final def to_string(self) -> str: """ Dump a container into a serialized format. """ data = StringIO() self.to_file(data) return data.getvalue() @final def to_file(self, path_or_file: str | Path | TextIO) -> None: """ Dump a container into a serialized format. This is used to **save** cached objects from this container. """ if isinstance(path_or_file, (str, Path)): with open(path_or_file, "w", encoding="utf-8") as f: self.to_file(f) return else: json.dump( self.to_dict(), path_or_file, indent=4, sort_keys=True, ) ContainerSet = Annotated[ Dict[ContainerDeclaration, Container], """A set of bindings between container declarations and container instances.""", ]