WIP accelerated encoding with orjson

plotly · nicolaskruchten · May 27, 2021 · Dec 5, 2020 · Dec 5, 2020 · Dec 5, 2020
commit 40b9af19edc60e2d5b1eb5630f321938bbb21dee
diff --git a/packages/python/plotly/_plotly_utils/utils.py b/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

diff --git a/packages/python/plotly/plotly/basedatatypes.py b/packages/python/plotly/plotly/basedatatypes.py
@@ -3273,7 +3273,7 @@ def _perform_batch_animate(self, animation_opts):
 
     # Exports
     # -------
-    def to_dict(self):
+    def to_dict(self, clone=True):
         """
         Convert figure to a dictionary
 
@@ -3286,31 +3286,41 @@ def to_dict(self):
         """
         # Handle data
         # -----------
-        data = deepcopy(self._data)
+        if clone:
+            data = deepcopy(self._data)
+        else:
+            data = self._data
 
         # Handle layout
         # -------------
-        layout = deepcopy(self._layout)
+        if clone:
+            layout = deepcopy(self._layout)
+        else:
+            layout = self._layout
 
         # Handle frames
         # -------------
         # 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 clone:
+            frames = deepcopy([frame._props for frame in self._frame_objs])
+        else:
+            frames = [frame._props for frame in self._frame_objs]
+
         if frames:
             res["frames"] = frames
 
         return res
 
-    def to_plotly_json(self):
+    def to_plotly_json(self, clone=True):
         """
         Convert figure to a JSON representation as a Python dict
 
         Returns
         -------
         dict
         """
-        return self.to_dict()
+        return self.to_dict(clone=clone)
 
     @staticmethod
     def _to_ordered_dict(d, skip_uid=False):
@@ -5524,15 +5534,18 @@ def on_change(self, callback, *args, **kwargs):
         # -----------------
         self._change_callbacks[arg_tuples].append(callback)
 
-    def to_plotly_json(self):
+    def to_plotly_json(self, clone=False):
         """
         Return plotly JSON representation of object as a Python dict
 
         Returns
         -------
         dict
         """
-        return deepcopy(self._props if self._props is not None else {})
+        if clone:
+            return deepcopy(self._props if self._props is not None else {})
+        else:
+            return self._props if self._props is not None else {}
 
     @staticmethod
     def _vals_equal(v1, v2):

diff --git a/packages/python/plotly/plotly/io/_json.py b/packages/python/plotly/plotly/io/_json.py
@@ -2,12 +2,28 @@
 
 from six import string_types
 import json
+import decimal
 
 
 from plotly.io._utils import validate_coerce_fig_to_dict, validate_coerce_output_type
+from _plotly_utils.utils import iso_to_plotly_time_string
+from _plotly_utils.optional_imports import get_module
+from _plotly_utils.basevalidators import ImageUriValidator
 
 
-def to_json(fig, validate=True, pretty=False, remove_uids=True):
+def coerce_to_strict(const):
+    """
+    This is used to ultimately *encode* into strict JSON, see `encode`
+
+    """
+    # before python 2.7, 'true', 'false', 'null', were include here.
+    if const in ("Infinity", "-Infinity", "NaN"):
+        return None
+    else:
+        return const
+
+
+def to_json(fig, validate=True, pretty=False, remove_uids=True, engine="auto"):
     """
     Convert a figure to a JSON string representation
 
@@ -32,7 +48,7 @@ def to_json(fig, validate=True, pretty=False, remove_uids=True):
     str
         Representation of figure as a JSON string
     """
-    from _plotly_utils.utils import PlotlyJSONEncoder
+    orjson = get_module("orjson", should_load=True)
 
     # Validate figure
     # ---------------
@@ -44,16 +60,77 @@ def to_json(fig, validate=True, pretty=False, remove_uids=True):
         for trace in fig_dict.get("data", []):
             trace.pop("uid", None)
 
+    # Determine json engine
+    if engine == "auto":
+        if orjson is not None:
+            engine = "orjson"
+        else:
+            engine = "json"
+    elif engine not in ["orjson", "json", "legacy"]:
+        raise ValueError("Invalid json engine: %s" % engine)
+
+    modules = {"sage_all": get_module("sage.all", should_load=False),
+               "np": get_module("numpy", should_load=False),
+               "pd": get_module("pandas", should_load=False),
+               "image": get_module("PIL.Image", should_load=False)}
+
+    orjson = get_module("orjson", should_load=True)
+
     # Dump to a JSON string and return
     # --------------------------------
-    opts = {"sort_keys": True}
-    if pretty:
-        opts["indent"] = 2
-    else:
-        # Remove all whitespace
-        opts["separators"] = (",", ":")
-
-    return json.dumps(fig_dict, cls=PlotlyJSONEncoder, **opts)
+    if engine in ("json", "legacy"):
+        opts = {"sort_keys": True}
+        if pretty:
+            opts["indent"] = 2
+        else:
+            # Remove all whitespace
+            opts["separators"] = (",", ":")
+
+        if engine == "json":
+            cleaned = clean_to_json_compatible(
+                fig, numpy_allowed=False,
+                non_finite_allowed=False,
+                datetime_allowed=False,
+                modules=modules,
+            )
+            encoded_o = json.dumps(cleaned, **opts)
+
+            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
+            try:
+                new_o = json.loads(encoded_o, parse_constant=coerce_to_strict)
+            except ValueError:
+
+                # invalid separators will fail here. raise a helpful exception
+                raise ValueError(
+                    "Encoding into strict JSON failed. Did you set the separators "
+                    "valid JSON separators?"
+                )
+            else:
+                return json.dumps(new_o, **opts)
+        else:
+            from _plotly_utils.utils import PlotlyJSONEncoder
+            return json.dumps(fig_dict, cls=PlotlyJSONEncoder, **opts)
+    elif engine == "orjson":
+        opts = (orjson.OPT_SORT_KEYS
+                | orjson.OPT_SERIALIZE_NUMPY
+                | orjson.OPT_OMIT_MICROSECONDS
+                )
+
+        if pretty:
+            opts |= orjson.OPT_INDENT_2
+
+        cleaned = clean_to_json_compatible(
+            fig, numpy_allowed=True,
+            non_finite_allowed=True,
+            datetime_allowed=True,
+            modules=modules,
+        )
+        return orjson.dumps(cleaned, option=opts).decode("utf8")
 
 
 def write_json(fig, file, validate=True, pretty=False, remove_uids=True):
@@ -194,3 +271,86 @@ def read_json(file, output_type="Figure", skip_invalid=False):
     # Construct and return figure
     # ---------------------------
     return from_json(json_str, skip_invalid=skip_invalid, output_type=output_type)
+
+
+def clean_to_json_compatible(obj, **kwargs):
+    # Try handling value as a scalar value that we have a conversion for.
+    # Return immediately if we know we've hit a primitive value
+
+    # unpack kwargs
+    numpy_allowed = kwargs.get("numpy_allowed", False)
+    non_finite_allowed = kwargs.get("non_finite_allowed", False)
+    datetime_allowed = kwargs.get("datetime_allowed", False)
+
+    modules = kwargs.get("modules", {})
+    sage_all = modules["sage_all"]
+    np = modules["np"]
+    pd = modules["pd"]
+    image = modules["image"]
+
+    # Plotly
+    try:
+        obj = obj.to_plotly_json(clone=False)
+    except (TypeError, NameError, ValueError):
+        # Try without clone for backward compatibility
+        obj = obj.to_plotly_json()
+    except AttributeError:
+        pass
+
+    # Sage
+    if sage_all is not None:
+        if obj in sage_all.RR:
+            return float(obj)
+        elif obj in sage_all.ZZ:
+            return int(obj)
+
+    # numpy
+    if np is not None:
+        if obj is np.ma.core.masked:
+            return float("nan")
+        elif numpy_allowed and isinstance(obj, np.ndarray) and obj.dtype.kind in ("b", "i", "u", "f"):
+            return obj
+
+    # pandas
+    if pd is not None:
+        if obj is pd.NaT:
+            return None
+        elif isinstance(obj, pd.Series):
+            if numpy_allowed and obj.dtype.kind in ("b", "i", "u", "f"):
+                return obj.values
+            elif datetime_allowed and obj.dtype.kind == "M":
+                return obj.dt.to_pydatetime().tolist()
+
+
+    # datetime and date
+    if not datetime_allowed:
+        try:
+            # Is this cleanup still needed?
+            return iso_to_plotly_time_string(obj.isoformat())
+        except AttributeError:
+            pass
+
+    # Try .tolist() convertible
+    try:
+        # obj = obj.tolist()
+        return obj.tolist()
+    except AttributeError:
+        pass
+
+    # Do best we can with decimal
+    if isinstance(obj, decimal.Decimal):
+        return float(obj)
+
+    # PIL
+    if image is not None and isinstance(obj, image.Image):
+        return ImageUriValidator.pil_image_to_uri(obj)
+
+    # Recurse into lists and dictionaries
+    if isinstance(obj, dict):
+        return {k: clean_to_json_compatible(v, **kwargs) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        if obj:
+            # Must process list recursively even though it may be slow
+            return [clean_to_json_compatible(v, **kwargs) for v in obj]
+
+    return obj