Commit | Line | Data |
---|---|---|
f3760aea SM |
1 | # SPDX-License-Identifier: MIT |
2 | # | |
3 | # Copyright (C) 2023 EfficiOS, inc. | |
4 | # | |
5 | # pyright: strict, reportTypeCommentUsage=false | |
6 | ||
7 | ||
8 | import re | |
9 | import json | |
10 | import typing | |
11 | from typing import ( | |
12 | Any, | |
13 | Dict, | |
14 | List, | |
15 | Type, | |
16 | Union, | |
17 | TextIO, | |
18 | Generic, | |
19 | Mapping, | |
20 | TypeVar, | |
21 | Optional, | |
22 | Sequence, | |
23 | overload, | |
24 | ) | |
25 | ||
26 | # Internal type aliases and variables | |
27 | _RawArrayT = List["_RawValT"] | |
28 | _RawObjT = Dict[str, "_RawValT"] | |
29 | _RawValT = Union[None, bool, int, float, str, _RawArrayT, _RawObjT] | |
30 | _RawValTV = TypeVar("_RawValTV", bound="_RawValT") | |
31 | _ValTV = TypeVar("_ValTV", bound="Val") | |
32 | ||
33 | ||
34 | # Type of a single JSON value path element. | |
35 | PathElemT = Union[str, int] | |
36 | ||
37 | ||
38 | # A JSON value path. | |
39 | class Path: | |
40 | def __init__(self, elems: Optional[List[PathElemT]] = None): | |
41 | if elems is None: | |
42 | elems = [] | |
43 | ||
44 | self._elems = elems | |
45 | ||
46 | # Elements of this path. | |
47 | @property | |
48 | def elems(self): | |
49 | return self._elems | |
50 | ||
51 | # Returns a new path containing the current elements plus `elem`. | |
52 | def __truediv__(self, elem: PathElemT): | |
53 | return Path(self._elems + [elem]) | |
54 | ||
55 | # Returns a valid jq filter. | |
56 | def __str__(self): | |
57 | s = "" | |
58 | ||
59 | for elem in self._elems: | |
60 | if type(elem) is str: | |
61 | if re.match(r"[a-zA-Z]\w*$", elem): | |
62 | s += ".{}".format(elem) | |
63 | else: | |
64 | s += '."{}"'.format(elem) | |
65 | else: | |
66 | assert type(elem) is int | |
67 | s += "[{}]".format(elem) | |
68 | ||
69 | if not s.startswith("."): | |
70 | s = "." + s | |
71 | ||
72 | return s | |
73 | ||
74 | ||
75 | # Base of any JSON value. | |
76 | class Val: | |
77 | _name = "a value" | |
78 | ||
79 | def __init__(self, path: Optional[Path] = None): | |
80 | if path is None: | |
81 | path = Path() | |
82 | ||
83 | self._path = path | |
84 | ||
85 | # Path to this JSON value. | |
86 | @property | |
87 | def path(self): | |
88 | return self._path | |
89 | ||
90 | ||
91 | # JSON null value. | |
92 | class NullVal(Val): | |
93 | _name = "a null" | |
94 | ||
95 | ||
96 | # JSON scalar value. | |
97 | class _ScalarVal(Val, Generic[_RawValTV]): | |
98 | def __init__(self, raw_val: _RawValTV, path: Optional[Path] = None): | |
99 | super().__init__(path) | |
100 | self._raw_val = raw_val | |
101 | ||
102 | # Raw value. | |
103 | @property | |
104 | def val(self): | |
105 | return self._raw_val | |
106 | ||
107 | ||
108 | # JSON boolean value. | |
109 | class BoolVal(_ScalarVal[bool]): | |
110 | _name = "a boolean" | |
111 | ||
112 | def __bool__(self): | |
113 | return self.val | |
114 | ||
115 | ||
116 | # JSON integer value. | |
117 | class IntVal(_ScalarVal[int]): | |
118 | _name = "an integer" | |
119 | ||
120 | def __int__(self): | |
121 | return self.val | |
122 | ||
123 | ||
124 | # JSON floating point number value. | |
125 | class FloatVal(_ScalarVal[float]): | |
126 | _name = "a floating point number" | |
127 | ||
128 | def __float__(self): | |
129 | return self.val | |
130 | ||
131 | ||
132 | # JSON string value. | |
133 | class StrVal(_ScalarVal[str]): | |
134 | _name = "a string" | |
135 | ||
136 | def __str__(self): | |
137 | return self.val | |
138 | ||
139 | ||
140 | # JSON array value. | |
141 | class ArrayVal(Val, Sequence[Val]): | |
142 | _name = "an array" | |
143 | ||
144 | def __init__(self, raw_val: _RawArrayT, path: Optional[Path] = None): | |
145 | super().__init__(path) | |
146 | self._raw_val = raw_val | |
147 | ||
148 | # Returns the value at index `index`. | |
149 | # | |
150 | # Raises `TypeError` if the type of the returned value isn't | |
151 | # `expected_elem_type`. | |
152 | def at(self, index: int, expected_elem_type: Type[_ValTV]): | |
153 | try: | |
154 | elem = self._raw_val[index] | |
155 | except IndexError: | |
156 | raise IndexError( | |
157 | "`{}`: array index {} out of range".format(self._path, index) | |
158 | ) | |
159 | ||
160 | return wrap(elem, self._path / index, expected_elem_type) | |
161 | ||
162 | # Returns an iterator yielding the values of this array value. | |
163 | # | |
164 | # Raises `TypeError` if the type of any yielded value isn't | |
165 | # `expected_elem_type`. | |
166 | def iter(self, expected_elem_type: Type[_ValTV]): | |
167 | for i in range(len(self._raw_val)): | |
168 | yield self.at(i, expected_elem_type) | |
169 | ||
170 | @overload | |
171 | def __getitem__(self, index: int) -> Val: | |
172 | ... | |
173 | ||
174 | @overload | |
175 | def __getitem__(self, index: slice) -> Sequence[Val]: | |
176 | ... | |
177 | ||
178 | def __getitem__(self, index: Union[int, slice]) -> Union[Val, Sequence[Val]]: | |
179 | if type(index) is slice: | |
180 | raise NotImplementedError | |
181 | ||
182 | return self.at(index, Val) | |
183 | ||
184 | def __len__(self): | |
185 | return len(self._raw_val) | |
186 | ||
187 | ||
188 | # JSON object value. | |
189 | class ObjVal(Val, Mapping[str, Val]): | |
190 | _name = "an object" | |
191 | ||
192 | def __init__(self, raw_val: _RawObjT, path: Optional[Path] = None): | |
193 | super().__init__(path) | |
194 | self._raw_val = raw_val | |
195 | ||
196 | # Returns the value having the key `key`. | |
197 | # | |
198 | # Raises `TypeError` if the type of the returned value isn't | |
199 | # `expected_type`. | |
200 | def at(self, key: str, expected_type: Type[_ValTV]): | |
201 | try: | |
202 | val = self._raw_val[key] | |
203 | except KeyError: | |
204 | raise KeyError("`{}`: no value has the key `{}`".format(self._path, key)) | |
205 | ||
206 | return wrap(val, self._path / key, expected_type) | |
207 | ||
208 | def __getitem__(self, key: str) -> Val: | |
209 | return self.at(key, Val) | |
210 | ||
211 | def __len__(self): | |
212 | return len(self._raw_val) | |
213 | ||
214 | def __iter__(self): | |
215 | return iter(self._raw_val) | |
216 | ||
217 | ||
218 | # Raises `TypeError` if the type of `val` is not `expected_type`. | |
219 | def _check_type(val: Val, expected_type: Type[Val]): | |
220 | if not isinstance(val, expected_type): | |
221 | raise TypeError( | |
e45fa35b SM |
222 | "`{}`: expecting {} value, got {}".format( |
223 | val.path, | |
224 | expected_type._name, # pyright: ignore [reportPrivateUsage] | |
225 | type(val)._name, # pyright: ignore [reportPrivateUsage] | |
f3760aea SM |
226 | ) |
227 | ) | |
228 | ||
229 | ||
230 | # Wraps the raw value `raw_val` into an equivalent instance of some | |
231 | # `Val` subclass having the path `path` and returns it. | |
232 | # | |
233 | # If the resulting JSON value type isn't `expected_type`, then this | |
234 | # function raises `TypeError`. | |
235 | def wrap( | |
236 | raw_val: _RawValT, path: Optional[Path] = None, expected_type: Type[_ValTV] = Val | |
237 | ) -> _ValTV: | |
238 | val = None | |
239 | ||
240 | if raw_val is None: | |
241 | val = NullVal(path) | |
242 | elif isinstance(raw_val, bool): | |
243 | val = BoolVal(raw_val, path) | |
244 | elif isinstance(raw_val, int): | |
245 | val = IntVal(raw_val, path) | |
246 | elif isinstance(raw_val, float): | |
247 | val = FloatVal(raw_val, path) | |
248 | elif isinstance(raw_val, str): | |
249 | val = StrVal(raw_val, path) | |
250 | elif isinstance(raw_val, list): | |
251 | val = ArrayVal(raw_val, path) | |
252 | else: | |
253 | assert isinstance(raw_val, dict) | |
254 | val = ObjVal(raw_val, path) | |
255 | ||
256 | assert val is not None | |
257 | _check_type(val, expected_type) | |
258 | return typing.cast(_ValTV, val) | |
259 | ||
260 | ||
261 | # Like json.loads(), but returns a `Val` instance, raising `TypeError` | |
262 | # if its type isn't `expected_type`. | |
263 | def loads(s: str, expected_type: Type[_ValTV] = Val, **kwargs: Any) -> _ValTV: | |
264 | return wrap(json.loads(s, **kwargs), Path(), expected_type) | |
265 | ||
266 | ||
267 | # Like json.load(), but returns a `Val` instance, raising `TypeError` if | |
268 | # its type isn't `expected_type`. | |
269 | def load(fp: TextIO, expected_type: Type[_ValTV] = Val, **kwargs: Any) -> _ValTV: | |
270 | return wrap(json.load(fp, **kwargs), Path(), expected_type) |