Skip to content

Commit 06c1f05

Browse files
fix(spanner): implement dict protocol and nested unwrapping for JsonObject (#17915)
Fixes JsonObject array/scalar/null variants breaking standard Python container protocols (len, bool, iter, getitem, contains, eq). Adds _unwrap_for_json helper to safely serialize nested JsonObject instances without data erasure. Fixes #15870
1 parent ec949e9 commit 06c1f05

3 files changed

Lines changed: 612 additions & 96 deletions

File tree

packages/google-cloud-spanner/google/cloud/spanner_v1/data_types.py

Lines changed: 202 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -31,31 +31,192 @@ class JsonObject(dict):
3131
"""
3232

3333
def __init__(self, *args, **kwargs):
34-
self._is_null = (args, kwargs) == ((), {}) or args == (None,)
35-
self._is_array = len(args) and isinstance(args[0], (list, tuple))
36-
self._is_scalar_value = len(args) == 1 and not isinstance(args[0], (list, dict))
34+
self._is_null = (
35+
(args, kwargs) == ((), {})
36+
or (len(args) == 1 and args[0] is None)
37+
or (len(args) == 1 and isinstance(args[0], JsonObject) and args[0]._is_null)
38+
)
39+
self._is_array = (
40+
len(args) == 1
41+
and not self._is_null
42+
and isinstance(
43+
args[0]._array_value
44+
if isinstance(args[0], JsonObject) and args[0]._is_array
45+
else args[0],
46+
(list, tuple),
47+
)
48+
)
49+
self._is_scalar_value = (
50+
len(args) == 1
51+
and not self._is_null
52+
and not self._is_array
53+
and not isinstance(
54+
args[0]._simple_value
55+
if isinstance(args[0], JsonObject) and args[0]._is_scalar_value
56+
else args[0],
57+
dict,
58+
)
59+
)
60+
61+
if self._is_null:
62+
super().__init__({"__json_null__": True})
63+
return
64+
65+
if len(args) and isinstance(args[0], JsonObject):
66+
if args[0]._is_array:
67+
self._array_value = list(args[0]._array_value)
68+
return
69+
elif args[0]._is_scalar_value:
70+
self._simple_value = args[0]._simple_value
71+
return
72+
else:
73+
super().__init__(args[0].copy())
74+
return
3775

3876
# if the JSON object is represented with an array,
3977
# the value is contained separately
4078
if self._is_array:
41-
self._array_value = args[0]
79+
self._array_value = list(args[0])
4280
return
4381

4482
# If it's a scalar value, set _simple_value and return early
4583
if self._is_scalar_value:
4684
self._simple_value = args[0]
4785
return
4886

49-
if len(args) and isinstance(args[0], JsonObject):
50-
self._is_array = args[0]._is_array
51-
self._is_scalar_value = args[0]._is_scalar_value
87+
super().__init__(*args, **kwargs)
88+
89+
def get(self, key, default=None):
90+
if self._is_array:
91+
try:
92+
return self._array_value[key]
93+
except (IndexError, TypeError):
94+
return default
95+
if self._is_scalar_value or self._is_null:
96+
return default
97+
return super().get(key, default)
98+
99+
def copy(self):
100+
if self._is_array:
101+
return JsonObject(list(self._array_value))
102+
if self._is_scalar_value:
103+
return JsonObject(self._simple_value)
104+
if self._is_null:
105+
return JsonObject()
106+
return JsonObject(super().copy())
107+
108+
def keys(self):
109+
if self._is_array:
110+
return dict(enumerate(self._array_value)).keys()
111+
if self._is_scalar_value or self._is_null:
112+
return {}.keys()
113+
return super().keys()
114+
115+
def values(self):
116+
if self._is_array:
117+
return dict(enumerate(self._array_value)).values()
118+
if self._is_scalar_value:
119+
return {0: self._simple_value}.values()
120+
if self._is_null:
121+
return {}.values()
122+
return super().values()
123+
124+
def items(self):
125+
if self._is_array:
126+
return dict(enumerate(self._array_value)).items()
127+
if self._is_scalar_value:
128+
return {0: self._simple_value}.items()
129+
if self._is_null:
130+
return {}.items()
131+
return super().items()
132+
133+
def pop(self, key, *args):
134+
if self._is_array:
135+
try:
136+
return self._array_value.pop(key)
137+
except (IndexError, TypeError):
138+
if args:
139+
return args[0]
140+
raise KeyError(key) from None
141+
if self._is_scalar_value or self._is_null:
142+
if args:
143+
return args[0]
144+
raise KeyError(key)
145+
return super().pop(key, *args)
146+
147+
def __len__(self):
148+
if self._is_null:
149+
return 0
150+
if self._is_array:
151+
return len(self._array_value)
152+
if self._is_scalar_value:
153+
return 1
154+
return super().__len__()
155+
156+
def __bool__(self):
157+
if self._is_null:
158+
return False
159+
if self._is_array:
160+
return bool(self._array_value)
161+
if self._is_scalar_value:
162+
return bool(self._simple_value)
163+
return super().__len__() > 0
164+
165+
def __iter__(self):
166+
if self._is_null:
167+
return iter([])
168+
if self._is_array:
169+
return iter(self._array_value)
170+
if self._is_scalar_value:
171+
raise TypeError(
172+
f"'{type(self._simple_value).__name__}' object is not iterable"
173+
)
174+
return super().__iter__()
175+
176+
def __getitem__(self, key):
177+
if self._is_array:
178+
return self._array_value[key]
179+
if self._is_scalar_value:
180+
raise TypeError(
181+
f"'{type(self._simple_value).__name__}' object is not subscriptable"
182+
)
183+
return super().__getitem__(key)
184+
185+
def __contains__(self, item):
186+
if self._is_null:
187+
return False
188+
if self._is_array:
189+
return item in self._array_value
190+
if self._is_scalar_value:
191+
raise TypeError(
192+
f"argument of type '{type(self._simple_value).__name__}' "
193+
"is not iterable"
194+
)
195+
return super().__contains__(item)
196+
197+
def __eq__(self, other):
198+
if isinstance(other, JsonObject):
52199
if self._is_array:
53-
self._array_value = args[0]._array_value
54-
elif self._is_scalar_value:
55-
self._simple_value = args[0]._simple_value
200+
return other._is_array and self._array_value == other._array_value
201+
if self._is_scalar_value:
202+
return (
203+
other._is_scalar_value and self._simple_value == other._simple_value
204+
)
205+
if self._is_null:
206+
return other._is_null
207+
return not (
208+
other._is_array or other._is_scalar_value or other._is_null
209+
) and super().__eq__(other)
210+
if self._is_array:
211+
return self._array_value == other
212+
if self._is_scalar_value:
213+
return self._simple_value == other
214+
if self._is_null:
215+
return other is None
216+
return super().__eq__(other)
56217

57-
if not self._is_null:
58-
super(JsonObject, self).__init__(*args, **kwargs)
218+
def __ne__(self, other):
219+
return not (self == other)
59220

60221
def __repr__(self):
61222
if self._is_array:
@@ -64,7 +225,7 @@ def __repr__(self):
64225
if self._is_scalar_value:
65226
return str(self._simple_value)
66227

67-
return super(JsonObject, self).__repr__()
228+
return super().__repr__()
68229

69230
@classmethod
70231
def from_str(cls, str_repr):
@@ -90,17 +251,33 @@ def serialize(self):
90251
if self._is_null:
91252
return None
92253

254+
raw = _unwrap_for_json(self)
93255
if self._is_scalar_value:
94-
return json.dumps(self._simple_value)
256+
return json.dumps(raw)
257+
258+
return json.dumps(raw, sort_keys=True, separators=(",", ":"))
95259

96-
if self._is_array:
97-
return json.dumps(self._array_value, sort_keys=True, separators=(",", ":"))
98260

99-
return json.dumps(self, sort_keys=True, separators=(",", ":"))
261+
def _unwrap_for_json(val):
262+
"""Recursively unwrap JsonObject instances for safe json.dumps serialization."""
263+
if isinstance(val, JsonObject):
264+
if val._is_null:
265+
return None
266+
if val._is_array:
267+
return [_unwrap_for_json(item) for item in val._array_value]
268+
if val._is_scalar_value:
269+
return val._simple_value
270+
return {k: _unwrap_for_json(v) for k, v in val.items()}
271+
if isinstance(val, dict):
272+
return {k: _unwrap_for_json(v) for k, v in val.items()}
273+
if isinstance(val, (list, tuple)):
274+
return [_unwrap_for_json(item) for item in val]
275+
return val
100276

101277

102278
_INTERVAL_PATTERN = re.compile(
103-
r"^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$"
279+
r"^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?"
280+
r"(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$"
104281
)
105282

106283

@@ -109,7 +286,7 @@ class Interval:
109286
"""Represents a Spanner INTERVAL type.
110287
111288
An interval is a combination of months, days and nanoseconds.
112-
Internally, Spanner supports Interval value with the following range of individual fields:
289+
Internally, Spanner supports Interval value with individual fields range:
113290
months: [-120000, 120000]
114291
days: [-3660000, 3660000]
115292
nanoseconds: [-316224000000000000000, 316224000000000000000]
@@ -199,12 +376,14 @@ def from_str(cls, s: str) -> "Interval":
199376
parts = match.groups()
200377
if not any(parts[:3]) and not parts[3]:
201378
raise ValueError(
202-
f"Invalid interval format: at least one component (Y/M/D/H/M/S) is required: {s}"
379+
"Invalid interval format: at least one component "
380+
f"(Y/M/D/H/M/S) is required: {s}"
203381
)
204382

205383
if parts[3] == "T" and not any(parts[4:7]):
206384
raise ValueError(
207-
f"Invalid interval format: time designator 'T' present but no time components specified: {s}"
385+
"Invalid interval format: time designator 'T' present "
386+
f"but no time components specified: {s}"
208387
)
209388

210389
def parse_num(s: str, suffix: str) -> int:
@@ -298,14 +477,14 @@ def _proto_enum(int_val, proto_enum_object):
298477

299478

300479
def get_proto_message(bytes_string, proto_message_object):
301-
"""parses serialized protocol buffer bytes' data or its list into proto message or list of proto message.
480+
"""Parses serialized protocol buffer bytes data or list into proto message.
302481
303482
Args:
304483
bytes_string (bytes or list[bytes]): bytes object.
305484
proto_message_object (Message): Message object for parsing
306485
307486
Returns:
308-
Message or list[Message]: parses serialized protocol buffer data into this message.
487+
Message or list[Message]: Parsed protocol buffer message(s).
309488
310489
Raises:
311490
ValueError: if the input proto_message_object is not of type Message

packages/google-cloud-spanner/tests/unit/test__helpers.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -868,6 +868,7 @@ def test_w_json(self):
868868
from google.protobuf.struct_pb2 import Value
869869

870870
from google.cloud.spanner_v1 import Type, TypeCode
871+
from google.cloud.spanner_v1.data_types import JsonObject
871872

872873
VALUE = {"id": 27863, "Name": "Anamika"}
873874
str_repr = json.dumps(VALUE, sort_keys=True, separators=(",", ":"))
@@ -884,7 +885,9 @@ def test_w_json(self):
884885
field_type = Type(code=TypeCode.JSON)
885886
value_pb = Value(string_value=str_repr)
886887

887-
self.assertEqual(self._callFUT(value_pb, field_type, field_name), {})
888+
self.assertEqual(
889+
self._callFUT(value_pb, field_type, field_name), JsonObject(None)
890+
)
888891

889892
def test_w_unknown_type(self):
890893
from google.protobuf.struct_pb2 import Value

0 commit comments

Comments
 (0)