turbo_broccoli.turbo_broccoli

  1"""Main module containing the JSON encoder and decoder methods."""
  2
  3import json
  4import zlib
  5from pathlib import Path
  6from typing import Any
  7
  8from turbo_broccoli import user
  9from turbo_broccoli.context import Context
 10from turbo_broccoli.custom import get_decoders, get_encoders
 11from turbo_broccoli.exceptions import TypeIsNodecode, TypeNotSupported
 12
 13
 14def _from_jsonable(obj: Any, ctx: Context) -> Any:
 15    """
 16    Takes an object fresh from `json.load` or `json.loads` and loads types that
 17    are supported by TurboBroccoli therein.
 18    """
 19    if isinstance(obj, dict):
 20        obj = {k: _from_jsonable(v, ctx / k) for k, v in obj.items()}
 21        if "__type__" in obj:
 22            try:
 23                ctx.raise_if_nodecode(obj["__type__"])
 24                base = obj["__type__"].split(".")[0]
 25                if base == "user":
 26                    name = ".".join(obj["__type__"].split(".")[1:])
 27                    if decoder := user.decoders.get(name):
 28                        obj = decoder(obj, ctx)
 29                else:
 30                    obj = get_decoders()[base](obj, ctx)
 31            except TypeIsNodecode:
 32                pass
 33    elif isinstance(obj, list):
 34        return [_from_jsonable(v, ctx / str(i)) for i, v in enumerate(obj)]
 35    elif isinstance(obj, tuple):
 36        return tuple(
 37            _from_jsonable(v, ctx / str(i)) for i, v in enumerate(obj)
 38        )
 39    return obj
 40
 41
 42def _make_or_set_ctx(
 43    file_path: str | Path | None, ctx: Context | None, **kwargs
 44) -> Context:
 45    """
 46    Generate a context object that is consistent with the inputs.
 47    """
 48    if file_path is None and ctx is None:
 49        raise ValueError(
 50            "Either a file path or a context (or both) must be provided."
 51        )
 52    if file_path is not None and ctx is not None:
 53        if ctx.file_path is not None and ctx.file_path != file_path:
 54            raise ValueError(
 55                "The file path in the context does not match the provided "
 56                "file path."
 57            )
 58        assert isinstance(file_path, (str, Path))  # for typechecking
 59        ctx.file_path = Path(file_path)
 60    if ctx is None:
 61        ctx = Context(file_path=file_path, **kwargs)
 62    return ctx
 63
 64
 65def _to_jsonable(obj: Any, ctx: Context) -> Any:
 66    """
 67    Transforms an object (dict, list, primitive) that possibly contains types
 68    that TurboBroccoli's custom encoders support, and returns an object that is
 69    readily vanilla JSON-serializable.
 70    """
 71    name = obj.__class__.__name__
 72    if name in user.encoders:
 73        obj = user.encoders[name](obj, ctx)
 74    for encoder in get_encoders():
 75        try:
 76            obj = encoder(obj, ctx)
 77            break
 78        except TypeNotSupported:
 79            pass
 80    if isinstance(obj, dict):
 81        return {k: _to_jsonable(v, ctx / k) for k, v in obj.items()}
 82    if isinstance(obj, list):
 83        return [_to_jsonable(v, ctx / str(i)) for i, v in enumerate(obj)]
 84    if isinstance(obj, tuple):
 85        return tuple(_to_jsonable(v, ctx / str(i)) for i, v in enumerate(obj))
 86    return obj
 87
 88
 89def from_json(doc: str, ctx: Context | None = None) -> Any:
 90    """
 91    Deserializes a JSON string. The context's file path and compression setting
 92    will be ignored.
 93    """
 94    return _from_jsonable(json.loads(doc), Context() if ctx is None else ctx)
 95
 96
 97def load_json(
 98    file_path: str | Path | None = None, ctx: Context | None = None, **kwargs
 99) -> Any:
100    """
101    Loads a JSON file.
102
103    Args:
104        file_path (str | Path | None): If left to `None`, a context with a file
105            path must be provided
106        ctx (Context | None): The context to use. If `None`, a new context will
107            be created with the kwargs.
108        **kwargs: Forwarded to the `turbo_broccoli.context.Context`
109            constructor. If `ctx` is provided, the kwargs are ignored.
110    """
111    ctx = _make_or_set_ctx(file_path, ctx, **kwargs)
112    assert isinstance(ctx.file_path, Path)  # for typechecking
113    if ctx.compress:
114        with ctx.file_path.open(mode="rb") as fp:
115            s = zlib.decompress(fp.read()).decode("utf-8")
116        return _from_jsonable(json.loads(s), ctx)
117    with ctx.file_path.open(mode="r", encoding="utf-8") as fp:
118        return _from_jsonable(json.load(fp), ctx)
119
120
121def save_json(
122    obj: Any,
123    file_path: str | Path | None = None,
124    ctx: Context | None = None,
125    **kwargs,
126) -> None:
127    """
128    Serializes an object and writes the result to a file. The artifact path and
129    the output file's parent folder will be created if they don't exist.
130
131    Args:
132        obj (Any):
133        file_path (str | Path):
134        ctx (Context | None): The context to use. If `None`, a new context will
135            be created with the kwargs.
136        **kwargs: Forwarded to the `turbo_broccoli.context.Context`
137            constructor.
138    """
139    ctx = _make_or_set_ctx(file_path, ctx, **kwargs)
140    data = json.dumps(_to_jsonable(obj, ctx))
141    assert isinstance(ctx.file_path, Path)  # for typechecking
142    if not ctx.file_path.parent.exists():
143        ctx.file_path.parent.mkdir(parents=True)
144    if ctx.compress:
145        with ctx.file_path.open(mode="wb") as fp:
146            fp.write(zlib.compress(data.encode("utf-8")))
147    else:
148        with ctx.file_path.open(mode="w", encoding="utf-8") as fp:
149            fp.write(data)
150
151
152def to_json(obj: Any, ctx: Context | None = None) -> str:
153    """
154    Converts an object to a JSON string. The context's artifact folder will be
155    created if it doesn't exist. The context's file path and compression
156    setting will be ignored.
157    """
158    ctx = Context() if ctx is None else ctx
159    if not ctx.artifact_path.exists():
160        ctx.artifact_path.mkdir(parents=True)
161    return json.dumps(_to_jsonable(obj, ctx))