fix: decode tuple/frozenset payloads with non-finite floats (#9365)

Python's `json.dumps` emits bare `NaN`/`Infinity`/`-Infinity` for
non-finite floats inside embedded tuple/frozenset payloads (e.g.
`text/plain+tuple:[NaN]`). The outer JSON still parses because those
tokens live inside a JSON string, but `formatTuplePayload` and
`formatFrozensetPayload` used strict `JSON.parse` and threw.

Swap both to `jsonParseWithSpecialChar` (same parser we already use at
the outer level) and render non-finite elements as `float('nan')` /
`float('inf')` / `-float('inf')`, matching the scalar-key form.

Author: mscolnick

frontend/src/components/editor/output/JsonOutput.tsx

@@ -18,6 +18,7 @@ import { memo, useState } from "react";
 import type { OutputMessage } from "@/core/kernel/messages";
 import { cn } from "@/utils/cn";
 import { copyToClipboard } from "@/utils/copy";
+import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { isUrl } from "@/utils/urls";
 import { useTheme } from "../../../theme/useTheme";
 import { logNever } from "../../../utils/assertNever";
@@ -386,14 +387,41 @@ function renderLeaf(leaf: string, render: LeafRenderer): React.ReactNode {
 // See `_key_formatter` in marimo/_output/formatters/structures.py.
 const KEY_ENCODED_PREFIX = "text/plain+";
 
+// Format elements for a Python collection literal. Non-finite floats
+// (NaN / Infinity / -Infinity) parse as JS `number` via
+// `jsonParseWithSpecialChar`; `JSON.stringify` on those returns `null`,
+// so render them as the same `float(...)` literals we use for scalar
+// float keys (see `decodeKeyForCopy`).
+function formatCollectionItems(items: unknown[]): string {
+  return items
+    .map((x) => {
+      if (typeof x === "number" && !Number.isFinite(x)) {
+        if (Number.isNaN(x)) {
+          return "float('nan')";
+        }
+        return x > 0 ? "float('inf')" : "-float('inf')";
+      }
+      return JSON.stringify(x);
+    })
+    .join(", ");
+}
+
 // Format a JSON-list payload as a Python tuple literal. 1-element tuples
 // need a trailing comma — `(1)` is just `1` in Python, `(1,)` is the tuple.
+// Uses `jsonParseWithSpecialChar` so bare `NaN`/`Infinity`/`-Infinity`
+// emitted by Python's json.dumps round-trip cleanly.
 function formatTuplePayload(jsonList: string): string {
-  const items = JSON.parse(jsonList) as unknown[];
-  const inner = items.map((x) => JSON.stringify(x)).join(", ");
+  const items = jsonParseWithSpecialChar<unknown[]>(jsonList);
+  // `jsonParseWithSpecialChar` returns `{}` when both parse passes fail;
+  // fall back to the raw payload so a malformed wire form doesn't crash
+  // rendering/copy. Matches the defensive pattern in `formatSetPayload`.
+  if (!Array.isArray(items)) {
+    return jsonList;
+  }
   if (items.length === 0) {
     return "()";
   }
+  const inner = formatCollectionItems(items);
   if (items.length === 1) {
     return `(${inner},)`;
   }
@@ -403,29 +431,31 @@ function formatTuplePayload(jsonList: string): string {
 // Format a JSON-list payload as a Python frozenset literal. Empty → `frozenset()`
 // rather than `frozenset({})` (which reads like a dict).
 function formatFrozensetPayload(jsonList: string): string {
-  const items = JSON.parse(jsonList) as unknown[];
+  const items = jsonParseWithSpecialChar<unknown[]>(jsonList);
+  if (!Array.isArray(items)) {
+    return jsonList;
+  }
   if (items.length === 0) {
     return "frozenset()";
   }
-  const inner = items.map((x) => JSON.stringify(x)).join(", ");
+  const inner = formatCollectionItems(items);
   return `frozenset({${inner}})`;
 }
 
 // Format a JSON-list payload as a Python set literal. Empty → `set()`
 // (not `{}`, which is a dict literal in Python).
 function formatSetPayload(jsonList: string): string {
-  try {
-    const items = JSON.parse(jsonList) as unknown[];
-    if (items.length === 0) {
-      return "set()";
-    }
-    const inner = items.map((x) => JSON.stringify(x)).join(", ");
-    return `{${inner}}`;
-  } catch {
+  const items = jsonParseWithSpecialChar<unknown[]>(jsonList);
+  if (!Array.isArray(items)) {
     // Back-compat: older wire form was `text/plain+set:{1, 2, 3}` (Python
     // set-literal string, not JSON). Pass it through as-is rather than crash.
     return jsonList;
   }
+  if (items.length === 0) {
+    return "set()";
+  }
+  const inner = formatCollectionItems(items);
+  return `{${inner}}`;
 }
 
 // Renderers for decoded non-string keys. Visual affordances match Python:

frontend/src/components/editor/output/__tests__/json-output.test.ts

@@ -417,6 +417,44 @@ describe("getCopyValue with encoded non-string keys", () => {
     `);
   });
 
+  it("parses tuple/frozenset payloads containing bare NaN/Infinity", () => {
+    // Python's json.dumps emits bare `NaN`/`Infinity` inside the embedded
+    // tuple/frozenset payload strings (JSON spec violation, but ECMA-262-
+    // friendly via the fallback in jsonParseWithSpecialChar). The outer
+    // JSON stays strict because those tokens live inside a JSON string
+    // key/value. Regression for tuple-key payloads that previously broke
+    // the frontend's `JSON.parse` and threw.
+    const value = {
+      "text/plain+tuple:[NaN]": "tn",
+      "text/plain+tuple:[Infinity, -Infinity]": "ti",
+      k: "text/plain+frozenset:[Infinity, 1]",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        (float('nan'),): "tn",
+        (float('inf'), -float('inf')): "ti",
+        "k": frozenset({float('inf'), 1})
+      }"
+    `);
+  });
+
+  it("falls back to the raw payload for malformed tuple/frozenset", () => {
+    // `jsonParseWithSpecialChar` returns `{}` on parse failure rather
+    // than throwing; without an `Array.isArray` guard, the formatters
+    // would crash on `.length`/`.map`. Pass the raw payload through so
+    // a malformed wire form doesn't break the whole render.
+    const value = {
+      "text/plain+tuple:not a json list": "t",
+      k: "text/plain+frozenset:also broken",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        not a json list: "t",
+        "k": also broken
+      }"
+    `);
+  });
+
   it("unescapes string keys that looked encoded", () => {
     const value = {
       "text/plain+str:text/plain+int:2": "hello",

tests/_output/formatters/test_structures.py

@@ -671,6 +671,85 @@ def test_format_structure_dict_plain_string_keys_unchanged() -> None:
     )
 
 
+def _reject_non_finite(literal: str) -> float:
+    # `parse_constant` fires for bare `NaN`, `Infinity`, and `-Infinity`.
+    # Python's `json.loads` accepts these by default (non-spec) — passing a
+    # raising hook makes the test match the JS `JSON.parse` behavior we
+    # actually care about.
+    raise AssertionError(f"outer JSON contained bare {literal!r}")
+
+
+def test_format_structure_tuple_key_with_nan_outer_json_is_strict() -> None:
+    """Tuple keys with non-finite floats are embedded strings — the outer
+    JSON must be strict per the JS `JSON.parse` spec (no bare `NaN` /
+    `Infinity` at the top level).
+
+    Those tokens live inside the embedded tuple payload *string*, not at
+    the outer JSON level, so the frontend's outer `JSON.parse` succeeds
+    and it then calls `jsonParseWithSpecialChar` on the embedded
+    payload. Previously the bare token appeared at the top level and
+    broke the outer parse.
+    """
+    StructuresFormatter().register()
+
+    _, data = get_and_format(
+        {(float("nan"),): "n", (float("inf"), -float("inf")): "i"}
+    )
+    # `parse_constant` raises on bare NaN/Infinity at the JSON layer —
+    # matching JS `JSON.parse` strictness. `json.loads` alone accepts
+    # them by default, which is too lenient to test the contract.
+    parsed = json.loads(data, parse_constant=_reject_non_finite)
+    assert parsed == {
+        "text/plain+tuple:[NaN]": "n",
+        "text/plain+tuple:[Infinity, -Infinity]": "i",
+    }
+
+
+def test_format_structure_frozenset_value_with_nan_outer_json_is_strict() -> (
+    None
+):
+    """Frozenset values with non-finite floats parse strictly at the outer level."""
+    StructuresFormatter().register()
+
+    _, data = get_and_format({"k": frozenset({float("inf"), 1})})
+    # Outer parse is strict (JS-`JSON.parse`-compatible).
+    parsed = json.loads(data, parse_constant=_reject_non_finite)
+    key = parsed["k"]
+    assert key.startswith("text/plain+frozenset:")
+    # The embedded payload contains bare `Infinity`; the frontend parses
+    # it with `jsonParseWithSpecialChar`. Python's permissive `json.loads`
+    # is fine here because we're just inspecting the embedded content.
+    payload = json.loads(key[len("text/plain+frozenset:") :])
+    assert set(payload) == {1, float("inf")}
+
+
+def test_format_structure_frozenset_key_with_nan_outer_json_is_strict() -> (
+    None
+):
+    """Frozenset keys with non-finite floats parse strictly at the outer level."""
+    StructuresFormatter().register()
+
+    _, data = get_and_format({frozenset({float("nan")}): "v"})
+    # Outer parse is strict — the bare `NaN` lives inside the key string.
+    parsed = json.loads(data, parse_constant=_reject_non_finite)
+    (key,) = parsed
+    assert key == "text/plain+frozenset:[NaN]"
+
+
+def test_format_structure_tuple_value_with_nan_is_strict_json() -> None:
+    """Tuple values with non-finite floats round-trip via scalar sentinels.
+
+    Tuple values don't hit the tuple-encoder path because `flatten`
+    recurses into tuples before leaf formatting — each float is handled
+    by `_leaf_formatter` and emitted as its own `text/plain+float:`
+    sentinel string.
+    """
+    StructuresFormatter().register()
+
+    formatted = format_structure([(float("nan"), float("inf"))])
+    assert formatted == [("text/plain+float:nan", "text/plain+float:inf")]
+
+
 def test_format_structure_bigint() -> None:
     bigint = 2**64
     assert format_structure([bigint]) == ([f"text/plain+bigint:{bigint}"])