fixes

frontend/src/core/errors/__tests__/errors.test.ts

@@ -1,6 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import { describe, it, expect } from "vitest";
-import { getImportCode } from "../errors";
+import { getAutoFixes, getImportCode } from "../errors";
+import type { MarimoError } from "@/core/kernel/messages";
 
 describe("getImportCode", () => {
   it("returns simple import for same name", () => {
@@ -14,3 +15,48 @@ describe("getImportCode", () => {
     expect(getImportCode("plt")).toBe("import matplotlib.pyplot as plt");
   });
 });
+
+describe("getAutoFixes", () => {
+  it("returns wrap in function fix for multiple-defs error", () => {
+    const error: MarimoError = {
+      type: "multiple-defs",
+      name: "foo",
+      cells: ["foo"],
+    };
+
+    const fixes = getAutoFixes(error);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("Wrap in a function");
+  });
+
+  it("returns import fix for NameError with known import", () => {
+    const error: MarimoError = {
+      type: "exception",
+      exception_type: "NameError",
+      msg: "name 'np' is not defined",
+    };
+
+    const fixes = getAutoFixes(error);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("Add 'import numpy as np'");
+  });
+
+  it("returns no fixes for NameError with unknown import", () => {
+    const error: MarimoError = {
+      type: "exception",
+      exception_type: "NameError",
+      msg: "name 'unknown_module' is not defined",
+    };
+
+    expect(getAutoFixes(error)).toHaveLength(0);
+  });
+
+  it("returns no fixes for other error types", () => {
+    const error: MarimoError = {
+      type: "syntax",
+      msg: "invalid syntax",
+    };
+
+    expect(getAutoFixes(error)).toHaveLength(0);
+  });
+});

frontend/src/core/errors/__tests__/utils.test.ts

@@ -35,6 +35,37 @@ _()`;
     expect(wrapInFunction(input)).toBe(expected);
   });
 
+  test("preserves existing indentation", () => {
+    const input = `def foo():
+    x = 1
+    y = 2`;
+    const expected = `def _():
+    def foo():
+        x = 1
+        y = 2
+    return
+_()`;
+    expect(wrapInFunction(input)).toBe(expected);
+  });
+
+  test("multi-line parentheses", () => {
+    const input = `x = "foo"
+y = "bar"
+(alt.Chart(df).mark_line().encode(
+    x="x",
+    y="y",
+))`;
+    const expected = `def _():
+    x = "foo"
+    y = "bar"
+    return (alt.Chart(df).mark_line().encode(
+        x="x",
+        y="y",
+    ))
+_()`;
+    expect(wrapInFunction(input)).toBe(expected);
+  });
+
   test("preserves empty lines", () => {
     const input = `x = 1
 

frontend/src/core/errors/errors.ts

@@ -21,7 +21,7 @@ export function getAutoFixes(error: MarimoError): AutoFix[] {
       {
         title: "Wrap in a function",
         description:
-          "Wrap the cell contents in a function so they are marked as definitions of the cell.",
+          "Make this cell's variables local by wrapping the cell in a function.",
         onFix: async (ctx) => {
           invariant(ctx.editor, "Editor is null");
           const code = wrapInFunction(ctx.editor.state.doc.toString());