Merge pull request #2955 from plotly/orjson_encoding

JSON encoding refactor and orjson encoding

Author: nicolaskruchten

doc/requirements.txt

@@ -31,3 +31,4 @@ umap-learn==0.5.1
 pooch
 wget
 nbconvert==5.6.1
+orjson

packages/python/plotly/_plotly_utils/basevalidators.py

@@ -53,7 +53,7 @@ def to_scalar_or_list(v):
         return v
 
 
-def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
+def copy_to_readonly_numpy_array_or_list(v, kind=None, force_numeric=False):
     """
     Convert an array-like value into a read-only numpy array
 
@@ -89,7 +89,13 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
 
     # u: unsigned int, i: signed int, f: float
     numeric_kinds = {"u", "i", "f"}
-    kind_default_dtypes = {"u": "uint32", "i": "int32", "f": "float64", "O": "object"}
+    kind_default_dtypes = {
+        "u": "uint32",
+        "i": "int32",
+        "f": "float64",
+        "O": "object",
+        "U": "U",
+    }
 
     # Handle pandas Series and Index objects
     if pd and isinstance(v, (pd.Series, pd.Index)):
@@ -113,18 +119,12 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
     if not isinstance(v, np.ndarray):
         # v has its own logic on how to convert itself into a numpy array
         if is_numpy_convertable(v):
-            return copy_to_readonly_numpy_array(
+            return copy_to_readonly_numpy_array_or_list(
                 np.array(v), kind=kind, force_numeric=force_numeric
             )
         else:
             # v is not homogenous array
-            v_list = [to_scalar_or_list(e) for e in v]
-
-            # Lookup dtype for requested kind, if any
-            dtype = kind_default_dtypes.get(first_kind, None)
-
-            # construct new array from list
-            new_v = np.array(v_list, order="C", dtype=dtype)
+            return [to_scalar_or_list(e) for e in v]
     elif v.dtype.kind in numeric_kinds:
         # v is a homogenous numeric array
         if kind and v.dtype.kind not in kind:
@@ -135,6 +135,12 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
         else:
             # Either no kind was requested or requested kind is satisfied
             new_v = np.ascontiguousarray(v.copy())
+    elif v.dtype.kind == "O":
+        if kind:
+            dtype = kind_default_dtypes.get(first_kind, None)
+            return np.array(v, dtype=dtype)
+        else:
+            return v.tolist()
     else:
         # v is a non-numeric homogenous array
         new_v = v.copy()
@@ -149,12 +155,12 @@ def copy_to_readonly_numpy_array(v, kind=None, force_numeric=False):
     if "U" not in kind:
         # Force non-numeric arrays to have object type
         # --------------------------------------------
-        # Here we make sure that non-numeric arrays have the object
-        # datatype. This works around cases like np.array([1, 2, '3']) where
+        # Here we make sure that non-numeric arrays become lists
+        # This works around cases like np.array([1, 2, '3']) where
         # numpy converts the integers to strings and returns array of dtype
         # '<U21'
         if new_v.dtype.kind not in ["u", "i", "f", "O", "M"]:
-            new_v = np.array(v, dtype="object")
+            return v.tolist()
 
     # Set new array to be read-only
     # -----------------------------
@@ -191,7 +197,7 @@ def is_homogeneous_array(v):
             if v_numpy.shape == ():
                 return False
             else:
-                return True
+                return True  # v_numpy.dtype.kind in ["u", "i", "f", "M", "U"]
     return False
 
 
@@ -393,7 +399,7 @@ def validate_coerce(self, v):
             # Pass None through
             pass
         elif is_homogeneous_array(v):
-            v = copy_to_readonly_numpy_array(v)
+            v = copy_to_readonly_numpy_array_or_list(v)
         elif is_simple_array(v):
             v = to_scalar_or_list(v)
         else:
@@ -598,7 +604,7 @@ def validate_coerce(self, v):
                 self.raise_invalid_elements(invalid_els[:10])
 
             if is_homogeneous_array(v):
-                v = copy_to_readonly_numpy_array(v)
+                v = copy_to_readonly_numpy_array_or_list(v)
             else:
                 v = to_scalar_or_list(v)
         else:
@@ -754,7 +760,7 @@ def validate_coerce(self, v):
         elif self.array_ok and is_homogeneous_array(v):
             np = get_module("numpy")
             try:
-                v_array = copy_to_readonly_numpy_array(v, force_numeric=True)
+                v_array = copy_to_readonly_numpy_array_or_list(v, force_numeric=True)
             except (ValueError, TypeError, OverflowError):
                 self.raise_invalid_val(v)
 
@@ -881,7 +887,7 @@ def validate_coerce(self, v):
             pass
         elif self.array_ok and is_homogeneous_array(v):
             np = get_module("numpy")
-            v_array = copy_to_readonly_numpy_array(
+            v_array = copy_to_readonly_numpy_array_or_list(
                 v, kind=("i", "u"), force_numeric=True
             )
 
@@ -1042,26 +1048,7 @@ def validate_coerce(self, v):
                 if invalid_els:
                     self.raise_invalid_elements(invalid_els)
 
-            if is_homogeneous_array(v):
-                np = get_module("numpy")
-
-                # If not strict, let numpy cast elements to strings
-                v = copy_to_readonly_numpy_array(v, kind="U")
-
-                # Check no_blank
-                if self.no_blank:
-                    invalid_els = v[v == ""][:10].tolist()
-                    if invalid_els:
-                        self.raise_invalid_elements(invalid_els)
-
-                # Check values
-                if self.values:
-                    invalid_inds = np.logical_not(np.isin(v, self.values))
-                    invalid_els = v[invalid_inds][:10].tolist()
-                    if invalid_els:
-                        self.raise_invalid_elements(invalid_els)
-
-            elif is_simple_array(v):
+            if is_simple_array(v) or is_homogeneous_array(v):
                 if not self.strict:
                     v = [StringValidator.to_str_or_unicode_or_none(e) for e in v]
 
@@ -1338,8 +1325,12 @@ def validate_coerce(self, v, should_raise=True):
             # Pass None through
             pass
         elif self.array_ok and is_homogeneous_array(v):
-            v = copy_to_readonly_numpy_array(v)
-            if self.numbers_allowed() and v.dtype.kind in ["u", "i", "f"]:
+            v = copy_to_readonly_numpy_array_or_list(v)
+            if (
+                not isinstance(v, list)
+                and self.numbers_allowed()
+                and v.dtype.kind in ["u", "i", "f"]
+            ):
                 # Numbers are allowed and we have an array of numbers.
                 # All good
                 pass
@@ -1353,9 +1344,9 @@ def validate_coerce(self, v, should_raise=True):
 
                 # ### Check that elements have valid colors types ###
                 elif self.numbers_allowed() or invalid_els:
-                    v = copy_to_readonly_numpy_array(validated_v, kind="O")
+                    v = copy_to_readonly_numpy_array_or_list(validated_v, kind="O")
                 else:
-                    v = copy_to_readonly_numpy_array(validated_v, kind="U")
+                    v = copy_to_readonly_numpy_array_or_list(validated_v, kind="U")
         elif self.array_ok and is_simple_array(v):
             validated_v = [self.validate_coerce(e, should_raise=False) for e in v]
 
@@ -1870,7 +1861,7 @@ def validate_coerce(self, v):
                 self.raise_invalid_elements(invalid_els)
 
             if is_homogeneous_array(v):
-                v = copy_to_readonly_numpy_array(validated_v, kind="U")
+                v = copy_to_readonly_numpy_array_or_list(validated_v, kind="U")
             else:
                 v = to_scalar_or_list(v)
         else:
@@ -1918,7 +1909,7 @@ def validate_coerce(self, v):
             # Pass None through
             pass
         elif self.array_ok and is_homogeneous_array(v):
-            v = copy_to_readonly_numpy_array(v, kind="O")
+            v = copy_to_readonly_numpy_array_or_list(v, kind="O")
         elif self.array_ok and is_simple_array(v):
             v = to_scalar_or_list(v)
         return v

packages/python/plotly/_plotly_utils/tests/validators/test_dataarray_validator.py

@@ -32,15 +32,29 @@ def test_validator_acceptance_simple(val, validator):
 
 
 @pytest.mark.parametrize(
-    "val",
-    [np.array([2, 3, 4]), pd.Series(["a", "b", "c"]), np.array([[1, 2, 3], [4, 5, 6]])],
+    "val", [np.array([2, 3, 4]), np.array([[1, 2, 3], [4, 5, 6]])],
 )
 def test_validator_acceptance_homogeneous(val, validator):
     coerce_val = validator.validate_coerce(val)
     assert isinstance(coerce_val, np.ndarray)
     assert np.array_equal(validator.present(coerce_val), val)
 
 
+# Accept object array as list
+@pytest.mark.parametrize(
+    "val",
+    [
+        ["A", "B", "C"],
+        np.array(["A", "B", "C"], dtype="object"),
+        pd.Series(["a", "b", "c"]),
+    ],
+)
+def test_validator_accept_object_array_as_list(val, validator):
+    coerce_val = validator.validate_coerce(val)
+    assert isinstance(coerce_val, list)
+    assert coerce_val == list(val)
+
+
 # ### Rejection ###
 @pytest.mark.parametrize("val", ["Hello", 23, set(), {}])
 def test_rejection(val, validator):

packages/python/plotly/_plotly_utils/tests/validators/test_enumerated_validator.py

@@ -126,7 +126,7 @@ def test_rejection_by_element_aok(val, validator_aok):
         [],
         ["bar12"],
         ("foo", "bar012", "baz"),
-        np.array([]),
+        np.array([], dtype="object"),
         np.array(["bar12"]),
         np.array(["foo", "bar012", "baz"]),
     ],
@@ -135,7 +135,7 @@ def test_acceptance_aok(val, validator_aok_re):
     # Values should be accepted and returned unchanged
     coerce_val = validator_aok_re.validate_coerce(val)
     if isinstance(val, (np.ndarray, pd.Series)):
-        assert np.array_equal(coerce_val, np.array(val, dtype=coerce_val.dtype))
+        assert coerce_val == list(np.array(val))
     elif isinstance(val, (list, tuple)):
         assert validator_aok_re.present(coerce_val) == tuple(val)
     else:

packages/python/plotly/_plotly_utils/tests/validators/test_pandas_series_input.py

@@ -149,13 +149,10 @@ def test_color_validator_object(color_validator, color_object_pandas):
     res = color_validator.validate_coerce(color_object_pandas)
 
     # Check type
-    assert isinstance(res, np.ndarray)
-
-    # Check dtype
-    assert res.dtype == "object"
+    assert isinstance(res, list)
 
     # Check values
-    np.testing.assert_array_equal(res, color_object_pandas)
+    assert res == color_object_pandas.tolist()
 
 
 def test_color_validator_categorical(color_validator, color_categorical_pandas):
@@ -164,13 +161,10 @@ def test_color_validator_categorical(color_validator, color_categorical_pandas):
 
     # Check type
     assert color_categorical_pandas.dtype == "category"
-    assert isinstance(res, np.ndarray)
-
-    # Check dtype
-    assert res.dtype == "object"
+    assert isinstance(res, list)
 
     # Check values
-    np.testing.assert_array_equal(res, np.array(color_categorical_pandas))
+    assert res == color_categorical_pandas.tolist()
 
 
 def test_data_array_validator_dates_series(
@@ -180,13 +174,10 @@ def test_data_array_validator_dates_series(
     res = data_array_validator.validate_coerce(datetime_pandas)
 
     # Check type
-    assert isinstance(res, np.ndarray)
-
-    # Check dtype
-    assert res.dtype == "object"
+    assert isinstance(res, list)
 
     # Check values
-    np.testing.assert_array_equal(res, dates_array)
+    assert res == dates_array.tolist()
 
 
 def test_data_array_validator_dates_dataframe(
@@ -197,10 +188,7 @@ def test_data_array_validator_dates_dataframe(
     res = data_array_validator.validate_coerce(df)
 
     # Check type
-    assert isinstance(res, np.ndarray)
-
-    # Check dtype
-    assert res.dtype == "object"
+    assert isinstance(res, list)
 
     # Check values
-    np.testing.assert_array_equal(res, dates_array.reshape(len(dates_array), 1))
+    assert res == dates_array.reshape(len(dates_array), 1).tolist()

packages/python/plotly/_plotly_utils/tests/validators/test_string_validator.py

@@ -138,8 +138,7 @@ def test_acceptance_aok_scalars(val, validator_aok):
 def test_acceptance_aok_list(val, validator_aok):
     coerce_val = validator_aok.validate_coerce(val)
     if isinstance(val, np.ndarray):
-        assert isinstance(coerce_val, np.ndarray)
-        assert np.array_equal(coerce_val, np.array(val, dtype=coerce_val.dtype))
+        assert coerce_val == val.tolist()
     elif isinstance(val, list):
         assert validator_aok.present(val) == tuple(val)
     else:
@@ -178,9 +177,7 @@ def test_rejection_aok_values(val, validator_aok_values):
 )
 def test_acceptance_no_blanks_aok(val, validator_no_blanks_aok):
     coerce_val = validator_no_blanks_aok.validate_coerce(val)
-    if isinstance(val, np.ndarray):
-        assert np.array_equal(coerce_val, np.array(val, dtype=coerce_val.dtype))
-    elif isinstance(val, list):
+    if isinstance(val, (list, np.ndarray)):
         assert validator_no_blanks_aok.present(coerce_val) == tuple(val)
     else:
         assert coerce_val == val

packages/python/plotly/_plotly_utils/tests/validators/test_xarray_input.py

@@ -126,10 +126,7 @@ def test_color_validator_object(color_validator, color_object_xarray):
     res = color_validator.validate_coerce(color_object_xarray)
 
     # Check type
-    assert isinstance(res, np.ndarray)
-
-    # Check dtype
-    assert res.dtype == "object"
+    assert isinstance(res, list)
 
     # Check values
-    np.testing.assert_array_equal(res, color_object_xarray)
+    assert res == list(color_object_xarray)

packages/python/plotly/_plotly_utils/utils.py

@@ -61,8 +61,10 @@ def encode(self, o):
         # We catch false positive cases (e.g. strings such as titles, labels etc.)
         # but this is ok since the intention is to skip the decoding / reencoding
         # step when it's completely safe
+
         if not ("NaN" in encoded_o or "Infinity" in encoded_o):
             return encoded_o
+
         # now:
         #    1. `loads` to switch Infinity, -Infinity, NaN to None
         #    2. `dumps` again so you get 'null' instead of extended JSON

packages/python/plotly/plotly/basedatatypes.py

@@ -3297,6 +3297,7 @@ def to_dict(self):
         # Frame key is only added if there are any frames
         res = {"data": data, "layout": layout}
         frames = deepcopy([frame._props for frame in self._frame_objs])
+
         if frames:
             res["frames"] = frames
 
@@ -3413,6 +3414,13 @@ def to_json(self, *args, **kwargs):
         remove_uids: bool (default True)
             True if trace UIDs should be omitted from the JSON representation
 
+        engine: str (default None)
+            The JSON encoding engine to use. One of:
+              - "json" for an encoder based on the built-in Python json module
+              - "orjson" for a fast encoder the requires the orjson package
+            If not specified, the default encoder is set to the current value of
+            plotly.io.json.config.default_encoder.
+
         Returns
         -------
         str
@@ -3469,6 +3477,13 @@ def write_json(self, *args, **kwargs):
         remove_uids: bool (default True)
             True if trace UIDs should be omitted from the JSON representation
 
+        engine: str (default None)
+            The JSON encoding engine to use. One of:
+              - "json" for an encoder based on the built-in Python json module
+              - "orjson" for a fast encoder the requires the orjson package
+            If not specified, the default encoder is set to the current value of
+            plotly.io.json.config.default_encoder.
+
         Returns
         -------
         None