@marimo-team/frontend 0.23.3-dev11 → 0.23.3-dev15

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

@@ -8,6 +8,7 @@ import {
   floatType,
   intType,
   JsonViewer,
+  type JsonViewerKeyRenderer,
   nullType,
   objectType,
   stringType,
@@ -130,6 +131,9 @@ export const JsonOutput: React.FC<Props> = memo(
             collapseStringsAfterLength={COLLAPSED_TEXT_LENGTH}
             // leave the default valueTypes as it was - 'python', only 'json' is changed
             valueTypes={valueTypesMap[valueTypes]}
+            // Render dict keys that carry Python type info (e.g. `int`, `tuple`).
+            // See `_key_formatter` in marimo/_output/formatters/structures.py.
+            keyRenderer={valueTypes === "python" ? keyRenderer : undefined}
             // Don't group arrays, it will make the tree view look like there are nested arrays
             groupArraysAfterLength={Number.MAX_SAFE_INTEGER}
             // Built-in clipboard shifts content on hover
@@ -207,7 +211,10 @@ const LEAF_RENDERERS: Record<string, LeafRenderer> = {
   ),
   "text/plain+float:": (value) => <span>{value}</span>,
   "text/plain+bigint:": (value) => <span>{value}</span>,
-  "text/plain+set:": (value) => <span>set{value}</span>,
+  "text/plain+set:": (value) => <span>{formatSetPayload(value)}</span>,
+  "text/plain+frozenset:": (value) => (
+    <span>{formatFrozensetPayload(value)}</span>
+  ),
   "text/plain+tuple:": (value) => <span>{value}</span>,
   "text/plain:": (value) => <CollapsibleTextOutput text={value} />,
   "application/json:": (value) => (
@@ -375,6 +382,99 @@ function renderLeaf(leaf: string, render: LeafRenderer): React.ReactNode {
   return <span>{leaf}</span>;
 }
 
+// Prefix marking keys that carry encoded type information from Python.
+// See `_key_formatter` in marimo/_output/formatters/structures.py.
+const KEY_ENCODED_PREFIX = "text/plain+";
+
+// 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.
+function formatTuplePayload(jsonList: string): string {
+  const items = JSON.parse(jsonList) as unknown[];
+  const inner = items.map((x) => JSON.stringify(x)).join(", ");
+  if (items.length === 0) {
+    return "()";
+  }
+  if (items.length === 1) {
+    return `(${inner},)`;
+  }
+  return `(${inner})`;
+}
+
+// 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[];
+  if (items.length === 0) {
+    return "frozenset()";
+  }
+  const inner = items.map((x) => JSON.stringify(x)).join(", ");
+  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 {
+    // 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;
+  }
+}
+
+// Renderers for decoded non-string keys. Visual affordances match Python:
+// unquoted primitives, parens for tuple, `frozenset({...})` for frozenset,
+// and the `text/plain+str:` escape re-quotes the original string.
+const KEY_DECODERS: Record<string, (data: string) => React.ReactNode> = {
+  "text/plain+int:": (v) => <span>{v}</span>,
+  "text/plain+float:": (v) => <span>{v}</span>,
+  "text/plain+bool:": (v) => <span>{v === "True" ? "True" : "False"}</span>,
+  "text/plain+none:": () => <span>None</span>,
+  "text/plain+tuple:": (v) => <span>{formatTuplePayload(v)}</span>,
+  "text/plain+frozenset:": (v) => <span>{formatFrozensetPayload(v)}</span>,
+  "text/plain+str:": (v) => <span>"{v}"</span>,
+};
+
+function isEncodedKey(key: unknown): key is string {
+  return typeof key === "string" && key.startsWith(KEY_ENCODED_PREFIX);
+}
+
+// `@textea/json-viewer` drops quotes from integer-like string keys, which
+// makes the string `"2"` visually identical to the decoded int `2`. Match
+// the same keys the viewer strips and render them with explicit quotes.
+const INT_LIKE_STRING = /^-?\d+$/;
+
+const keyRenderer: JsonViewerKeyRenderer = Object.assign(
+  ({ path }: DataItemProps) => {
+    const key = path[path.length - 1];
+    if (typeof key !== "string") {
+      return <span>{String(key)}</span>;
+    }
+    if (isEncodedKey(key)) {
+      const [data, mimeType] = leafDataAndMimeType(key);
+      const render = KEY_DECODERS[`${mimeType}:`];
+      return render ? render(data) : <span>{key}</span>;
+    }
+    // Plain integer-like string — quote it so it's distinct from a decoded int.
+    return <span>"{key}"</span>;
+  },
+  {
+    when: ({ path }: DataItemProps) => {
+      const key = path[path.length - 1];
+      return (
+        isEncodedKey(key) ||
+        (typeof key === "string" && INT_LIKE_STRING.test(key))
+      );
+    },
+  },
+);
+
 const MIME_PREFIXES = Object.keys(LEAF_RENDERERS);
 const REPLACE_PREFIX = "<marimo-replace>";
 const REPLACE_SUFFIX = "</marimo-replace>";
@@ -413,8 +513,10 @@ function pythonJsonReplacer(_key: string, value: unknown): unknown {
       return `${REPLACE_PREFIX}(${leafData(value).slice(1, -1)})${REPLACE_SUFFIX}`;
     }
     if (value.startsWith("text/plain+set:")) {
-      // replace first and last characters [] with {}
-      return `${REPLACE_PREFIX}{${leafData(value).slice(1, -1)}}${REPLACE_SUFFIX}`;
+      return `${REPLACE_PREFIX}${formatSetPayload(leafData(value))}${REPLACE_SUFFIX}`;
+    }
+    if (value.startsWith("text/plain+frozenset:")) {
+      return `${REPLACE_PREFIX}${formatFrozensetPayload(leafData(value))}${REPLACE_SUFFIX}`;
     }
 
     if (MIME_PREFIXES.some((prefix) => value.startsWith(prefix))) {
@@ -428,10 +530,61 @@ function pythonJsonReplacer(_key: string, value: unknown): unknown {
   return value;
 }
 
+// Rewrite an encoded key string into the Python literal that should appear
+// unquoted in the copy output. Wrapping in REPLACE_PREFIX/SUFFIX makes the
+// final regex pass strip the surrounding JSON quotes.
+function decodeKeyForCopy(key: string): string {
+  const [data, mimeType] = leafDataAndMimeType(key);
+  const wrap = (s: string) => `${REPLACE_PREFIX}${s}${REPLACE_SUFFIX}`;
+  switch (`${mimeType}:`) {
+    case "text/plain+int:":
+      return wrap(data);
+    case "text/plain+float:":
+      if (data === "nan") {
+        return wrap("float('nan')");
+      }
+      if (data === "inf") {
+        return wrap("float('inf')");
+      }
+      if (data === "-inf") {
+        return wrap("-float('inf')");
+      }
+      return wrap(data);
+    case "text/plain+bool:":
+      return wrap(data === "True" ? "True" : "False");
+    case "text/plain+none:":
+      return wrap("None");
+    case "text/plain+tuple:":
+      return wrap(formatTuplePayload(data));
+    case "text/plain+frozenset:":
+      return wrap(formatFrozensetPayload(data));
+    case "text/plain+str:":
+      // `data` is the original Python string; it stays quoted.
+      return data;
+    default:
+      return key;
+  }
+}
+
+function rewriteEncodedKeys(value: unknown): unknown {
+  if (Array.isArray(value)) {
+    return value.map(rewriteEncodedKeys);
+  }
+  if (typeof value === "object" && value !== null) {
+    const out: Record<string, unknown> = {};
+    for (const [k, v] of Object.entries(value)) {
+      const newKey = isEncodedKey(k) ? decodeKeyForCopy(k) : k;
+      out[newKey] = rewriteEncodedKeys(v);
+    }
+    return out;
+  }
+  return value;
+}
+
 export function getCopyValue(value: unknown): string {
   // Because this results in valid json, it adds quotes around None and True/False.
   // but we want to make this look like Python, so we remove the quotes.
-  return JSON.stringify(value, pythonJsonReplacer, 2)
+  return JSON.stringify(rewriteEncodedKeys(value), pythonJsonReplacer, 2)
     .replaceAll(`"${REPLACE_PREFIX}`, "")
     .replaceAll(`${REPLACE_SUFFIX}"`, "");
 }

package/src/components/editor/output/__tests__/JsonOutput-mimetype.test.tsx

@@ -39,4 +39,84 @@ describe("JsonOutput with enhanced mimetype handling", () => {
     expect(container).toBeInTheDocument();
     expect(container.querySelector(".marimo-json-output")).toBeInTheDocument();
   });
+
+  it("renders encoded non-string keys with Python-style affordances", () => {
+    // Server-side `_key_formatter` encodes non-string dict keys with
+    // mimetype prefixes; the frontend `keyRenderer` must decode them
+    // so users see the original Python types (unquoted ints, parens for
+    // tuples, etc.) instead of the raw encoded strings.
+    const data = {
+      "text/plain+int:2": "int_val",
+      "text/plain+float:2.5": "float_val",
+      "text/plain+bool:True": "bool_val",
+      "text/plain+none:": "none_val",
+      "text/plain+tuple:[1, 2]": "tuple_val",
+      "text/plain+frozenset:[3, 4]": "fs_val",
+      "text/plain+str:text/plain+int:2": "escaped_str_val",
+      plain: "unchanged",
+    };
+
+    const { container } = render(<JsonOutput data={data} format="tree" />);
+    const text = container.textContent ?? "";
+
+    // `text/plain+str:` is the escape prefix — must never survive in output.
+    expect(text).not.toContain("text/plain+str:");
+    // Other encoded prefixes must not leak as-is. (They can still appear
+    // inside the unescaped original string key `"text/plain+int:2"`,
+    // which is intentional — but not for types *other* than int.)
+    expect(text).not.toContain("text/plain+bool:True");
+    expect(text).not.toContain("text/plain+tuple:[");
+    expect(text).not.toContain("text/plain+frozenset:[");
+    expect(text).not.toContain("text/plain+none:");
+
+    // Decoded visual forms are present with Python-style affordances.
+    expect(text).toContain('None:"none_val"');
+    expect(text).toContain('True:"bool_val"');
+    expect(text).toContain('2:"int_val"');
+    expect(text).toContain('2.5:"float_val"');
+    expect(text).toContain('(1, 2):"tuple_val"');
+    expect(text).toContain('frozenset({3, 4}):"fs_val"');
+    // Escaped str key renders as the original literal string (quoted).
+    expect(text).toContain('"text/plain+int:2":"escaped_str_val"');
+    // Plain string key unchanged.
+    expect(text).toContain('"plain":"unchanged"');
+  });
+
+  it("renders 1-tuple and empty-frozenset keys with correct Python syntax", () => {
+    // Regressions caught in review: `(1)` is not a tuple (needs `(1,)`),
+    // and `frozenset({})` reads like constructing from an empty dict
+    // (should be `frozenset()`). Locks in the tree-view rendering so these
+    // don't slip back.
+    const data = {
+      "text/plain+tuple:[42]": "one_tuple",
+      "text/plain+tuple:[]": "empty_tuple",
+      "text/plain+frozenset:[]": "empty_fs",
+      "text/plain+frozenset:[1]": "one_fs",
+    };
+
+    const { container } = render(<JsonOutput data={data} format="tree" />);
+    const text = container.textContent ?? "";
+
+    expect(text).toContain('(42,):"one_tuple"'); // trailing comma
+    expect(text).toContain('():"empty_tuple"');
+    expect(text).toContain('frozenset():"empty_fs"'); // not `frozenset({})`
+    expect(text).toContain('frozenset({1}):"one_fs"');
+  });
+
+  it("quotes integer-like string keys to distinguish them from int keys", () => {
+    // Without this, `"2"` and the decoded int `2` look identical — the
+    // textea viewer drops quotes from integer-like string keys by default.
+    const data = {
+      "2": "string_two",
+      "text/plain+int:2": "int_two",
+    };
+
+    const { container } = render(<JsonOutput data={data} format="tree" />);
+    const text = container.textContent ?? "";
+
+    expect(text).toContain('"2":"string_two"'); // quoted
+    expect(text).toContain('2:"int_two"'); // unquoted
+    // Non-integer string keys still render without our intervention.
+    expect(text).not.toContain("text/plain+"); // prefix stripped from int key
+  });
 });

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

@@ -174,7 +174,21 @@ describe("getCopyValue", () => {
   it("should handle sets", () => {
     const value = "text/plain+set:[1,2,3]";
     const result = getCopyValue(value);
-    expect(result).toMatchInlineSnapshot(`"{1,2,3}"`);
+    expect(result).toMatchInlineSnapshot(`"{1, 2, 3}"`);
+  });
+
+  it("should handle empty set", () => {
+    // Empty set literal in Python is `set()`, not `{}` (which is a dict).
+    expect(getCopyValue("text/plain+set:[]")).toMatchInlineSnapshot(`"set()"`);
+  });
+
+  it("should handle frozenset values", () => {
+    expect(getCopyValue("text/plain+frozenset:[1,2]")).toMatchInlineSnapshot(
+      `"frozenset({1, 2})"`,
+    );
+    expect(getCopyValue("text/plain+frozenset:[]")).toMatchInlineSnapshot(
+      `"frozenset()"`,
+    );
   });
 
   it("should handle sets in mixed types", () => {
@@ -188,7 +202,7 @@ describe("getCopyValue", () => {
       `
       "{
         "key1": 42,
-        "key2": {1,2,3},
+        "key2": {1, 2, 3},
         "key3": True
       }"
     `,
@@ -311,6 +325,137 @@ describe("determineMaxDisplayLength", () => {
   });
 });
 
+describe("getCopyValue with encoded non-string keys", () => {
+  // Keys are encoded by _key_formatter in
+  // marimo/_output/formatters/structures.py. Frontend must round-trip them
+  // to Python literals in the copy output.
+
+  it("decodes int keys unquoted", () => {
+    // JS reorders integer-like string keys to the front of object iteration
+    // (spec-mandated), so `"2"` appears before `"text/plain+int:2"` here.
+    // This is pre-existing and unrelated to the encoding — both entries
+    // survive, which is the regression this guards.
+    const value = { "text/plain+int:2": "no", "2": "oh" };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        "2": "oh",
+        2: "no"
+      }"
+    `);
+  });
+
+  it("decodes large int keys unquoted (no BigInt precision concern)", () => {
+    const value = { "text/plain+int:18446744073709551616": "v" };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        18446744073709551616: "v"
+      }"
+    `);
+  });
+
+  it("decodes float, bool, None, tuple, frozenset keys", () => {
+    const value = {
+      "text/plain+float:2.5": "f",
+      "text/plain+bool:True": "t",
+      "text/plain+bool:False": "b",
+      "text/plain+none:": "n",
+      "text/plain+tuple:[1, 2]": "tup",
+      "text/plain+frozenset:[3, 4]": "fs",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        2.5: "f",
+        True: "t",
+        False: "b",
+        None: "n",
+        (1, 2): "tup",
+        frozenset({3, 4}): "fs"
+      }"
+    `);
+  });
+
+  it("emits 1-element tuple keys with a trailing comma (Python syntax)", () => {
+    // `(1)` is just `1` in Python — a 1-tuple needs `(1,)`.
+    const value = {
+      "text/plain+tuple:[1]": "one",
+      "text/plain+tuple:[]": "empty",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        (1,): "one",
+        (): "empty"
+      }"
+    `);
+  });
+
+  it("emits empty frozenset keys as `frozenset()` not `frozenset({})`", () => {
+    // `frozenset({})` reads like it's constructing from an empty dict.
+    const value = {
+      "text/plain+frozenset:[]": "empty",
+      "text/plain+frozenset:[1]": "single",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        frozenset(): "empty",
+        frozenset({1}): "single"
+      }"
+    `);
+  });
+
+  it("decodes NaN/Inf float keys to valid Python literals", () => {
+    const value = {
+      "text/plain+float:nan": "n",
+      "text/plain+float:inf": "p",
+      "text/plain+float:-inf": "m",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        float('nan'): "n",
+        float('inf'): "p",
+        -float('inf'): "m"
+      }"
+    `);
+  });
+
+  it("unescapes string keys that looked encoded", () => {
+    const value = {
+      "text/plain+str:text/plain+int:2": "hello",
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        "text/plain+int:2": "hello"
+      }"
+    `);
+  });
+
+  it("decodes keys at every nesting level", () => {
+    const value = {
+      outer: {
+        "text/plain+int:1": "inner",
+        "text/plain+tuple:[2, 3]": "tup",
+      },
+    };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        "outer": {
+          1: "inner",
+          (2, 3): "tup"
+        }
+      }"
+    `);
+  });
+
+  it("leaves plain string keys untouched", () => {
+    const value = { foo: 1, bar: 2 };
+    expect(getCopyValue(value)).toMatchInlineSnapshot(`
+      "{
+        "foo": 1,
+        "bar": 2
+      }"
+    `);
+  });
+});
+
 describe("getCopyValue with application/ mimetypes", () => {
   it("should strip application/ mimetype prefix from leaf data", () => {
     expect(getCopyValue("application/json:{data}")).toBe('"{data}"');