turbo_broccoli.context

  1"""
  2A context object holds information about the (de)serialization process, such as
  3the current position in the document, output paths, etc.
  4"""
  5
  6import tempfile
  7from os import environ as ENV
  8from pathlib import Path
  9from typing import Literal
 10from uuid import uuid4
 11
 12from .exceptions import TypeIsNodecode
 13
 14
 15def _list_of_types_to_dict(lot: list[type]) -> dict[str, type]:
 16    """
 17    Converts a list of types `[T1, T2, ...]` to a dict that looks like `{"T1":
 18    T1, "T2": T2, ...}`.
 19    """
 20    return {t.__name__: t for t in lot}
 21
 22
 23class Context:
 24    """
 25    (De)Serialization context, which is an object that contains various
 26    information and parameters about the ongoing operation. If you want your
 27    (de)serialization to behave a certain way, create a context object and pass
 28    it to
 29    [`turbo_broccoli.to_json`](https://altaris.github.io/turbo-broccoli/turbo_broccoli/turbo_broccoli.html#to_json)
 30    or
 31    [`turbo_broccoli.from_json`](https://altaris.github.io/turbo-broccoli/turbo_broccoli/turbo_broccoli.html#from_json).
 32    For convenience,
 33    [`turbo_broccoli.save_json`](https://altaris.github.io/turbo-broccoli/turbo_broccoli/turbo_broccoli.html#save_json)
 34    and
 35    [`turbo_broccoli.load_json`](https://altaris.github.io/turbo-broccoli/turbo_broccoli/turbo_broccoli.html#load_json)
 36    take the context parameter's as kwargs.
 37    """
 38
 39    artifact_path: Path
 40    dataclass_types: dict[str, type]
 41    file_path: Path | None
 42    json_path: str
 43    keras_format: str
 44    min_artifact_size: int = 8000
 45    nacl_shared_key: bytes | None
 46    nodecode_types: list[str]
 47    pandas_format: str
 48    pandas_kwargs: dict
 49    pytorch_module_types: dict[str, type]
 50    compress: bool
 51
 52    def __init__(
 53        self,
 54        file_path: str | Path | None = None,
 55        artifact_path: str | Path | None = None,
 56        min_artifact_size: int | None = None,
 57        nodecode_types: list[str] | None = None,
 58        keras_format: Literal["keras", "tf", "h5"] | None = None,
 59        pandas_format: (
 60            Literal[
 61                "csv",
 62                "excel",
 63                "feather",
 64                "html",
 65                "json",
 66                "latex",
 67                "orc",
 68                "parquet",
 69                "pickle",
 70                "sql",
 71                "stata",
 72                "xml",
 73            ]
 74            | None
 75        ) = None,
 76        pandas_kwargs: dict | None = None,
 77        nacl_shared_key: bytes | None = None,
 78        dataclass_types: dict[str, type] | list[type] | None = None,
 79        pytorch_module_types: dict[str, type] | list[type] | None = None,
 80        json_path: str = "$",
 81        compress: bool = False,
 82    ) -> None:
 83        """
 84        Args:
 85            file_path (str | Path | None, optional): Output JSON file path.
 86            artifact_path (str | Path | None, optional): Artifact path.
 87                Defaults to the parent directory of `file_path`, or a new
 88                temporary directory if `file_path` is `None`.
 89            min_artifact_size (int, optional): Byte strings (and everything
 90                that serialize to byte strings such as numpy arrays) larget
 91                than this will be stored in artifact rather than be embedded in
 92                the output JSON string/file.
 93            nodecode_types (list[str], optional): List of type names which
 94                shall be deserialized to `None` rather than their true value.
 95                See
 96                [`TB_NODECODE`](https://altaris.github.io/turbo-broccoli/turbo_broccoli.html#environment-variables)
 97            keras_format ("keras", "tf", "h5", optional): Format for Keras
 98                artifacts
 99            pandas_format ("csv", "excel", "feather", "html", "json", "latex",
100                "orc", "parquet", "pickle", "sql", "stata", "xml", optional):
101                Format for pandas artifacts
102            pandas_kwargs (dict, optional): kwargs to forward to the pandas
103                `to_*` and `read_*` function. For example, if
104                `pandas_format="parquet"`, then the content of `pandas.kwargs`
105                will be forwarded to
106                [`pandas.DataFrame.to_parquet`](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.to_parquet.html)
107                and
108                [`pandas.read_parquet`](https://pandas.pydata.org/docs/reference/api/pandas.read_parquet.html)
109            nacl_shared_key (bytes, optional): PyNaCl shared key. See also
110                [PyNaCl's
111                documentation](https://pynacl.readthedocs.io/en/latest/secret/#key)
112            dataclass_types (dict[str, type] | list[type], optional): List of
113                dataclass types for deserialization. See the
114                [README](https://altaris.github.io/turbo-broccoli/turbo_broccoli.html#supported-types).
115            pytorch_module_types (dict[str, type] | list[type], optional): List
116                of pytorch module types for deserialization. See the
117                [README](https://altaris.github.io/turbo-broccoli/turbo_broccoli.html#supported-types).
118            json_path (str, optional): Current JSONpath. Don't use.
119            compress (bool, optional): Wether to compress the output JSON file/
120                string. Defaults to `False`. If `file_path` is provided and
121                ends in `.json.gz`, then this parameter is overrode to `True`.
122        """
123        self.json_path = json_path
124        self.file_path = (
125            Path(file_path) if isinstance(file_path, str) else file_path
126        )
127        if artifact_path is None:
128            if p := ENV.get("TB_ARTIFACT_PATH"):
129                self.artifact_path = Path(p)
130            else:
131                self.artifact_path = (
132                    self.file_path.parent
133                    if self.file_path is not None
134                    else Path(tempfile.mkdtemp())
135                )
136        else:
137            self.artifact_path = Path(artifact_path)
138        self.min_artifact_size = (
139            min_artifact_size
140            if min_artifact_size is not None
141            else int(ENV.get("TB_MAX_NBYTES", 8000))
142        )
143        self.nodecode_types = nodecode_types or ENV.get(
144            "TB_NODECODE", ""
145        ).split(",")
146        self.keras_format = keras_format or str(
147            ENV.get("TB_KERAS_FORMAT", "tf")
148        )
149        self.pandas_format = pandas_format or str(
150            ENV.get("TB_PANDAS_FORMAT", "csv")
151        )
152        self.pandas_kwargs = pandas_kwargs or {}
153        if isinstance(nacl_shared_key, bytes):
154            self.nacl_shared_key = nacl_shared_key
155        elif "TB_SHARED_KEY" in ENV:
156            self.nacl_shared_key = str(ENV["TB_SHARED_KEY"]).encode("utf-8")
157        else:
158            self.nacl_shared_key = None
159        self.dataclass_types = (
160            _list_of_types_to_dict(dataclass_types)
161            if isinstance(dataclass_types, list)
162            else (dataclass_types or {})
163        )
164        self.pytorch_module_types = (
165            _list_of_types_to_dict(pytorch_module_types)
166            if isinstance(pytorch_module_types, list)
167            else (pytorch_module_types or {})
168        )
169        self.compress = (
170            True
171            if (
172                self.file_path is not None
173                and self.file_path.name.endswith(".json.gz")
174            )
175            else compress
176        )
177
178    def __repr__(self) -> str:
179        fp, ap = str(self.file_path), str(self.artifact_path)
180        return (
181            f"Context(file_path={fp}, artifact_path={ap}, "
182            f"json_path={self.json_path})"
183        )
184
185    def __truediv__(self, x: str | int) -> "Context":
186        """
187        Returns a copy of the current context but where the `json_path`
188        attribute is `self.json_path + "." + str(x)`. Use this when you're
189        going down the document.
190        """
191        kwargs = self.__dict__.copy()
192        kwargs["json_path"] = self.json_path + "." + str(x)
193        return Context(**kwargs)
194
195    def id_to_artifact_path(self, art_id: str, extension: str = "tb") -> Path:
196        """
197        Takes an artifact id (which is an UUID4 string) and returns the
198        absolute path to the corresponding artifact file.
199        """
200        art_fn = art_id + "." + extension
201        if self.file_path is not None:
202            art_fn = self.file_path.stem + "." + art_fn
203        return self.artifact_path / art_fn
204
205    def new_artifact_path(self, extension: str = "tb") -> tuple[Path, str]:
206        """Returns the path to a new artifact alongside the artifact's ID"""
207        art_id = str(uuid4())
208        return self.id_to_artifact_path(art_id, extension), art_id
209
210    def raise_if_nodecode(self, type_name: str) -> None:
211        """
212        Raises a `turbo_broccoli.exceptions.TypeIsNodecode` exception if
213        `type_name` or any prefix is set to not be decoded in this context (see
214        `nodecode_types` constructor argument).
215
216        For example, if `type_name` is `a.b.c`, then this method raises
217        `turbo_broccoli.exceptions.TypeIsNodecode` if either `a`, `a.b`, or
218        `a.b.c` is set as a nodecode type.
219        """
220        parts = type_name.split(".")
221        for i in range(1, len(parts) + 1):
222            t = ".".join(parts[:i])
223            if t in self.nodecode_types:
224                raise TypeIsNodecode(t)