fix(spanner): implement dict protocol and nested unwrapping for JsonObject

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

@@ -55,7 +55,80 @@ def __init__(self, *args, **kwargs):
                 self._simple_value = args[0]._simple_value
 
         if not self._is_null:
-            super(JsonObject, self).__init__(*args, **kwargs)
+            super().__init__(*args, **kwargs)
+
+    def __len__(self):
+        if self._is_null:
+            return 0
+        if self._is_array:
+            return len(self._array_value)
+        if self._is_scalar_value:
+            return 1
+        return super().__len__()
+
+    def __bool__(self):
+        if self._is_null:
+            return False
+        if self._is_array:
+            return bool(self._array_value)
+        if self._is_scalar_value:
+            return True
+        return super().__len__() > 0
+
+    def __iter__(self):
+        if self._is_array:
+            return iter(self._array_value)
+        if self._is_scalar_value:
+            raise TypeError(
+                f"'{type(self._simple_value).__name__}' object is not iterable"
+            )
+        return super().__iter__()
+
+    def __getitem__(self, key):
+        if self._is_array:
+            return self._array_value[key]
+        if self._is_scalar_value:
+            raise TypeError(
+                f"'{type(self._simple_value).__name__}' object is not subscriptable"
+            )
+        return super().__getitem__(key)
+
+    def __contains__(self, item):
+        if self._is_array:
+            return item in self._array_value
+        if self._is_scalar_value:
+            raise TypeError(
+                f"argument of type '{type(self._simple_value).__name__}' "
+                "is not iterable"
+            )
+        return super().__contains__(item)
+
+    def __eq__(self, other):
+        if isinstance(other, JsonObject):
+            if self._is_array:
+                return (
+                    getattr(other, "_is_array", False)
+                    and self._array_value == other._array_value
+                )
+            if self._is_scalar_value:
+                return (
+                    getattr(other, "_is_scalar_value", False)
+                    and self._simple_value == other._simple_value
+                )
+            if self._is_null:
+                return getattr(other, "_is_null", False)
+            return not (
+                getattr(other, "_is_array", False)
+                or getattr(other, "_is_scalar_value", False)
+                or getattr(other, "_is_null", False)
+            ) and super().__eq__(other)
+        if self._is_array:
+            return self._array_value == other
+        if self._is_scalar_value:
+            return self._simple_value == other
+        if self._is_null:
+            return other is None or (isinstance(other, dict) and len(other) == 0)
+        return super().__eq__(other)
 
     def __repr__(self):
         if self._is_array:
@@ -64,7 +137,7 @@ def __repr__(self):
         if self._is_scalar_value:
             return str(self._simple_value)
 
-        return super(JsonObject, self).__repr__()
+        return super().__repr__()
 
     @classmethod
     def from_str(cls, str_repr):
@@ -90,17 +163,33 @@ def serialize(self):
         if self._is_null:
             return None
 
+        raw = _unwrap_for_json(self)
         if self._is_scalar_value:
-            return json.dumps(self._simple_value)
+            return json.dumps(raw)
+
+        return json.dumps(raw, sort_keys=True, separators=(",", ":"))
 
-        if self._is_array:
-            return json.dumps(self._array_value, sort_keys=True, separators=(",", ":"))
 
-        return json.dumps(self, sort_keys=True, separators=(",", ":"))
+def _unwrap_for_json(val):
+    """Recursively unwrap JsonObject instances for safe json.dumps serialization."""
+    if isinstance(val, JsonObject):
+        if val._is_null:
+            return None
+        if val._is_array:
+            return [_unwrap_for_json(item) for item in val._array_value]
+        if val._is_scalar_value:
+            return val._simple_value
+        return {k: _unwrap_for_json(v) for k, v in val.items()}
+    if isinstance(val, dict):
+        return {k: _unwrap_for_json(v) for k, v in val.items()}
+    if isinstance(val, (list, tuple)):
+        return [_unwrap_for_json(item) for item in val]
+    return val
 
 
 _INTERVAL_PATTERN = re.compile(
-    r"^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$"
+    r"^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?"
+    r"(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$"
 )
 
 
@@ -109,7 +198,7 @@ class Interval:
     """Represents a Spanner INTERVAL type.
 
     An interval is a combination of months, days and nanoseconds.
-    Internally, Spanner supports Interval value with the following range of individual fields:
+    Internally, Spanner supports Interval value with individual fields range:
     months: [-120000, 120000]
     days: [-3660000, 3660000]
     nanoseconds: [-316224000000000000000, 316224000000000000000]
@@ -199,12 +288,14 @@ def from_str(cls, s: str) -> "Interval":
         parts = match.groups()
         if not any(parts[:3]) and not parts[3]:
             raise ValueError(
-                f"Invalid interval format: at least one component (Y/M/D/H/M/S) is required: {s}"
+                "Invalid interval format: at least one component "
+                f"(Y/M/D/H/M/S) is required: {s}"
             )
 
         if parts[3] == "T" and not any(parts[4:7]):
             raise ValueError(
-                f"Invalid interval format: time designator 'T' present but no time components specified: {s}"
+                "Invalid interval format: time designator 'T' present "
+                f"but no time components specified: {s}"
             )
 
         def parse_num(s: str, suffix: str) -> int:
@@ -298,14 +389,14 @@ def _proto_enum(int_val, proto_enum_object):
 
 
 def get_proto_message(bytes_string, proto_message_object):
-    """parses serialized protocol buffer bytes' data or its list into proto message or list of proto message.
+    """Parses serialized protocol buffer bytes data or list into proto message.
 
     Args:
         bytes_string (bytes or list[bytes]): bytes object.
         proto_message_object (Message): Message object for parsing
 
     Returns:
-        Message or list[Message]: parses serialized protocol buffer data into this message.
+        Message or list[Message]: Parsed protocol buffer message(s).
 
     Raises:
         ValueError: if the input proto_message_object is not of type Message