Add a msgspec JSON engine

Author: T4rk1n

CHANGELOG.md

@@ -4,6 +4,9 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 
+### Added
+- Add a `msgspec` JSON engine, used by the `auto` engine when `msgspec` is installed and `orjson` is not [[#5813](https://github.com/plotly/plotly.py/pull/5813)]
+
 ### Updated
 - Speed up `plotly.io.to_json` and `PlotlyJSONEncoder` for figures and for payloads that mix plotly objects with plain data: figures are no longer deep copied, and with orjson only the values orjson can't serialize go through the slow cleaning path [[#5813](https://github.com/plotly/plotly.py/pull/5813)]
 

doc/python/renderers.md

@@ -271,6 +271,6 @@ No matter the approach chosen to display a figure, [the figure data structure](h
 
 *New in v5.0*
 
-The default JSON serialization mechanism can be slow for figures with many data points or with large `numpy` arrays or data frames. **If [the `orjson` package](https://github.com/ijl/orjson) is installed**, `plotly` will use that instead of the built-in `json` package, which can lead to **5-10x** speedups for large figures.
+The default JSON serialization mechanism can be slow for figures with many data points or with large `numpy` arrays or data frames. **If [the `orjson` package](https://github.com/ijl/orjson) is installed**, `plotly` will use that instead of the built-in `json` package, which can lead to **5-10x** speedups for large figures. If `orjson` isn't installed but [`msgspec`](https://github.com/jcrist/msgspec) is (Dash installs it), `plotly` uses `msgspec`, which gives most of the same speedup.
 
 Once a figure is serialized to JSON, it must be rendered by a browser, either immediately in the user's browser, at some later point if the figure is exported to HTML, or immediately in Kaleido's internal headless browser for static image export. Rendering time is generally proportional to the total number of data points in the figure, the number of traces and the number of subplots. In situations where rendering performance is slow, we recommend considering [the use of `plotly` WebGL traces](/python/webgl-vs-svg/) to exploit GPU-accelerated rendering in the browser, or [using the Datashader library to do Python-side rendering](/python/datashader/) before using `px.imshow()` to render the figure.

plotly/basedatatypes.py

@@ -3458,6 +3458,7 @@ def to_json(self, *args, **kwargs):
             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
+              - "msgspec" for a fast encoder that requires the msgspec package
             If not specified, the default encoder is set to the current value of
             plotly.io.json.config.default_encoder.
 
@@ -3518,6 +3519,7 @@ def write_json(self, *args, **kwargs):
             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
+              - "msgspec" for a fast encoder that requires the msgspec package
             If not specified, the default encoder is set to the current value of
             plotly.io.json.config.default_encoder.
 
@@ -5620,6 +5622,7 @@ def to_json(self, *args, **kwargs):
             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
+              - "msgspec" for a fast encoder that requires the msgspec package
             If not specified, the default encoder is set to the current value of
             plotly.io.json.config.default_encoder.
 

plotly/io/_json.py

@@ -10,7 +10,7 @@
 
 
 class JsonConfig(object):
-    _valid_engines = ("json", "orjson", "auto")
+    _valid_engines = ("json", "orjson", "msgspec", "auto")
 
     def __init__(self):
         self._default_engine = "auto"
@@ -30,6 +30,8 @@ def default_engine(self, val):
 
         if val == "orjson":
             self.validate_orjson()
+        elif val == "msgspec":
+            self.validate_msgspec()
 
         self._default_engine = val
 
@@ -39,6 +41,12 @@ def validate_orjson(cls):
         if orjson is None:
             raise ValueError("The orjson engine requires the orjson package")
 
+    @classmethod
+    def validate_msgspec(cls):
+        msgspec = get_module("msgspec")
+        if msgspec is None:
+            raise ValueError("The msgspec engine requires the msgspec package")
+
 
 config = JsonConfig()
 
@@ -66,6 +74,20 @@ def coerce_to_strict(const):
 )
 
 
+def _resolve_engine(engine):
+    if engine is None:
+        engine = config.default_engine
+    if engine == "auto":
+        if get_module("orjson", should_load=True) is not None:
+            return "orjson"
+        if get_module("msgspec", should_load=True) is not None:
+            return "msgspec"
+        return "json"
+    if engine not in ("orjson", "msgspec", "json"):
+        raise ValueError("Invalid json engine: %s" % engine)
+    return engine
+
+
 def _safe(json_str, _swap):
     out = json_str
     for unsafe_char, safe_char in _swap:
@@ -91,7 +113,9 @@ def to_json_plotly(plotly_object, pretty=False, engine=None):
         The JSON encoding engine to use. One of:
           - "json" for an engine based on the built-in Python json module
           - "orjson" for a faster engine that requires the orjson package
-          - "auto" for the "orjson" engine if available, otherwise "json"
+          - "msgspec" for a faster engine that requires the msgspec package
+          - "auto" for the "orjson" engine if available, then "msgspec",
+            otherwise "json"
         If not specified, the default engine is set to the current value of
         plotly.io.json.config.default_engine.
 
@@ -104,19 +128,7 @@ def to_json_plotly(plotly_object, pretty=False, engine=None):
     --------
     to_json : Convert a plotly Figure to JSON with validation
     """
-    orjson = get_module("orjson", should_load=True)
-
-    # Determine json engine
-    if engine is None:
-        engine = config.default_engine
-
-    if engine == "auto":
-        if orjson is not None:
-            engine = "orjson"
-        else:
-            engine = "json"
-    elif engine not in ["orjson", "json"]:
-        raise ValueError("Invalid json engine: %s" % engine)
+    engine = _resolve_engine(engine)
 
     modules = {
         "sage_all": get_module("sage.all", should_load=False),
@@ -140,8 +152,60 @@ def to_json_plotly(plotly_object, pretty=False, engine=None):
         return _safe(
             json.dumps(plotly_object, cls=PlotlyJSONEncoder, **opts), _swap_json
         )
+    elif engine == "msgspec":
+        JsonConfig.validate_msgspec()
+        msgspec = get_module("msgspec", should_load=True)
+        pd = modules["pd"]
+
+        def enc_hook(value):
+            # msgspec takes only exact types: numpy scalars subclass these
+            if isinstance(value, float):
+                return float(value)
+            if isinstance(value, int):
+                return int(value)
+            if isinstance(value, str):
+                return str.__str__(value)
+            if isinstance(value, (dict, list, tuple)):
+                # only asked for when they sit where JSON can't hold them
+                # (a tuple dict key), converting would loop forever
+                raise TypeError(
+                    f"Type is not JSON serializable: {type(value).__name__}"
+                )
+            if pd is not None:
+                if value is pd.NaT:
+                    return None
+                # isoformat keeps nanoseconds, as in the json engine
+                if isinstance(value, pd.Timestamp):
+                    return value.isoformat()
+                if isinstance(value, (pd.Series, pd.Index)) and value.dtype.kind == "M":
+                    return value.tolist()
+            to_json_dict = getattr(value, "_to_json_dict", None)
+            if to_json_dict is not None:
+                return to_json_dict()
+            cleaned = clean_to_json_compatible(
+                value,
+                numpy_allowed=False,
+                datetime_allowed=False,
+                modules=modules,
+            )
+            if cleaned is value:
+                raise TypeError(
+                    f"Type is not JSON serializable: {type(value).__name__}"
+                )
+            return cleaned
+
+        encoder = msgspec.json.Encoder(enc_hook=enc_hook, decimal_format="number")
+        try:
+            out = encoder.encode(plotly_object)
+        except (TypeError, ValueError, OverflowError, msgspec.EncodeError):
+            # Errors and values msgspec rejects come out as with the json engine
+            return to_json_plotly(plotly_object, pretty=pretty, engine="json")
+        if pretty:
+            out = msgspec.json.format(out, indent=2)
+        return _safe(out.decode("utf8"), _swap_orjson)
     elif engine == "orjson":
         JsonConfig.validate_orjson()
+        orjson = get_module("orjson", should_load=True)
         opts = orjson.OPT_NON_STR_KEYS | orjson.OPT_SERIALIZE_NUMPY
 
         if pretty:
@@ -213,7 +277,9 @@ def to_json(fig, validate=True, pretty=False, remove_uids=True, engine=None):
         The JSON encoding engine to use. One of:
           - "json" for an engine based on the built-in Python json module
           - "orjson" for a faster engine that requires the orjson package
-          - "auto" for the "orjson" engine if available, otherwise "json"
+          - "msgspec" for a faster engine that requires the msgspec package
+          - "auto" for the "orjson" engine if available, then "msgspec",
+            otherwise "json"
         If not specified, the default engine is set to the current value of
         plotly.io.json.config.default_engine.
 
@@ -266,7 +332,9 @@ def write_json(fig, file, validate=True, pretty=False, remove_uids=True, engine=
         The JSON encoding engine to use. One of:
           - "json" for an engine based on the built-in Python json module
           - "orjson" for a faster engine that requires the orjson package
-          - "auto" for the "orjson" engine if available, otherwise "json"
+          - "msgspec" for a faster engine that requires the msgspec package
+          - "auto" for the "orjson" engine if available, then "msgspec",
+            otherwise "json"
         If not specified, the default engine is set to the current value of
         plotly.io.json.config.default_engine.
     Returns
@@ -329,7 +397,10 @@ def from_json_plotly(value, engine=None):
           - if "json", parse JSON using built in json module
           - if "orjson", parse using the faster orjson module, requires the orjson
             package
-          - if "auto" use orjson module if available, otherwise use the json module
+          - if "msgspec", parse using the faster msgspec module, requires the
+            msgspec package
+          - if "auto" use orjson module if available, then msgspec, otherwise
+            use the json module
 
         If not specified, the default engine is set to the current value of
         plotly.io.json.config.default_engine.
@@ -342,8 +413,6 @@ def from_json_plotly(value, engine=None):
     --------
     from_json_plotly : Parse JSON with plotly conventions into a dict
     """
-    orjson = get_module("orjson", should_load=True)
-
     # Validate value
     # --------------
     if not isinstance(value, (str, bytes)):
@@ -353,22 +422,15 @@ def from_json_plotly(value, engine=None):
     Received value: {value}""".format(typ=type(value), value=value)
         )
 
-    # Determine json engine
-    if engine is None:
-        engine = config.default_engine
-
-    if engine == "auto":
-        if orjson is not None:
-            engine = "orjson"
-        else:
-            engine = "json"
-    elif engine not in ["orjson", "json"]:
-        raise ValueError("Invalid json engine: %s" % engine)
+    engine = _resolve_engine(engine)
 
     if engine == "orjson":
         JsonConfig.validate_orjson()
         # orjson handles bytes input natively
-        value_dict = orjson.loads(value)
+        value_dict = get_module("orjson", should_load=True).loads(value)
+    elif engine == "msgspec":
+        JsonConfig.validate_msgspec()
+        value_dict = get_module("msgspec", should_load=True).json.decode(value)
     else:
         # decode bytes to str for built-in json module
         if isinstance(value, bytes):
@@ -400,7 +462,10 @@ def from_json(value, output_type="Figure", skip_invalid=False, engine=None):
           - if "json", parse JSON using built in json module
           - if "orjson", parse using the faster orjson module, requires the orjson
             package
-          - if "auto" use orjson module if available, otherwise use the json module
+          - if "msgspec", parse using the faster msgspec module, requires the
+            msgspec package
+          - if "auto" use orjson module if available, then msgspec, otherwise
+            use the json module
 
         If not specified, the default engine is set to the current value of
         plotly.io.json.config.default_engine.
@@ -456,7 +521,10 @@ def read_json(file, output_type="Figure", skip_invalid=False, engine=None):
           - if "json", parse JSON using built in json module
           - if "orjson", parse using the faster orjson module, requires the orjson
             package
-          - if "auto" use orjson module if available, otherwise use the json module
+          - if "msgspec", parse using the faster msgspec module, requires the
+            msgspec package
+          - if "auto" use orjson module if available, then msgspec, otherwise
+            use the json module
 
         If not specified, the default engine is set to the current value of
         plotly.io.json.config.default_engine.

pyproject.toml

@@ -72,6 +72,7 @@ dev_optional = [
     "fiona<=1.9.6;python_version<='3.8'",
     "geopandas",
     "jupyterlab",
+    "msgspec",
     "orjson",
     "pandas",
     "pdfrw",

tests/test_io/test_to_from_plotly_json.py

@@ -13,6 +13,7 @@
 from _plotly_utils.optional_imports import get_module
 
 orjson = get_module("orjson")
+msgspec = get_module("msgspec")
 
 eastern = timezone("US/Eastern")
 
@@ -67,10 +68,11 @@ def check_roundtrip(value, engine, pretty):
 
 
 # Fixtures
+engines = ["json", "auto"]
 if orjson is not None:
-    engines = ["json", "orjson", "auto"]
-else:
-    engines = ["json", "auto"]
+    engines.append("orjson")
+if msgspec is not None:
+    engines.append("msgspec")
 
 
 @pytest.fixture(scope="module", params=engines)
@@ -322,3 +324,41 @@ def to_plotly_json(self):
     )
     assert result["fig"]["layout"]["title"] == {"text": "from to_dict"}
     assert result["trace"]["name"] == "from to_plotly_json"
+
+
+def test_datetime_series_with_nat(engine):
+    value = {"v": pd.Series([pd.Timestamp("2024-01-01"), pd.NaT])}
+    result = pio.to_json_plotly(value, engine=engine)
+    assert json.loads(result) == {"v": ["2024-01-01T00:00:00", None]}
+
+
+@pytest.mark.parametrize("value", [object(), 1j, {(1, 2): 1}, [{(1,): 2}]])
+def test_unserializable_raises_type_error(engine, value):
+    with pytest.raises(TypeError):
+        pio.to_json_plotly({"v": value}, engine=engine)
+
+
+@pytest.mark.skipif(msgspec is None, reason="requires msgspec")
+def test_auto_engine_uses_msgspec_without_orjson(monkeypatch):
+    def without_orjson(name, **kwargs):
+        return None if name == "orjson" else get_module(name, **kwargs)
+
+    monkeypatch.setattr("plotly.io._json.get_module", without_orjson)
+    utc = datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc)
+    # msgspec writes UTC as Z, the json engine as +00:00
+    assert (
+        pio.to_json_plotly({"t": utc}, engine="auto") == '{"t":"2024-01-01T00:00:00Z"}'
+    )
+    assert pio.from_json_plotly('{"a":[1]}', engine="auto") == {"a": [1]}
+
+
+@pytest.mark.skipif(msgspec is None, reason="requires msgspec")
+def test_msgspec_default_engine():
+    original = pio.config.default_engine
+    try:
+        pio.config.default_engine = "msgspec"
+        assert pio.to_json_plotly(go.Figure(go.Bar(y=[1]))) == pio.to_json_plotly(
+            go.Figure(go.Bar(y=[1])), engine="json"
+        )
+    finally:
+        pio.config.default_engine = original