@marimo-team/frontend 0.23.16-dev9 → 0.23.17-dev0

package/src/core/codemirror/code-lens/analyzer.ts

@@ -0,0 +1,284 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { syntaxTree } from "@codemirror/language";
+import type { EditorState } from "@codemirror/state";
+import type { SyntaxNode, Tree } from "@lezer/common";
+import {
+  NON_CODE_NODES,
+  PyKeyword,
+  PyNode,
+  SCOPE_CREATING_NODES,
+} from "../python-node-names";
+import { hasSyntaxErrors } from "../reactive-references/analyzer";
+
+export interface CodeLensTarget {
+  from: number;
+  to: number;
+  name: string;
+}
+
+/**
+ * Finds the first module-scope assignment site for each name in `names`.
+ */
+export function findDeclarationSites(options: {
+  state: EditorState;
+  names: ReadonlySet<string>;
+}): CodeLensTarget[] {
+  const { state, names } = options;
+  if (names.size === 0) {
+    return [];
+  }
+  const tree = syntaxTree(state);
+  if (!tree || hasSyntaxErrors(tree)) {
+    return [];
+  }
+
+  const firstSites = new Map<string, CodeLensTarget>();
+
+  const collectTargetNames = (node: SyntaxNode) => {
+    if (node.name === PyNode.VariableName) {
+      const name = state.doc.sliceString(node.from, node.to);
+      if (names.has(name) && !firstSites.has(name)) {
+        firstSites.set(name, { from: node.from, to: node.to, name });
+      }
+      return;
+    }
+    // Recurse into tuple/list unpacking, e.g. `(a, b) = ...` or `[a, b] = ...`,
+    // and parenthesized targets, e.g. `(df) = ...`.
+    // Member/subscript targets (`obj.attr = ...`) are not declarations.
+    if (
+      node.name === PyNode.TupleExpression ||
+      node.name === PyNode.ArrayExpression ||
+      node.name === PyNode.ParenthesizedExpression
+    ) {
+      for (let child = node.firstChild; child; child = child.nextSibling) {
+        collectTargetNames(child);
+      }
+    }
+  };
+
+  const collectAssignmentTargets = (assign: SyntaxNode) => {
+    // Targets are everything before the last AssignOp (handles `a = b = ...`).
+    // A bare annotation (`x: int`) has no AssignOp and is skipped.
+    let lastAssignOp = -1;
+    for (let child = assign.firstChild; child; child = child.nextSibling) {
+      if (child.name === PyNode.AssignOp) {
+        lastAssignOp = child.from;
+      }
+    }
+    if (lastAssignOp === -1) {
+      return;
+    }
+    for (
+      let child = assign.firstChild;
+      child && child.from < lastAssignOp;
+      child = child.nextSibling
+    ) {
+      collectTargetNames(child);
+    }
+  };
+
+  const visit = (node: SyntaxNode) => {
+    // Only module-scope assignments count as creation sites; top-level
+    // `if`/`for`/`try` bodies are still module scope.
+    if (SCOPE_CREATING_NODES.has(node.name)) {
+      return;
+    }
+    if (node.name === PyNode.AssignStatement) {
+      collectAssignmentTargets(node);
+      return;
+    }
+    for (let child = node.firstChild; child; child = child.nextSibling) {
+      visit(child);
+    }
+  };
+  visit(tree.topNode);
+
+  return [...firstSites.values()];
+}
+
+// Note that the cache pattern assumes users `import marimo as mo`.
+// This is a common convention for marimo notebooks.
+const CACHE_PATTERN = /\bmo\.(?:persistent_cache|cache)\b/g;
+
+export interface CacheSite {
+  from: number;
+  to: number;
+  /**
+   * Variable the cache is bound to: the decorated function name
+   * (`@mo.cache def f`), the assignment target (`g = mo.cache(f)`), or the
+   * `as` binding (`with mo.persistent_cache("k") as c`).
+   */
+  boundName: string | null;
+  /** First string argument, e.g. the persistent_cache name. */
+  cacheName: string | null;
+}
+
+function stripStringQuotes(text: string): string {
+  const match = text.match(/^[A-Za-z]*("""|'''|"|')([\s\S]*)\1$/);
+  return match ? match[2] : text;
+}
+
+/**
+ * Extracts a static cache name from the first positional argument or an
+ * explicit `name=` argument. Other string-valued options, such as
+ * `hash_type="content"` and `save_path="cache"`, are not cache names.
+ */
+function findStaticCacheName(
+  argList: SyntaxNode | null,
+  text: (node: SyntaxNode) => string,
+): string | null {
+  if (!argList) {
+    return null;
+  }
+
+  const isArgumentBoundary = (node: SyntaxNode | null) =>
+    node?.name === "," || node?.name === ")";
+
+  let firstArgument: SyntaxNode | null = null;
+  for (let child = argList.firstChild; child; child = child.nextSibling) {
+    if (child.name === "(" || child.name === ")" || child.name === ",") {
+      continue;
+    }
+    firstArgument ??= child;
+
+    if (
+      child.name === PyNode.VariableName &&
+      text(child) === "name" &&
+      child.nextSibling?.name === PyNode.AssignOp
+    ) {
+      const value = child.nextSibling.nextSibling;
+      if (
+        value?.name === PyNode.String &&
+        isArgumentBoundary(value.nextSibling)
+      ) {
+        return stripStringQuotes(text(value));
+      }
+    }
+  }
+
+  if (!firstArgument) {
+    return null;
+  }
+  if (
+    firstArgument.name === PyNode.String &&
+    isArgumentBoundary(firstArgument.nextSibling)
+  ) {
+    return stripStringQuotes(text(firstArgument));
+  }
+  return null;
+}
+
+/**
+ * Resolves a `mo.cache` / `mo.persistent_cache` match into a CacheSite:
+ * extends `to` past call arguments (`mo.cache(pin_modules=True)`), so an
+ * icon placed at `to` lands after the closing paren, and extracts the bound
+ * variable name and the cache-name string argument.
+ */
+function analyzeCacheSite(
+  state: EditorState,
+  tree: Tree,
+  from: number,
+  to: number,
+): CacheSite {
+  const text = (node: SyntaxNode) => state.doc.sliceString(node.from, node.to);
+  const node = tree.resolveInner(from, 1);
+  let boundName: string | null = null;
+  let argList: SyntaxNode | null = null;
+  let end = to;
+
+  // `mo.cache(...)` as an expression wraps in a CallExpression that starts
+  // at the match
+  let call: SyntaxNode | null = null;
+  for (
+    let ancestor: SyntaxNode | null = node.parent;
+    ancestor && ancestor.from === from;
+    ancestor = ancestor.parent
+  ) {
+    if (ancestor.name === PyNode.CallExpression) {
+      call = ancestor;
+      break;
+    }
+  }
+
+  if (call) {
+    end = call.to;
+    argList = call.getChild(PyNode.ArgList);
+    const parent = call.parent;
+    if (parent?.name === PyNode.AssignStatement) {
+      // `g = mo.cache(f)` — the direct VariableName child is the target
+      const target = parent.getChild(PyNode.VariableName);
+      boundName = target && target.from < call.from ? text(target) : null;
+    } else if (parent?.name === PyNode.WithStatement) {
+      // `with mo.persistent_cache("k") as c:` — `as c` follows the call
+      let child = parent.firstChild;
+      while (child) {
+        if (child.from < call.to || child.name === PyKeyword.As) {
+          child = child.nextSibling;
+          continue;
+        }
+        if (child.name === PyNode.VariableName) {
+          boundName = text(child);
+        }
+        break;
+      }
+    }
+  } else if (node.parent?.name === PyNode.Decorator) {
+    // `@mo.cache(...)` decorators are flat: the ArgList is a direct sibling
+    const decorator = node.parent;
+    const decoratorArgs = decorator.getChild(PyNode.ArgList);
+    if (decoratorArgs?.from === to) {
+      argList = decoratorArgs;
+      end = decoratorArgs.to;
+    }
+    // The decorated function's name
+    const fn = decorator.parent?.getChild(PyNode.FunctionDefinition);
+    const fnName = fn?.getChild(PyNode.VariableName);
+    boundName = fnName ? text(fnName) : null;
+  }
+
+  const cacheName = findStaticCacheName(argList, text);
+  return { from, to: end, boundName, cacheName };
+}
+
+/**
+ * True when the `mo` at `from` is a standalone module reference (the object of
+ * `mo.cache`), not an attribute of something else. This filters out chains
+ * like `obj.mo.cache(f)` — where `mo` is a `PropertyName` rather than a
+ * `VariableName` — and mentions inside comments or strings.
+ */
+function isStandaloneMoReference(
+  state: EditorState,
+  tree: Tree,
+  from: number,
+): boolean {
+  const node = tree.resolveInner(from, 1);
+  if (NON_CODE_NODES.has(node.name)) {
+    return false;
+  }
+  // In `obj.mo.cache`, `mo` parses as a `PropertyName`; a standalone `mo`
+  // (including whitespace-separated `mo . cache`) parses as a `VariableName`.
+  return (
+    node.name === PyNode.VariableName &&
+    state.doc.sliceString(node.from, node.to) === "mo"
+  );
+}
+
+/**
+ * Finds occurrences of `mo.cache` / `mo.persistent_cache` (as a decorator,
+ * call, or context manager) with a simple text match, skipping mentions in
+ * comments, strings, and attribute chains (`obj.mo.cache`). `to` extends past
+ * call arguments when present.
+ */
+export function findCacheSites(state: EditorState): CacheSite[] {
+  const tree = syntaxTree(state);
+  if (!tree || hasSyntaxErrors(tree)) {
+    return [];
+  }
+  return [...state.doc.toString().matchAll(CACHE_PATTERN)]
+    .filter((match) => isStandaloneMoReference(state, tree, match.index ?? 0))
+    .map((match) => {
+      const from = match.index ?? 0;
+      return analyzeCacheSite(state, tree, from, from + match[0].length);
+    });
+}

package/src/core/codemirror/code-lens/entities.ts

@@ -0,0 +1,123 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import type { CellId } from "@/core/cells/ids";
+import { dataConnectionsMapAtom } from "@/core/datasets/data-source-connections";
+import { INTERNAL_SQL_ENGINES } from "@/core/datasets/engines";
+import { datasetTablesAtom } from "@/core/datasets/state";
+import { store } from "@/core/state/jotai";
+import { storageNamespacesAtom } from "@/core/storage/state";
+import { variablesAtom } from "@/core/variables/state";
+import type { VariableName } from "@/core/variables/types";
+import { Logger } from "@/utils/Logger";
+
+export type CodeLensKind = "table" | "connection" | "bucket" | "cache";
+
+export interface CodeLensSpec {
+  /** Document position the icon is anchored at */
+  pos: number;
+  kind: CodeLensKind;
+  /** Variable name for entities; a stable identity key for cache sites */
+  name: string;
+  /** Extra context for cache sites */
+  cache?: {
+    boundName: string | null;
+    cacheName: string | null;
+  };
+}
+
+/**
+ * Sets `name` to `kind`, logging (and overwriting) if the name was already
+ * claimed by a different kind. This should be rare.
+ */
+function setEntity(
+  entities: Map<string, CodeLensKind>,
+  name: string,
+  kind: CodeLensKind,
+): void {
+  const existing = entities.get(name);
+  if (existing !== undefined && existing !== kind) {
+    Logger.error(
+      `[code-lens] "${name}" is claimed by both a ${existing} and a ${kind} entity; keeping ${kind}`,
+    );
+  }
+  entities.set(name, kind);
+}
+
+/**
+ * All datasource (dataframe/SQL engine) and storage bucket names in the
+ * notebook, keyed by name, independent of any particular cell.
+ */
+const lensEntityKindsAtom = atom((get) => {
+  const entities = new Map<string, CodeLensKind>();
+  for (const table of get(datasetTablesAtom)) {
+    if (table.variable_name) {
+      setEntity(entities, table.variable_name, "table");
+    }
+  }
+  for (const name of get(dataConnectionsMapAtom).keys()) {
+    if (!INTERNAL_SQL_ENGINES.has(name)) {
+      setEntity(entities, name, "connection");
+    }
+  }
+  for (const namespace of get(storageNamespacesAtom)) {
+    setEntity(entities, namespace.name, "bucket");
+  }
+  return entities;
+});
+
+interface LensEntitiesByCell {
+  byCell: Map<CellId, Map<string, CodeLensKind>>;
+  /**
+   * Entities the kernel hasn't reported a variable for yet (still executing,
+   * or just typed), so they can't be attributed to a declaring cell.
+   */
+  pending: Map<string, CodeLensKind>;
+}
+
+/**
+ * Groups `lensEntityKindsAtom` by declaring cell, once per store update.
+ * With C cells and T datasource/bucket entities, this turns
+ * an O(C x T) per-round cost (every cell re-scanning every entity) into O(T) here.
+ */
+const lensEntitiesByCellAtom = atom((get): LensEntitiesByCell => {
+  const kinds = get(lensEntityKindsAtom);
+  const variables = get(variablesAtom);
+  const byCell = new Map<CellId, Map<string, CodeLensKind>>();
+  const pending = new Map<string, CodeLensKind>();
+
+  for (const [name, kind] of kinds) {
+    const variable = variables[name as VariableName];
+    if (variable == null) {
+      pending.set(name, kind);
+      continue;
+    }
+    for (const declaringCell of variable.declaredBy) {
+      let forCell = byCell.get(declaringCell);
+      if (!forCell) {
+        forCell = new Map();
+        byCell.set(declaringCell, forCell);
+      }
+      forCell.set(name, kind);
+    }
+  }
+
+  return { byCell, pending };
+});
+
+/**
+ * Variables declared in `cellId` that are datasources (dataframes and SQL
+ * engines) or storage buckets, keyed by variable name.
+ */
+export function getLensEntities(cellId: CellId): Map<string, CodeLensKind> {
+  const { byCell, pending } = store.get(lensEntitiesByCellAtom);
+  const entities = new Map(byCell.get(cellId));
+  // Be permissive when the kernel hasn't reported a variable yet: any cell
+  // could turn out to be the declaring one.
+  for (const [name, kind] of pending) {
+    if (!entities.has(name)) {
+      entities.set(name, kind);
+    }
+  }
+  return entities;
+}

package/src/core/codemirror/code-lens/extension.ts

@@ -0,0 +1,298 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { type Extension, StateEffect, StateField } from "@codemirror/state";
+import {
+  Decoration,
+  type DecorationSet,
+  EditorView,
+  showTooltip,
+  type Tooltip,
+  ViewPlugin,
+  type ViewUpdate,
+  WidgetType,
+} from "@codemirror/view";
+import { type DebouncedFunc, debounce } from "lodash-es";
+import type { CellId } from "@/core/cells/ids";
+import { getFeatureFlag } from "@/core/config/feature-flag";
+import { dataConnectionsMapAtom } from "@/core/datasets/data-source-connections";
+import { datasetTablesAtom } from "@/core/datasets/state";
+import { store } from "@/core/state/jotai";
+import { storageNamespacesAtom } from "@/core/storage/state";
+import { variablesAtom } from "@/core/variables/state";
+import { languageAdapterState } from "../language/extension";
+import { openLensTarget } from "./actions";
+import { findCacheSites, findDeclarationSites } from "./analyzer";
+import { type CodeLensSpec, getLensEntities } from "./entities";
+import { LENS_ICONS, LENS_TOOLTIPS } from "./icons";
+import { mountLensPopover } from "./popover";
+
+// Delay (in ms) before showing the hover popover, matching the app's
+// tooltip delay
+export const CODE_LENS_HOVER_DELAY_MS = 400;
+
+const setCodeLenses = StateEffect.define<CodeLensSpec[]>();
+const setHoveredLens = StateEffect.define<CodeLensSpec | null>();
+
+// Pending hover timers, keyed by widget element so `destroy` can cancel them
+const HOVER_TIMERS = new WeakMap<HTMLElement, number>();
+
+class CodeLensWidget extends WidgetType {
+  private readonly spec: CodeLensSpec;
+
+  constructor(spec: CodeLensSpec) {
+    super();
+    this.spec = spec;
+  }
+
+  override eq(other: CodeLensWidget): boolean {
+    return (
+      // `pos` is captured by the DOM hover/click handlers, so a reused widget
+      // whose anchor moved must not be treated as equal
+      this.spec.pos === other.spec.pos &&
+      this.spec.kind === other.spec.kind &&
+      this.spec.name === other.spec.name &&
+      this.spec.cache?.boundName === other.spec.cache?.boundName &&
+      this.spec.cache?.cacheName === other.spec.cache?.cacheName
+    );
+  }
+
+  override toDOM(view: EditorView): HTMLElement {
+    const { spec } = this;
+    const element = document.createElement("span");
+    element.className = "mo-code-lens";
+    element.setAttribute("role", "button");
+    // Focusable so the action is reachable and activatable by keyboard
+    element.tabIndex = 0;
+    // No `title`: the native tooltip is replaced by the hover popover
+    element.setAttribute("aria-label", LENS_TOOLTIPS[spec.kind]);
+    // Static, trusted markup (see icons.ts)
+    element.innerHTML = LENS_ICONS[spec.kind];
+    const hidePopover = () => {
+      window.clearTimeout(HOVER_TIMERS.get(element));
+      HOVER_TIMERS.delete(element);
+      if (view.state.field(codeLensHoverField, false)) {
+        view.dispatch({ effects: setHoveredLens.of(null) });
+      }
+    };
+    element.onmouseenter = () => {
+      const timer = window.setTimeout(() => {
+        view.dispatch({ effects: setHoveredLens.of(spec) });
+      }, CODE_LENS_HOVER_DELAY_MS);
+      HOVER_TIMERS.set(element, timer);
+    };
+    element.onmouseleave = hidePopover;
+    element.onmousemove = (event) => {
+      // Keep the editor's built-in hover documentation tooltip from also
+      // triggering over the icon
+      event.stopPropagation();
+    };
+    element.onmousedown = (event) => {
+      // Don't move the cursor or steal focus from the editor
+      event.preventDefault();
+      event.stopPropagation();
+    };
+    element.onclick = (event) => {
+      event.preventDefault();
+      event.stopPropagation();
+      hidePopover();
+      openLensTarget(spec.kind);
+    };
+    element.onkeydown = (event) => {
+      if (event.key === "Enter" || event.key === " ") {
+        event.preventDefault();
+        event.stopPropagation();
+        hidePopover();
+        openLensTarget(spec.kind);
+      }
+    };
+    return element;
+  }
+
+  override destroy(dom: HTMLElement): void {
+    window.clearTimeout(HOVER_TIMERS.get(dom));
+    HOVER_TIMERS.delete(dom);
+  }
+
+  override ignoreEvent(): boolean {
+    // The widget handles its own events
+    return true;
+  }
+}
+
+function createLensTooltip(spec: CodeLensSpec): Tooltip {
+  return {
+    pos: spec.pos,
+    above: true,
+    create: () => {
+      const dom = document.createElement("div");
+      // Same chrome as the SQL completion/hover popovers
+      dom.classList.add("mo-cm-tooltip", "docs-documentation");
+      const unmount = mountLensPopover(dom, spec);
+      return { dom, resize: false, destroy: unmount };
+    },
+  };
+}
+
+const codeLensHoverField = StateField.define<Tooltip | null>({
+  create() {
+    return null;
+  },
+  update(tooltip, tr) {
+    for (const effect of tr.effects) {
+      if (effect.is(setHoveredLens)) {
+        return effect.value ? createLensTooltip(effect.value) : null;
+      }
+      if (effect.is(setCodeLenses)) {
+        // Lenses were rebuilt; the hovered widget may be gone
+        return null;
+      }
+    }
+    return tr.docChanged ? null : tooltip;
+  },
+  provide: (field) => showTooltip.from(field),
+});
+
+const codeLensField = StateField.define<DecorationSet>({
+  create() {
+    return Decoration.none;
+  },
+  update(decorations, tr) {
+    decorations = decorations.map(tr.changes);
+    for (const effect of tr.effects) {
+      if (effect.is(setCodeLenses)) {
+        decorations = Decoration.set(
+          effect.value.map((spec) =>
+            Decoration.widget({
+              widget: new CodeLensWidget(spec),
+              side: 1,
+            }).range(spec.pos),
+          ),
+          // Variable and cache lenses are collected separately, so sort
+          true,
+        );
+      }
+    }
+    return decorations;
+  },
+  provide: (field) => EditorView.decorations.from(field),
+});
+
+/**
+ * Plugin that keeps code lens decorations in sync with the document and the
+ * datasource/bucket state.
+ */
+class CodeLensPlugin {
+  private readonly view: EditorView;
+  private readonly cellId: CellId;
+  private readonly includeCache: boolean;
+  private readonly unsubscribes: Array<() => void>;
+
+  // Delay (in ms) before recomputing lenses after user changes or store updates
+  private readonly debounceMs = 300;
+  private readonly scheduleUpdate: DebouncedFunc<() => void>;
+
+  constructor(view: EditorView, cellId: CellId, includeCache: boolean) {
+    this.view = view;
+    this.cellId = cellId;
+    this.includeCache = includeCache;
+    this.scheduleUpdate = debounce(() => this.run(), this.debounceMs);
+    const onStoreChange = () => this.scheduleUpdate();
+    this.unsubscribes = [
+      store.sub(datasetTablesAtom, onStoreChange),
+      store.sub(dataConnectionsMapAtom, onStoreChange),
+      store.sub(storageNamespacesAtom, onStoreChange),
+      // The declaring-cell filter depends on variables
+      store.sub(variablesAtom, onStoreChange),
+    ];
+    this.scheduleUpdate();
+  }
+
+  update(update: ViewUpdate) {
+    // Recompute on edits, and when the cell's language changes (e.g. Python ->
+    // SQL) so stale Python-only icons are cleared even if the text is unchanged
+    const adapterChanged =
+      update.startState.field(languageAdapterState, false)?.type !==
+      update.state.field(languageAdapterState, false)?.type;
+    if (update.docChanged || adapterChanged) {
+      this.scheduleUpdate();
+    }
+  }
+
+  destroy() {
+    this.scheduleUpdate.cancel();
+    for (const unsubscribe of this.unsubscribes) {
+      unsubscribe();
+    }
+  }
+
+  private run() {
+    const { state } = this.view;
+    const lenses: CodeLensSpec[] = [];
+
+    // Only python cells: SQL/markdown docs aren't python, and a cache icon
+    // inside a SQL string would be misleading.
+    const adapterType = state.field(languageAdapterState, false)?.type;
+    if (adapterType == null || adapterType === "python") {
+      const entities = getLensEntities(this.cellId);
+      const targets = findDeclarationSites({
+        state,
+        names: new Set(entities.keys()),
+      });
+      for (const target of targets) {
+        const kind = entities.get(target.name);
+        if (kind) {
+          lenses.push({ pos: target.to, kind, name: target.name });
+        }
+      }
+      if (this.includeCache) {
+        for (const site of findCacheSites(state)) {
+          lenses.push({
+            pos: site.to,
+            kind: "cache",
+            name: `cache:${site.from}`,
+            cache: { boundName: site.boundName, cacheName: site.cacheName },
+          });
+        }
+      }
+    }
+
+    // Defer dispatch to avoid triggering during an editor update cycle
+    queueMicrotask(() => {
+      this.view.dispatch({ effects: setCodeLenses.of(lenses) });
+    });
+  }
+}
+
+const codeLensTheme = EditorView.baseTheme({
+  ".mo-code-lens": {
+    display: "inline-flex",
+    verticalAlign: "baseline",
+    // Optically center the icon against the text
+    transform: "translateY(1.5px)",
+    marginLeft: "0.3em",
+    cursor: "pointer",
+    opacity: "0.5",
+  },
+  ".mo-code-lens:hover": {
+    opacity: "1",
+  },
+});
+
+/**
+ * Inline icons linking datasource/bucket variables and `mo.cache` /
+ * `mo.persistent_cache` calls to their panels.
+ * Gated behind the `editor_code_lens` feature flag.
+ */
+export function codeLensBundle(cellId: CellId): Extension {
+  if (!getFeatureFlag("editor_code_lens")) {
+    return [];
+  }
+  return [
+    codeLensField,
+    codeLensHoverField,
+    ViewPlugin.define(
+      (view) => new CodeLensPlugin(view, cellId, getFeatureFlag("cache_panel")),
+    ),
+    codeLensTheme,
+  ];
+}