feat: add auto-fix for multiple definctions and popular imports

frontend/package.json

@@ -37,6 +37,7 @@
     "@lezer/common": "^1.2.1",
     "@lezer/highlight": "^1.2.1",
     "@lezer/lr": "^1.4.2",
+    "@lezer/python": "^1.1.15",
     "@marimo-team/marimo-api": "file:../openapi",
     "@marimo-team/react-slotz": "^0.1.8",
     "@open-rpc/client-js": "^1.8.1",

frontend/src/components/editor/Output.tsx

@@ -48,9 +48,10 @@ type MimeBundleOrTuple = MimeBundle | [MimeBundle, { [key: string]: unknown }];
  */
 export const OutputRenderer: React.FC<{
   message: Pick<OutputMessage, "channel" | "data" | "mimetype">;
+  cellId?: CellId;
   onRefactorWithAI?: (opts: { prompt: string }) => void;
 }> = memo((props) => {
-  const { message, onRefactorWithAI } = props;
+  const { message, onRefactorWithAI, cellId } = props;
   const { theme } = useTheme();
 
   // Memoize parsing the json data
@@ -123,7 +124,7 @@ export const OutputRenderer: React.FC<{
 
     case "application/vnd.marimo+error":
       invariant(Array.isArray(data), "Expected array data");
-      return <MarimoErrorOutput errors={data} />;
+      return <MarimoErrorOutput cellId={cellId} errors={data} />;
 
     case "application/vnd.marimo+traceback":
       invariant(
@@ -190,7 +191,8 @@ OutputRenderer.displayName = "OutputRenderer";
 const MimeBundleOutputRenderer: React.FC<{
   channel: OutputMessage["channel"];
   data: MimeBundleOrTuple;
-}> = memo(({ data, channel }) => {
+  cellId?: CellId;
+}> = memo(({ data, channel, cellId }) => {
   const mimebundle = Array.isArray(data) ? data[0] : data;
 
   // If there is none, return null
@@ -203,6 +205,7 @@ const MimeBundleOutputRenderer: React.FC<{
   if (Object.keys(mimebundle).length === 1) {
     return (
       <OutputRenderer
+        cellId={cellId}
         message={{
           channel: channel,
           data: mimebundle[first],
@@ -242,6 +245,7 @@ const MimeBundleOutputRenderer: React.FC<{
             <TabsContent key={mime} value={mime}>
               <ErrorBoundary>
                 <OutputRenderer
+                  cellId={cellId}
                   message={{
                     channel: channel,
                     data: output,
@@ -307,7 +311,7 @@ export const OutputArea = React.memo(
           id={CellOutputId.create(cellId)}
           className={cn(stale && "marimo-output-stale", className)}
         >
-          <OutputRenderer message={output} />
+          <OutputRenderer cellId={cellId} message={output} />
         </Container>
       </ErrorBoundary>
     );

frontend/src/components/editor/chrome/panels/error-panel.tsx

@@ -21,7 +21,11 @@ export const ErrorsPanel: React.FC = () => {
             <CellLinkError cellId={error.cellId} />
           </div>
           <div key={error.cellId} className="px-2">
-            <MarimoErrorOutput key={error.cellId} errors={error.output.data} />
+            <MarimoErrorOutput
+              key={error.cellId}
+              errors={error.output.data}
+              cellId={error.cellId}
+            />
           </div>
         </div>
       ))}

frontend/src/components/editor/errors/auto-fix.tsx

@@ -0,0 +1,55 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { Button } from "@/components/ui/button";
+import { Tooltip } from "@/components/ui/tooltip";
+import { useCellActions, notebookAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { getAutoFixes } from "@/core/errors/errors";
+import type { MarimoError } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
+import { LightbulbIcon } from "lucide-react";
+
+export const AutoFixButton = ({
+  errors,
+  cellId,
+}: { errors: MarimoError[]; cellId: CellId }) => {
+  const { createNewCell } = useCellActions();
+  const autoFixes = errors.flatMap((error) => getAutoFixes(error));
+
+  if (autoFixes.length === 0) {
+    return null;
+  }
+
+  // TODO: Add a dropdown menu with the auto-fixes, when we need to support
+  // multiple fixes.
+  const firstFix = autoFixes[0];
+
+  return (
+    <Tooltip content={firstFix.description} align="start">
+      <Button
+        size="xs"
+        variant="secondary"
+        onClick={() => {
+          const editorView =
+            store.get(notebookAtom).cellHandles[cellId].current?.editorView;
+          firstFix.onFix({
+            addCodeBelow: (code: string) => {
+              createNewCell({
+                cellId: cellId,
+                autoFocus: false,
+                before: false,
+                code: code,
+              });
+            },
+            editor: editorView,
+            cellId: cellId,
+          });
+          // Focus the editor
+          editorView?.focus();
+        }}
+      >
+        <LightbulbIcon className="h-3 w-3 mr-2" />
+        {firstFix.title}
+      </Button>
+    </Tooltip>
+  );
+};

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

@@ -138,6 +138,7 @@ export const ConsoleOutput = (props: Props): React.ReactNode => {
         return (
           <React.Fragment key={idx}>
             <OutputRenderer
+              cellId={cellId}
               onRefactorWithAI={onRefactorWithAI}
               message={output}
             />

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

@@ -15,6 +15,7 @@
 import { Fragment } from "react";
 import { CellLinkError } from "../links/cell-link";
 import type { CellId } from "@/core/cells/ids";
+import { AutoFixButton } from "../errors/auto-fix";
 
 const Tip = (props: {
   className?: string;
@@ -31,6 +32,7 @@
 };
 
 interface Props {
+  cellId: CellId | undefined;
   errors: MarimoError[];
   className?: string;
 }
@@ -40,6 +42,7 @@
  */
 export const MarimoErrorOutput = ({
   errors,
+  cellId,
   className,
 }: Props): JSX.Element => {
   let titleContents = "This cell wasn't run because it has errors";
@@ -55,7 +58,7 @@
       case "cycle":
         return (
           <Fragment key={idx}>
-            <p className="mt-4">{"This cell is in a cycle:"}</p>
+            <p>{"This cell is in a cycle:"}</p>
             <ul className="list-disc">
               {error.edges_with_vars.map((edge) => (
                 <li className={liStyle} key={`${edge[0]}-${edge[1]}`}>
@@ -76,9 +79,7 @@
       case "multiple-defs":
         return (
           <Fragment key={idx}>
-            <p className="mt-4">
-              {`The variable '${error.name}' was defined by another cell:`}
-            </p>
+            <p>{`The variable '${error.name}' was defined by another cell:`}</p>
             <ul className="list-disc">
               {error.cells.map((cid) => (
                 <li className={liStyle} key={cid}>
@@ -97,7 +98,7 @@
       case "delete-nonlocal":
         return (
           <Fragment key={idx}>
-            <div className="mt-4">
+            <div>
               {`The variable '${error.name}' can't be deleted because it was defined by another cell (`}
               <CellLinkError cellId={error.cells[0] as CellId} />
               {")"}
@@ -110,7 +111,7 @@
         );
 
       case "interruption":
         titleContents = "Interrupted";
         return (
           <p key={idx}>{"This cell was interrupted and needs to be re-run."}</p>
         );
@@ -195,7 +196,7 @@
   });
 
   const title = (
-    <AlertTitle className="font-code font-bold mb-4 tracking-wide">
+    <AlertTitle className="font-code font-bold tracking-wide">
       {titleContents}
     </AlertTitle>
   );
@@ -204,14 +205,15 @@
     <Alert
       variant={alertVariant}
       className={cn(
-        `border-none font-code text-sm text-[0.84375rem] px-0 ${textColor} normal [&:has(svg)]:pl-0`,
+        `border-none font-code text-sm text-[0.84375rem] px-0 ${textColor} normal [&:has(svg)]:pl-0 space-y-4`,
         className,
       )}
     >
       {title}
       <div>
         <ul>{msgs}</ul>
       </div>
+      {cellId && <AutoFixButton errors={errors} cellId={cellId} />}
     </Alert>
   );
 };

frontend/src/core/cells/__tests__/runs.test.ts

@@ -15,7 +15,7 @@ const { reducer, initialState, isPureMarkdown } = exportedForTesting;
 
 function first<T>(map: Map<string, T> | undefined): T {
   invariant(map, "Map is undefined");
-  return map.values().next().value;
+  return map.values().next().value as T;
 }
 
 describe("RunsState Reducer", () => {

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

@@ -0,0 +1,62 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { describe, it, expect } from "vitest";
+import { getAutoFixes, getImportCode } from "../errors";
+import type { MarimoError } from "@/core/kernel/messages";
+
+describe("getImportCode", () => {
+  it("returns simple import for same name", () => {
+    expect(getImportCode("json")).toBe("import json");
+    expect(getImportCode("math")).toBe("import math");
+  });
+
+  it("returns aliased import for different names", () => {
+    expect(getImportCode("np")).toBe("import numpy as np");
+    expect(getImportCode("pd")).toBe("import pandas as pd");
+    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);
+  });
+});