diff --git a/extension/src/commands.ts b/extension/src/commands.ts index b05fce17..e3711b5c 100644 --- a/extension/src/commands.ts +++ b/extension/src/commands.ts @@ -122,6 +122,8 @@ export type VscodeBuiltinCommand = | "workbench.action.quickOpen" // Toggle Outputs | "notebook.cell.toggleOutputs" + // Collapse a cell's input (code) editor + | "notebook.cell.collapseCellInput" // Fold Cell | "notebook.fold" // Unfold Cell diff --git a/extension/src/features/HideCodeSync.ts b/extension/src/features/HideCodeSync.ts new file mode 100644 index 00000000..819cc16b --- /dev/null +++ b/extension/src/features/HideCodeSync.ts @@ -0,0 +1,103 @@ +import { Effect, HashSet, Layer, Option, Ref, Stream } from "effect"; + +import { VsCode } from "../platform/VsCode.ts"; +import { + type MarimoNotebookCell, + MarimoNotebookDocument, + type NotebookId, +} from "../schemas/MarimoNotebookDocument.ts"; + +/** + * A contiguous, end-exclusive range of cells, addressed by index. + * + * This matches VS Code's `ICellRange`, the shape `notebook.cell.collapseCellInput` + * accepts to pick which cells to act on. + */ +export interface CellRange { + readonly start: number; + readonly end: number; +} + +/** + * Returns one single-cell range per `hide_code` cell, in document order. + * + * It is a plain function rather than part of the layer below so the selection + * logic can be tested without standing up a notebook editor. + */ +export function hiddenInputCellRanges( + cells: readonly MarimoNotebookCell[], +): CellRange[] { + return cells + .filter((cell) => cell.isCodeHidden) + .map((cell) => ({ start: cell.index, end: cell.index + 1 })); +} + +/** + * Collapses the code input of `hide_code` cells when a marimo notebook opens. + * + * VS Code exposes the input-collapsed state as view-only: we can neither read + * it nor set it through cell metadata, only fire `notebook.cell.collapseCellInput`. + * So this is a one-way sync, not a binding. We collapse a notebook's `hide_code` + * cells the first time it becomes active and track it per session, so refocusing + * the tab does not re-hide a cell the user expanded to edit. See issue #326. + */ +export const HideCodeSyncLive = Layer.scopedDiscard( + Effect.gen(function* () { + const code = yield* VsCode; + + // Notebooks already collapsed this session; refocusing a tab must not + // re-hide a cell the user has since expanded. + const collapsed = yield* Ref.make(HashSet.empty()); + + yield* Effect.forkScoped( + // Prepend the currently-active editor: onDidChangeActiveNotebookEditor + // only emits future changes, so this covers a notebook open at startup. + Stream.concat( + Stream.fromEffect(code.window.getActiveNotebookEditor()), + code.window.activeNotebookEditorChanges(), + ).pipe( + Stream.runForEach( + Effect.fn("HideCodeSync.collapseHiddenCells")(function* (editor) { + const notebook = Option.filterMap(editor, (editor) => + MarimoNotebookDocument.tryFrom(editor.notebook), + ); + if (Option.isNone(notebook)) { + return; + } + + const id = notebook.value.id; + if (HashSet.has(yield* Ref.get(collapsed), id)) { + return; + } + + const ranges = hiddenInputCellRanges(notebook.value.getCells()); + if (ranges.length === 0) { + return; + } + + yield* Effect.annotateCurrentSpan({ + notebook: id, + hiddenCells: ranges.length, + }); + + yield* code.commands + .executeCommand("notebook.cell.collapseCellInput", { + ranges, + document: notebook.value.uri, + }) + .pipe( + // Mark only after a successful collapse, so a transient failure + // retries on the next activation. + Effect.andThen(Ref.update(collapsed, HashSet.add(id))), + Effect.catchAll((error) => + Effect.logWarning("Failed to collapse hidden cells").pipe( + Effect.annotateLogs({ error }), + ), + ), + ); + }), + ), + ), + ); + }), +); diff --git a/extension/src/features/Main.ts b/extension/src/features/Main.ts index 5dc6b256..6ea39098 100644 --- a/extension/src/features/Main.ts +++ b/extension/src/features/Main.ts @@ -53,6 +53,7 @@ import type { Telemetry } from "../telemetry/Telemetry.ts"; import { CellMetadataBindingsLive } from "./CellMetadataBindings.ts"; import { CellStatusBarProviderLive } from "./CellStatusBarProvider.ts"; import { DebugLayerLive } from "./DebugLayer.ts"; +import { HideCodeSyncLive } from "./HideCodeSync.ts"; import { MarimoCodeLensProviderLive } from "./MarimoCodeLensProvider.ts"; import { MarimoFileDetectorLive } from "./MarimoFileDetector.ts"; import { RegisterCommandsLive } from "./RegisterCommands.ts"; @@ -80,6 +81,7 @@ const MainLive = Layer.empty Layer.merge(CellMetadataBindingsLive), Layer.merge(ReloadOnConfigChangeLive), Layer.merge(ThemeSyncLive), + Layer.merge(HideCodeSyncLive), Layer.merge(DebugLayerLive), ) .pipe( diff --git a/extension/src/features/__tests__/HideCodeSync.test.ts b/extension/src/features/__tests__/HideCodeSync.test.ts new file mode 100644 index 00000000..9a5eff4b --- /dev/null +++ b/extension/src/features/__tests__/HideCodeSync.test.ts @@ -0,0 +1,132 @@ +import { describe, expect, it } from "@effect/vitest"; +import { Effect, Layer, Option, Ref, TestClock } from "effect"; + +import { + createNotebookCell, + createTestNotebookDocument, + TestVsCode, +} from "../../__mocks__/TestVsCode.ts"; +import { MarimoNotebookCell } from "../../schemas/MarimoNotebookDocument.ts"; +import { + type CellRange, + hiddenInputCellRanges, + HideCodeSyncLive, +} from "../HideCodeSync.ts"; + +const cell = (index: number, hideCode: boolean) => + MarimoNotebookCell.from( + createNotebookCell( + createTestNotebookDocument("/test/notebook_mo.py"), + { + kind: 2, + value: "", + languageId: "python", + metadata: { + stableId: `cell-${index}`, + options: { hide_code: hideCode }, + }, + }, + index, + ), + ); + +describe("hiddenInputCellRanges", () => { + it("returns one end-exclusive range per hide_code cell, by index", () => { + const cells = [ + cell(0, false), + cell(1, true), + cell(2, false), + cell(3, true), + ]; + expect(hiddenInputCellRanges(cells)).toEqual([ + { start: 1, end: 2 }, + { start: 3, end: 4 }, + ]); + }); + + it("returns no ranges when no cell hides its code", () => { + expect(hiddenInputCellRanges([cell(0, false)])).toEqual([]); + }); +}); + +const isCellRange = (x: unknown): x is CellRange => + typeof x === "object" && + x !== null && + "start" in x && + typeof x.start === "number" && + "end" in x && + typeof x.end === "number"; + +const collapseRanges = (arg: unknown): readonly CellRange[] | undefined => + typeof arg === "object" && + arg !== null && + "ranges" in arg && + Array.isArray(arg.ranges) && + arg.ranges.every(isCellRange) + ? arg.ranges + : undefined; + +/** The ranges passed to each `notebook.cell.collapseCellInput` call, in order. */ +const collapses = Effect.fn(function* (vscode: TestVsCode) { + const executions = yield* Ref.get(vscode.executions); + return executions + .filter((e) => e.command === "notebook.cell.collapseCellInput") + .map((e) => collapseRanges(e.args[0])); +}); + +const withTestCtx = Effect.fn(function* (hideCode: ReadonlyArray) { + const editor = TestVsCode.makeNotebookEditor("/test/notebook_mo.py", { + data: { + cells: hideCode.map((hide_code, index) => ({ + kind: 2, + value: "", + languageId: "python", + metadata: { stableId: `cell-${index}`, options: { hide_code } }, + })), + }, + }); + const vscode = yield* TestVsCode.make({ + initialDocuments: [editor.notebook], + }); + const layer = HideCodeSyncLive.pipe(Layer.provide(vscode.layer)); + return { vscode, editor, layer }; +}); + +describe("HideCodeSync", () => { + it.scoped( + "collapses hide_code cells when a notebook first becomes active", + Effect.fn(function* () { + const { vscode, editor, layer } = yield* withTestCtx([false, true, true]); + + yield* Effect.gen(function* () { + yield* vscode.setActiveNotebookEditor(Option.some(editor)); + yield* TestClock.adjust("1 millis"); + + expect(yield* collapses(vscode)).toEqual([ + [ + { start: 1, end: 2 }, + { start: 2, end: 3 }, + ], + ]); + }).pipe(Effect.provide(layer)); + }), + ); + + it.scoped( + "collapses once and does not re-collapse on tab refocus", + Effect.fn(function* () { + const { vscode, editor, layer } = yield* withTestCtx([true]); + + yield* Effect.gen(function* () { + yield* vscode.setActiveNotebookEditor(Option.some(editor)); + yield* TestClock.adjust("1 millis"); + yield* vscode.setActiveNotebookEditor(Option.none()); + yield* TestClock.adjust("1 millis"); + yield* vscode.setActiveNotebookEditor(Option.some(editor)); + yield* TestClock.adjust("1 millis"); + + expect(yield* collapses(vscode)).toEqual([[{ start: 0, end: 1 }]]); + }).pipe(Effect.provide(layer)); + }), + ); +}); diff --git a/extension/src/schemas/MarimoNotebookDocument.ts b/extension/src/schemas/MarimoNotebookDocument.ts index 1cc33b13..d56d4326 100644 --- a/extension/src/schemas/MarimoNotebookDocument.ts +++ b/extension/src/schemas/MarimoNotebookDocument.ts @@ -130,6 +130,22 @@ export class MarimoNotebookCell { ); } + /** + * Whether the cell's code is hidden via `@app.cell(hide_code=True)`. + * + * marimo stores the decorator's `hide_code` flag in the cell's config; the + * LSP deserialize path surfaces it as `metadata.options.hide_code`. VS Code + * has no API to read or set the input-collapsed state, so this is the source + * of truth the {@link HideCodeSyncLive} feature uses to one-way sync the + * editor's collapse state (issue #326). + */ + get isCodeHidden() { + return this.metadata.pipe( + Option.map((meta) => meta.options?.hide_code === true), + Option.getOrElse(() => false), + ); + } + /** * The cell's language metadata, if present. */