Merge pull request #131 from predict-idlab/orjson_werkzeug

💪🏼  making orjson serialization more robust, see #118

Author: jonasvdd

plotly_resampler/figure_resampler/figure_resampler_interface.py

@@ -291,8 +291,10 @@ def _check_update_trace_data(
             s_res: pd.Series = downsampler.aggregate(
                 hf_series, hf_trace_data["max_n_samples"]
             )
-            trace["x"] = s_res.index
-            trace["y"] = s_res.values
+            # Also parse the data types to an orjson compatible format
+            # Note this can be removed once orjson supports f16
+            trace["x"] = self._parse_dtype_orjson(s_res.index)
+            trace["y"] = self._parse_dtype_orjson(s_res.values)
             # todo -> first draft & not MP safe
 
             agg_prefix, agg_suffix = ' <i style="color:#fc9944">~', "</i>"
@@ -700,10 +702,6 @@ def _parse_get_trace_props(
                 except ValueError:
                     hf_y = hf_y.astype("str")
 
-            # orjson encoding doesn't like to encode with uint8 & uint16 dtype
-            if str(hf_y.dtype) in ["uint8", "uint16"]:
-                hf_y = hf_y.astype("uint32")
-
             assert len(hf_x) == len(hf_y), "x and y have different length!"
         else:
             self._print(f"trace {trace['type']} is not a high-frequency trace")
@@ -1283,6 +1281,17 @@ def construct_update_data(
             layout_traces_list.append(trace_reduced)
         return layout_traces_list
 
+    @staticmethod
+    def _parse_dtype_orjson(series: np.ndarray) -> np.ndarray:
+        """Verify the orjson compatibility of the series and convert it if needed."""
+        # NOTE:
+        #    * float16 and float128 aren't supported with latest orjson versions (3.8.1)
+        #    * this method assumes that the it will not get a float128 series
+        # -> this method can be removed if orjson supports float16
+        if series.dtype in [np.float16]:
+            return series.astype(np.float32)
+        return series
+
     @staticmethod
     def _re_matches(regex: re.Pattern, strings: Iterable[str]) -> List[str]:
         """Returns all the items in ``strings`` which regex.match(es) ``regex``."""

tests/test_figure_resampler.py

@@ -101,6 +101,48 @@ def test_add_trace_not_resampling(float_series):
         hf_hovertext="hovertext",
     )
 
+def test_various_dtypes(float_series):
+    # List of dtypes supported by orjson >= 3.8
+    valid_dtype_list = [
+        np.bool_,
+        # ---- uints
+        np.uint8,
+        np.uint16,
+        np.uint32,
+        np.uint64,
+        # -------- ints
+        np.int8,
+        np.int16,
+        np.int32,
+        np.int64,
+        # -------- floats
+        np.float16,  # currently not supported by orjson
+        np.float32,
+        np.float64,
+    ]
+    for dtype in valid_dtype_list:
+        fig = FigureResampler(go.Figure(), default_n_shown_samples=1000)
+        # nb. datapoints > default_n_shown_samples
+        fig.add_trace(
+            go.Scatter(name="float_series"),
+            hf_x=float_series.index,
+            hf_y=float_series.astype(dtype),
+        )
+        fig.full_figure_for_development()
+
+    # List of dtypes not supported by orjson >= 3.8
+    invalid_dtype_list = [ np.float16 ]
+    for invalid_dtype in invalid_dtype_list:
+        fig = FigureResampler(go.Figure(), default_n_shown_samples=1000)
+        # nb. datapoints < default_n_shown_samples
+        with pytest.raises(TypeError):  
+            # if this test fails -> orjson supports f16 => remove casting frome code
+            fig.add_trace(
+                go.Scatter(name="float_series"),
+                hf_x=float_series.index[:500],
+                hf_y=float_series.astype(invalid_dtype)[:500],
+            )
+            fig.full_figure_for_development()
 
 def test_max_n_samples(float_series):
     s = float_series[:5000]