Match the json engine on Decimal, NaT and circular values with msgspec

Author: T4rk1n

plotly/io/_json.py

@@ -155,6 +155,7 @@ def to_json_plotly(plotly_object, pretty=False, engine=None):
     elif engine == "msgspec":
         JsonConfig.validate_msgspec()
         msgspec = get_module("msgspec", should_load=True)
+        np = modules["np"]
         pd = modules["pd"]
 
         def enc_hook(value):
@@ -171,6 +172,8 @@ def enc_hook(value):
                 raise TypeError(
                     f"Type is not JSON serializable: {type(value).__name__}"
                 )
+            if np is not None and isinstance(value, np.datetime64) and np.isnat(value):
+                return None
             if pd is not None:
                 if value is pd.NaT:
                     return None
@@ -197,7 +200,16 @@ def enc_hook(value):
         encoder = msgspec.json.Encoder(enc_hook=enc_hook, decimal_format="number")
         try:
             out = encoder.encode(plotly_object)
-        except (TypeError, ValueError, OverflowError, msgspec.EncodeError):
+            # msgspec writes Decimal as is, so a NaN or Infinity value or a
+            # Decimal dict key gives invalid JSON. Checking the syntax is cheap.
+            msgspec.json.decode(out, type=msgspec.Raw)
+        except (
+            TypeError,
+            ValueError,
+            OverflowError,
+            RecursionError,
+            msgspec.MsgspecError,
+        ):
             # Errors and values msgspec rejects come out as with the json engine
             return to_json_plotly(plotly_object, pretty=pretty, engine="json")
         if pretty:
@@ -430,7 +442,12 @@ def from_json_plotly(value, engine=None):
         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)
+        msgspec = get_module("msgspec", should_load=True)
+        try:
+            value_dict = msgspec.json.decode(value)
+        except (ValueError, msgspec.MsgspecError):
+            # NaN and Infinity literals, numbers out of float range
+            value_dict = from_json_plotly(value, engine="json")
     else:
         # decode bytes to str for built-in json module
         if isinstance(value, bytes):

tests/test_io/test_to_from_plotly_json.py

@@ -6,6 +6,7 @@
 import pandas as pd
 import json
 import datetime
+import decimal
 import re
 import sys
 import warnings
@@ -362,3 +363,38 @@ def test_msgspec_default_engine():
         )
     finally:
         pio.config.default_engine = original
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        decimal.Decimal("NaN"),
+        decimal.Decimal("-Infinity"),
+    ],
+)
+def test_non_finite_values_are_null(engine, pretty, value):
+    result = pio.to_json_plotly({"v": [value]}, engine=engine, pretty=pretty)
+    assert json.loads(result) == {"v": [None]}
+
+
+@pytest.mark.skipif(msgspec is None, reason="requires msgspec")
+def test_msgspec_datetime64_nat_is_null():
+    value = {"v": np.datetime64("NaT")}
+    assert pio.to_json_plotly(value, engine="msgspec") == '{"v":null}'
+
+
+@pytest.mark.skipif(msgspec is None, reason="requires msgspec")
+def test_msgspec_errors_match_json():
+    circular = []
+    circular.append(circular)
+    with pytest.raises(ValueError, match="Circular reference"):
+        pio.to_json_plotly(circular, engine="msgspec")
+    with pytest.raises(TypeError):
+        pio.to_json_plotly({decimal.Decimal("1.5"): 1}, engine="msgspec")
+
+
+@pytest.mark.skipif(msgspec is None, reason="requires msgspec")
+def test_msgspec_reads_what_json_reads():
+    value = '{"a":NaN,"b":1e400}'
+    result = pio.from_json_plotly(value, engine="msgspec")
+    assert str(result) == str(pio.from_json_plotly(value, engine="json"))