fix: json data type for non object values (#1236)

surbhigarg92 · web-flow · commit 0007be37a65f · 2024-11-14T00:56:59.000-08:00
* fix: json data type for non object values

* review comments
diff --git a/google/cloud/spanner_v1/data_types.py b/google/cloud/spanner_v1/data_types.py
@@ -31,17 +31,26 @@ class JsonObject(dict):
     def __init__(self, *args, **kwargs):
         self._is_null = (args, kwargs) == ((), {}) or args == (None,)
         self._is_array = len(args) and isinstance(args[0], (list, tuple))
+        self._is_scalar_value = len(args) == 1 and not isinstance(args[0], (list, dict))
 
         # if the JSON object is represented with an array,
         # the value is contained separately
         if self._is_array:
             self._array_value = args[0]
             return
 
+        # If it's a scalar value, set _simple_value and return early
+        if self._is_scalar_value:
+            self._simple_value = args[0]
+            return
+
         if len(args) and isinstance(args[0], JsonObject):
             self._is_array = args[0]._is_array
+            self._is_scalar_value = args[0]._is_scalar_value
             if self._is_array:
                 self._array_value = args[0]._array_value
+            elif self._is_scalar_value:
+                self._simple_value = args[0]._simple_value
 
         if not self._is_null:
             super(JsonObject, self).__init__(*args, **kwargs)
@@ -50,6 +59,9 @@ def __repr__(self):
         if self._is_array:
             return str(self._array_value)
 
+        if self._is_scalar_value:
+            return str(self._simple_value)
+
         return super(JsonObject, self).__repr__()
 
     @classmethod
@@ -76,6 +88,9 @@ def serialize(self):
         if self._is_null:
             return None
 
+        if self._is_scalar_value:
+            return json.dumps(self._simple_value)
+
         if self._is_array:
             return json.dumps(self._array_value, sort_keys=True, separators=(",", ":"))
 
diff --git a/tests/unit/test_datatypes.py b/tests/unit/test_datatypes.py
@@ -43,3 +43,56 @@ def test_w_JsonObject_of_list_of_dict(self):
         expected = json.dumps(data, sort_keys=True, separators=(",", ":"))
         data_jsonobject = JsonObject(JsonObject(data))
         self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_simple_float_JsonData(self):
+        data = 1.1
+        expected = json.dumps(data)
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_simple_str_JsonData(self):
+        data = "foo"
+        expected = json.dumps(data)
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_empty_str_JsonData(self):
+        data = ""
+        expected = json.dumps(data)
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_None_JsonData(self):
+        data = None
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), None)
+
+    def test_w_list_of_simple_JsonData(self):
+        data = [1.1, "foo"]
+        expected = json.dumps(data, sort_keys=True, separators=(",", ":"))
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_empty_list(self):
+        data = []
+        expected = json.dumps(data)
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_empty_dict(self):
+        data = [{}]
+        expected = json.dumps(data)
+        data_jsonobject = JsonObject(data)
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_JsonObject_of_simple_JsonData(self):
+        data = 1.1
+        expected = json.dumps(data)
+        data_jsonobject = JsonObject(JsonObject(data))
+        self.assertEqual(data_jsonobject.serialize(), expected)
+
+    def test_w_JsonObject_of_list_of_simple_JsonData(self):
+        data = [1.1, "foo"]
+        expected = json.dumps(data, sort_keys=True, separators=(",", ":"))
+        data_jsonobject = JsonObject(JsonObject(data))
+        self.assertEqual(data_jsonobject.serialize(), expected)
