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

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

@@ -0,0 +1,224 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { python } from "@codemirror/lang-python";
+import { EditorState } from "@codemirror/state";
+import { describe, expect, it } from "vitest";
+import { findCacheSites, findDeclarationSites } from "../analyzer";
+
+function pythonState(code: string) {
+  return EditorState.create({ doc: code, extensions: [python()] });
+}
+
+function run(names: string[], code: string) {
+  return findDeclarationSites({
+    state: pythonState(code),
+    names: new Set(names),
+  });
+}
+
+describe("findDeclarationSites", () => {
+  it("finds a simple assignment", () => {
+    const code = 'df = pd.read_csv("a.csv")';
+    const targets = run(["df"], code);
+    expect(targets).toEqual([{ from: 0, to: 2, name: "df" }]);
+    expect(code.slice(targets[0].from, targets[0].to)).toBe("df");
+  });
+
+  it("only matches the given names", () => {
+    const targets = run(["df"], "df = 1\nother = 2");
+    expect(targets.map((t) => t.name)).toEqual(["df"]);
+  });
+
+  it("handles annotated assignments", () => {
+    const targets = run(["df", "DataFrame"], "df: DataFrame = load()");
+    expect(targets.map((t) => t.name)).toEqual(["df"]);
+  });
+
+  it("handles chained assignments", () => {
+    const targets = run(["a", "b"], "a = b = load()");
+    expect(targets.map((t) => t.name).toSorted()).toEqual(["a", "b"]);
+  });
+
+  it("handles tuple and list unpacking", () => {
+    expect(run(["x", "y"], "x, y = f()").map((t) => t.name)).toEqual([
+      "x",
+      "y",
+    ]);
+    expect(run(["a", "b"], "(a, b) = f()").map((t) => t.name)).toEqual([
+      "a",
+      "b",
+    ]);
+    expect(run(["a", "b"], "[a, b] = f()").map((t) => t.name)).toEqual([
+      "a",
+      "b",
+    ]);
+  });
+
+  it("handles a parenthesized single target", () => {
+    expect(run(["df"], "(df) = load()").map((t) => t.name)).toEqual(["df"]);
+  });
+
+  it("returns only the first assignment site per name", () => {
+    const targets = run(["df"], "df = 1\ndf = 2");
+    expect(targets).toEqual([{ from: 0, to: 2, name: "df" }]);
+  });
+
+  it("ignores attribute and subscript targets", () => {
+    expect(run(["df"], "obj.df = 1")).toEqual([]);
+    expect(run(["df"], "data[df] = 1")).toEqual([]);
+  });
+
+  it("ignores augmented assignments", () => {
+    expect(run(["df"], "df += 1")).toEqual([]);
+  });
+
+  it("ignores assignments in function, lambda, and class scopes", () => {
+    expect(run(["df"], "def f():\n    df = 1")).toEqual([]);
+    expect(run(["df"], "class A:\n    df = 1")).toEqual([]);
+  });
+
+  it("finds assignments in top-level control flow", () => {
+    const targets = run(["df"], "if cond:\n    df = load()");
+    expect(targets.map((t) => t.name)).toEqual(["df"]);
+  });
+
+  it("ignores usages that are not assignments", () => {
+    expect(run(["df"], "print(df)")).toEqual([]);
+  });
+
+  it("returns nothing for syntax errors", () => {
+    expect(run(["df"], "df = (")).toEqual([]);
+  });
+
+  it("returns nothing for an empty name set", () => {
+    expect(run([], "df = 1")).toEqual([]);
+  });
+});
+
+describe("findCacheSites", () => {
+  function runCache(code: string) {
+    return findCacheSites(pythonState(code)).map((site) =>
+      code.slice(site.from, site.to),
+    );
+  }
+
+  it("matches decorator usage", () => {
+    expect(runCache("@mo.cache\ndef f():\n    return 1")).toEqual(["mo.cache"]);
+  });
+
+  it("extends past decorator arguments", () => {
+    expect(
+      runCache("@mo.cache(pin_modules=True)\ndef f():\n    return 1"),
+    ).toEqual(["mo.cache(pin_modules=True)"]);
+  });
+
+  it("extends past call arguments", () => {
+    expect(runCache("g = mo.cache(f)")).toEqual(["mo.cache(f)"]);
+    expect(runCache("g = mo.cache(f, pin_modules=True)")).toEqual([
+      "mo.cache(f, pin_modules=True)",
+    ]);
+  });
+
+  it("does not extend a bare reference or an enclosing call", () => {
+    expect(runCache("g = mo.cache")).toEqual(["mo.cache"]);
+    expect(runCache("print(mo.cache, 1)")).toEqual(["mo.cache"]);
+  });
+
+  it("matches persistent_cache context manager usage", () => {
+    expect(runCache('with mo.persistent_cache("k"):\n    pass')).toEqual([
+      'mo.persistent_cache("k")',
+    ]);
+  });
+
+  it("matches multiple occurrences", () => {
+    expect(runCache("mo.cache(f)\nmo.persistent_cache(g)")).toHaveLength(2);
+  });
+
+  it("does not match similar names", () => {
+    expect(runCache("memo.cache(f)")).toEqual([]);
+    expect(runCache("mo.cached(f)")).toEqual([]);
+    expect(runCache("mo.cache_info()")).toEqual([]);
+  });
+
+  it("does not match attribute chains like obj.mo.cache", () => {
+    expect(runCache("obj.mo.cache(f)")).toEqual([]);
+    expect(runCache("self.mo.persistent_cache(g)")).toEqual([]);
+  });
+
+  it("does not match mentions in comments or strings", () => {
+    expect(runCache("# uses mo.cache internally")).toEqual([]);
+    expect(runCache('x = "mo.cache"')).toEqual([]);
+    expect(runCache("mo.cache(f)  # not mo.persistent_cache")).toEqual([
+      "mo.cache(f)",
+    ]);
+  });
+
+  it("returns nothing for syntax errors", () => {
+    expect(runCache("mo.cache(f)\ndf = (")).toEqual([]);
+  });
+
+  describe("bound name and cache name extraction", () => {
+    function sites(code: string) {
+      return findCacheSites(pythonState(code)).map(
+        ({ boundName, cacheName }) => ({ boundName, cacheName }),
+      );
+    }
+
+    it("extracts the decorated function name", () => {
+      expect(sites("@mo.cache\ndef add(a, b):\n    return a + b")).toEqual([
+        { boundName: "add", cacheName: null },
+      ]);
+      expect(
+        sites("@mo.cache(pin_modules=True)\ndef add(a, b):\n    return a + b"),
+      ).toEqual([{ boundName: "add", cacheName: null }]);
+    });
+
+    it("extracts the assignment target", () => {
+      expect(sites("g = mo.cache(f)")).toEqual([
+        { boundName: "g", cacheName: null },
+      ]);
+    });
+
+    it("extracts the with-statement binding and cache name", () => {
+      expect(sites('with mo.persistent_cache("k") as c:\n    pass')).toEqual([
+        { boundName: "c", cacheName: "k" },
+      ]);
+      expect(sites('with mo.persistent_cache("k"):\n    pass')).toEqual([
+        { boundName: null, cacheName: "k" },
+      ]);
+      expect(sites('with mo.persistent_cache(name="k"):\n    pass')).toEqual([
+        { boundName: null, cacheName: "k" },
+      ]);
+    });
+
+    it("ignores string-valued options that are not cache names", () => {
+      expect(
+        sites(
+          '@mo.cache(hash_type="content")\ndef add(a, b):\n    return a + b',
+        ),
+      ).toEqual([{ boundName: "add", cacheName: null }]);
+      expect(
+        sites(
+          '@mo.persistent_cache(save_path="cache")\ndef add(a, b):\n    return a + b',
+        ),
+      ).toEqual([{ boundName: "add", cacheName: null }]);
+      expect(
+        sites('with mo.persistent_cache(prefix, save_path="cache"):\n    pass'),
+      ).toEqual([{ boundName: null, cacheName: null }]);
+    });
+
+    it("extracts an explicit name after other keyword arguments", () => {
+      expect(
+        sites(
+          'with mo.persistent_cache(save_path="cache", name="k"):\n    pass',
+        ),
+      ).toEqual([{ boundName: null, cacheName: "k" }]);
+    });
+
+    it("extracts nothing for bare expressions", () => {
+      expect(sites("mo.cache(f)")).toEqual([
+        { boundName: null, cacheName: null },
+      ]);
+    });
+  });
+});

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

@@ -0,0 +1,175 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId, variableName } from "@/__tests__/branded";
+import {
+  type DataSourceConnection,
+  dataSourceConnectionsAtom,
+} from "@/core/datasets/data-source-connections";
+import { type ConnectionName, DUCKDB_ENGINE } from "@/core/datasets/engines";
+import { datasetsAtom } from "@/core/datasets/state";
+import type { QualifiedColumn } from "@/core/datasets/types";
+import type { DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
+import { storageAtom } from "@/core/storage/state";
+import type { StorageNamespace } from "@/core/storage/types";
+import { variablesAtom } from "@/core/variables/state";
+import type { Variables } from "@/core/variables/types";
+import { Logger } from "@/utils/Logger";
+import { getLensEntities } from "../entities";
+
+const CELL = cellId("cell1");
+const OTHER_CELL = cellId("other-cell");
+
+const table = (name: string): DataTable => ({
+  name,
+  source: "memory",
+  source_type: "local",
+  num_rows: 1,
+  num_columns: 1,
+  columns: [],
+  variable_name: variableName(name),
+});
+
+const namespace = (name: string): StorageNamespace => ({
+  name: variableName(name),
+  displayName: name,
+  protocol: "s3",
+  rootPath: `s3://${name}`,
+  backendType: "obstore",
+  storageEntries: [],
+});
+
+const connection = (name: string): DataSourceConnection => ({
+  name: name as ConnectionName,
+  source: "postgres",
+  dialect: "postgresql",
+  display_name: name,
+  databases: [],
+});
+
+function seedStore(opts: {
+  tables?: DataTable[];
+  namespaces?: StorageNamespace[];
+  connections?: DataSourceConnection[];
+  variables?: Variables;
+}) {
+  store.set(datasetsAtom, {
+    tables: opts.tables ?? [],
+    expandedTables: new Set<string>(),
+    expandedColumns: new Set<QualifiedColumn>(),
+    columnsPreviews: new Map(),
+  });
+  store.set(storageAtom, {
+    namespaces: opts.namespaces ?? [],
+    entriesByPath: new Map(),
+    pageMetadataByPath: new Map(),
+  });
+  store.set(dataSourceConnectionsAtom, {
+    latestEngineSelected: DUCKDB_ENGINE,
+    connectionsMap: new Map((opts.connections ?? []).map((c) => [c.name, c])),
+  });
+  store.set(variablesAtom, opts.variables ?? {});
+}
+
+describe("getLensEntities", () => {
+  beforeEach(() => {
+    seedStore({});
+    vi.restoreAllMocks();
+  });
+
+  it("returns an empty map when there are no entities", () => {
+    expect(getLensEntities(CELL)).toEqual(new Map());
+  });
+
+  it("includes tables, connections, and buckets declared in the cell", () => {
+    seedStore({
+      tables: [table("df")],
+      connections: [connection("engine")],
+      namespaces: [namespace("bucket")],
+    });
+
+    const entities = getLensEntities(CELL);
+    expect(entities).toEqual(
+      new Map([
+        ["df", "table"],
+        ["engine", "connection"],
+        ["bucket", "bucket"],
+      ]),
+    );
+  });
+
+  it("excludes the internal DuckDB engine from connections", () => {
+    seedStore({ connections: [connection(DUCKDB_ENGINE)] });
+    expect(getLensEntities(CELL)).toEqual(new Map());
+  });
+
+  it("excludes entities declared by a different cell", () => {
+    seedStore({
+      tables: [table("df")],
+      variables: {
+        [variableName("df")]: {
+          name: variableName("df"),
+          declaredBy: [OTHER_CELL],
+          usedBy: [],
+        },
+      },
+    });
+
+    expect(getLensEntities(CELL)).toEqual(new Map());
+  });
+
+  it("is permissive when the kernel hasn't reported the variable yet", () => {
+    // No entry in `variablesAtom` for "df" at all.
+    seedStore({ tables: [table("df")] });
+    expect(getLensEntities(CELL)).toEqual(new Map([["df", "table"]]));
+  });
+
+  it("last-write-wins and logs an error when a name is claimed by more than one kind", () => {
+    const errorSpy = vi
+      .spyOn(Logger, "error")
+      .mockImplementation(() => undefined);
+    seedStore({
+      tables: [table("shared")],
+      connections: [connection("shared")],
+    });
+
+    const entities = getLensEntities(CELL);
+
+    expect(entities.get("shared")).toBe("connection");
+    expect(errorSpy).toHaveBeenCalledExactlyOnceWith(
+      expect.stringContaining("shared"),
+    );
+  });
+
+  it("does not log when the same name repeats with the same kind", () => {
+    const errorSpy = vi
+      .spyOn(Logger, "error")
+      .mockImplementation(() => undefined);
+    seedStore({ tables: [table("df"), table("df")] });
+
+    const entities = getLensEntities(CELL);
+
+    expect(entities.get("df")).toBe("table");
+    expect(errorSpy).not.toHaveBeenCalled();
+  });
+
+  it("shares the underlying entity map across cells instead of rebuilding it per call", () => {
+    // Regression test: the table/connection/namespace scan (and its
+    // collision logging) must run once and be reused by every cell, not
+    // once per `getLensEntities` call.
+    const errorSpy = vi
+      .spyOn(Logger, "error")
+      .mockImplementation(() => undefined);
+    seedStore({
+      tables: [table("shared")],
+      connections: [connection("shared")],
+    });
+
+    getLensEntities(CELL);
+    getLensEntities(OTHER_CELL);
+    getLensEntities(CELL);
+
+    expect(errorSpy).toHaveBeenCalledTimes(1);
+  });
+});