edq.util.json

  1"""
  2This file standardizes how we write and read JSON.
  3Specifically, we try to be flexible when reading (using JSON5),
  4and strict when writing (using vanilla JSON).
  5"""
  6
  7import enum
  8import gzip
  9import io
 10import json
 11import os
 12import typing
 13
 14import json5
 15
 16import edq.util.common
 17import edq.util.dirent
 18
 19def load(
 20        file_obj: typing.IO,
 21        strict: bool = False,
 22        gzipped: bool = False,
 23        encoding: str = edq.util.dirent.DEFAULT_ENCODING,
 24        **kwargs: typing.Any) -> typing.Any:
 25    """
 26    Load a file object/handler as JSON.
 27    If strict is set, then use standard Python JSON,
 28    otherwise use JSON5.
 29
 30    If `gzipped` is set, the file object is treated as a gzipped bytes stream (e.g. `open('test.json.gz', 'rb')`).
 31    """
 32
 33    if (gzipped):
 34        binary_file_obj = gzip.GzipFile(fileobj = file_obj)
 35        file_obj = io.TextIOWrapper(binary_file_obj, encoding = encoding)
 36
 37    if (strict):
 38        return json.load(file_obj, **kwargs)
 39
 40    return json5.load(file_obj, **kwargs)
 41
 42def loads(text: str, strict: bool = False, **kwargs: typing.Any) -> typing.Any:
 43    """
 44    Load a string as JSON.
 45    If strict is set, then use standard Python JSON,
 46    otherwise use JSON5.
 47    """
 48
 49    if (strict):
 50        return json.loads(text, **kwargs)
 51
 52    return json5.loads(text, **kwargs)
 53
 54def load_path(
 55        path: str,
 56        strict: bool = False,
 57        gzipped: typing.Union[bool, None] = None,
 58        encoding: str = edq.util.dirent.DEFAULT_ENCODING,
 59        **kwargs: typing.Any) -> typing.Any:
 60    """
 61    Load a file path as JSON.
 62    If strict is set, then use standard Python JSON,
 63    otherwise use JSON5.
 64
 65    If `gzipped` is not set, the behavior is guessed from the extension (".gz").
 66    """
 67
 68    if (not os.path.exists(path)):
 69        raise FileNotFoundError(f"File does not exist: '{path}'.")
 70
 71    if (os.path.isdir(path)):
 72        raise IsADirectoryError(f"Cannot open JSON file, expected a file but got a directory at '{path}'.")
 73
 74    if (gzipped is None):
 75        gzipped = (os.path.splitext(path)[-1] == '.gz')
 76
 77    open_func = open
 78    if (gzipped):
 79        open_func = gzip.open  # type: ignore[assignment]
 80
 81    with open_func(path, 'rt', encoding = encoding) as file:
 82        try:
 83            return load(file, strict = strict, **kwargs)
 84        except Exception as ex:
 85            raise ValueError(f"Failed to read JSON file '{path}'.") from ex
 86
 87def json_serialization_handle(value: typing.Any) -> typing.Union[typing.Dict[str, typing.Any], str, typing.Any]:
 88    """
 89    Handle objects that are not JSON serializable by default,
 90    e.g., calling vars() on an object.
 91    This is meant to be used as the `default` argument to `json` stdlib dumping functions.
 92    """
 93
 94    # If this looks like a edq.util.serial.DictSerializer.
 95    if (hasattr(value, 'to_dict')):
 96        return value.to_dict(edq.util.common.SerializationContext())
 97
 98    # If this looks like a edq.util.serial.PODSerializer.
 99    if (hasattr(value, 'to_pod')):
100        return value.to_pod(edq.util.common.SerializationContext())
101
102    if (isinstance(value, enum.Enum)):
103        return str(value)
104
105    if (hasattr(value, '__dict__')):
106        return dict(vars(value))
107
108    raise ValueError(f"Could not JSON serial object: '{value}'.")
109
110def dump(
111        data: typing.Any,
112        file_obj: typing.TextIO,
113        default: typing.Union[typing.Callable, None] = json_serialization_handle,
114        sort_keys: bool = True,
115        **kwargs: typing.Any) -> None:
116    """ Dump an object as a JSON file object. """
117
118    json.dump(data, file_obj, default = default, sort_keys = sort_keys, **kwargs)
119
120def dumps(
121        data: typing.Any,
122        default: typing.Union[typing.Callable, None] = json_serialization_handle,
123        sort_keys: bool = True,
124        **kwargs: typing.Any) -> str:
125    """ Dump an object as a JSON string. """
126
127    return json.dumps(data, default = default, sort_keys = sort_keys, **kwargs)
128
129def dump_path(
130        data: typing.Any,
131        path: str,
132        default: typing.Union[typing.Callable, None] = json_serialization_handle,
133        sort_keys: bool = True,
134        gzipped: typing.Union[bool, None] = None,
135        encoding: str = edq.util.dirent.DEFAULT_ENCODING,
136        **kwargs: typing.Any) -> None:
137    """
138    Dump an object as a JSON file.
139
140    If `gzipped` is not set, the behavior is guessed from the extension (".gz").
141    """
142
143    if (gzipped is None):
144        gzipped = (os.path.splitext(path)[-1] == '.gz')
145
146    open_func = open
147    if (gzipped):
148        open_func = gzip.open  # type: ignore[assignment]
149
150    with open_func(path, 'wt', encoding = encoding) as file:
151        dump(data, file, default = default, sort_keys = sort_keys, **kwargs)