Init Reload & Preview Command/CodeLens

package.json

@@ -37,6 +37,11 @@
         "title": "Preview active Manim Cell",
         "category": "Manim Notebook"
       },
+      {
+        "command": "manim-notebook.reloadAndPreviewManimCell",
+        "title": "Reload and Preview active Manim Cell",
+        "category": "Manim Notebook"
+      },
       {
         "command": "manim-notebook.previewSelection",
         "title": "Preview selected Manim code",

src/extension.ts

@@ -2,8 +2,8 @@ import * as vscode from 'vscode';
 import { window } from 'vscode';
 import { ManimShell, NoActiveShellError } from './manimShell';
 import { ManimCell } from './manimCell';
+import { previewManimCell, reloadAndPreviewManimCell, previewCode } from './previewCode';
 import { ManimCellRanges } from './pythonParsing';
-import { previewCode } from './previewCode';
 import { startScene, exitScene } from './startStopScene';
 import { exportScene } from './export';
 import { Logger, Window, LogRecorder } from './logger';
@@ -20,6 +20,14 @@ export function activate(context: vscode.ExtensionContext) {
 			previewManimCell(cellCode, startLine);
 		});
 
+	const reloadAndPreviewManimCellCommand = vscode.commands.registerCommand(
+		'manim-notebook.reloadAndPreviewManimCell',
+		(cellCode?: string, startLine?: number) => {
+			Logger.info("💠 Command requested: Reload & Preview Manim Cell"
+				+ `, startLine=${startLine}`);
+			reloadAndPreviewManimCell(cellCode, startLine);
+		});
+
 	const previewSelectionCommand = vscode.commands.registerCommand(
 		'manim-notebook.previewSelection', () => {
 			Logger.info("💠 Command requested: Preview Selection");
@@ -99,51 +107,6 @@ export function deactivate() {
 	Logger.info("💠 Manim Notebook extension deactivated");
 }
 
-/**
- * Previews all code inside of a Manim cell.
- * 
- * A Manim cell starts with ##
- * 
- * This can be invoked by either:
- * - clicking the code lens (the button above the cell) -> this cell is previewed
- * - command pallette -> the 1 cell where the cursor is is previewed
- * 
- * If Manim isn't running, it will be automatically started
- * (at the start of the cell which will be previewed: on its starting ## line),
- * and then this cell is previewed.
- */
-async function previewManimCell(cellCode?: string, startLine?: number) {
-	let startLineFinal: number | undefined = startLine;
-
-	// User has executed the command via command pallette
-	if (cellCode === undefined) {
-		const editor = window.activeTextEditor;
-		if (!editor) {
-			Window.showErrorMessage(
-				'No opened file found. Place your cursor in a Manim cell.');
-			return;
-		}
-		const document = editor.document;
-
-		// Get the code of the cell where the cursor is placed
-		const cursorLine = editor.selection.active.line;
-		const range = ManimCellRanges.getCellRangeAtLine(document, cursorLine);
-		if (!range) {
-			Window.showErrorMessage('Place your cursor in a Manim cell.');
-			return;
-		}
-		cellCode = document.getText(range);
-		startLineFinal = range.start.line;
-	}
-
-	if (startLineFinal === undefined) {
-		Window.showErrorMessage('Internal error: Line number not found in `previewManimCell()`.');
-		return;
-	}
-
-	await previewCode(cellCode, startLineFinal);
-}
-
 /**
  * Previews the selected code.
  * 

src/manimCell.ts

@@ -35,35 +35,42 @@ export class ManimCell implements vscode.CodeLensProvider, vscode.FoldingRangePr
     }
 
     public provideCodeLenses(document: vscode.TextDocument, token: vscode.CancellationToken): vscode.CodeLens[] {
+        if (!window.activeTextEditor) {
+            return [];
+        }
+
         const codeLenses: vscode.CodeLens[] = [];
 
         const ranges = ManimCellRanges.calculateRanges(document);
-        for (const range of ranges) {
-            codeLenses.push(new vscode.CodeLens(range));
+        for (let range of ranges) {
+            range = new vscode.Range(range.start, range.end);
+            const codeLens = new vscode.CodeLens(range);
+            const codeLensReload = new vscode.CodeLens(range);
+
+            const document = window.activeTextEditor.document;
+            const cellCode = document.getText(range);
+
+            codeLens.command = {
+                title: "Preview Manim Cell",
+                command: "manim-notebook.previewManimCell",
+                tooltip: "Preview this Manim Cell",
+                arguments: [cellCode, range.start.line]
+            };
+
+            codeLensReload.command = {
+                title: "Reload & Preview",
+                command: "manim-notebook.reloadAndPreviewManimCell",
+                tooltip: "Reload & Preview this Manim Cell",
+                arguments: [cellCode, range.start.line]
+            };
+
+            codeLenses.push(codeLens);
+            codeLenses.push(codeLensReload);
         }
 
         return codeLenses;
     }
 
-    public resolveCodeLens(codeLens: vscode.CodeLens, token: vscode.CancellationToken): vscode.CodeLens {
-        if (!window.activeTextEditor) {
-            return codeLens;
-        }
-
-        const document = window.activeTextEditor?.document;
-        const range = new vscode.Range(codeLens.range.start, codeLens.range.end);
-        const cellCode = document.getText(range);
-
-        codeLens.command = {
-            title: "▶ Preview Manim Cell",
-            command: "manim-notebook.previewManimCell",
-            tooltip: "Preview this Manim Cell inside an interactive Manim environment",
-            arguments: [cellCode, codeLens.range.start.line]
-        };
-
-        return codeLens;
-    }
-
     public provideFoldingRanges(document: vscode.TextDocument, context: vscode.FoldingContext, token: vscode.CancellationToken): vscode.FoldingRange[] {
         const ranges = ManimCellRanges.calculateRanges(document);
         return ranges.map(range => new vscode.FoldingRange(range.start.line, range.end.line));

src/manimShell.ts

@@ -152,6 +152,13 @@ export class ManimShell {
      */
     private iPythonCellCount: number = 0;
 
+    /**
+     * Whether to wait for a restarted IPython instance, i.e. for an IPython
+     * cell count of 1. This is set to `true` before the `reload()` command is
+     * issued and set back to `false` after the IPython cell count is 1.
+     */
+    waitForRestartedIPythonInstance = false;
+
     /**
      * Whether the execution of a new command is locked. This is used to prevent
      * multiple new scenes from being started at the same time, e.g. when users
@@ -199,6 +206,20 @@ export class ManimShell {
         return ManimShell.#instance;
     }
 
+    /**
+     * Indicates that the next command should wait until a restarted IPython
+     * instance is detected, i.e. starting with cell 1 again. This should be
+     * called before the `reload()` command is issued.
+     */
+    public async nextTimeWaitForRestartedIPythonInstance() {
+        if (await this.isLocked()) {
+            return;
+        }
+
+        this.iPythonCellCount = 0;
+        this.waitForRestartedIPythonInstance = true;
+    }
+
     /**
      * Executes the given command. If no active terminal running Manim is found,
      * a new terminal is spawned, and a new Manim session is started in it
@@ -231,6 +252,36 @@ export class ManimShell {
             command, waitUntilFinished, forceExecute, true, undefined, undefined);
     }
 
+    /**
+     * Returns whether the command execution is currently locked, i.e. when
+     * Manim is starting up or another command is currently running.
+     * 
+     * @param forceExecute see `execCommand()`
+     * @returns true if the command execution is locked, false otherwise.
+     */
+    private async isLocked(forceExecute = false): Promise<boolean> {
+        if (this.lockDuringStartup) {
+            Window.showWarningMessage("Manim is currently starting. Please wait a moment.");
+            return true;
+        }
+
+        if (this.isExecutingCommand) {
+            // MacOS specific behavior
+            if (this.shouldLockDuringCommandExecution && !forceExecute) {
+                Window.showWarningMessage(
+                    `Simultaneous Manim commands are not currently supported on MacOS. `
+                    + `Please wait for the current operations to finish before initiating `
+                    + `a new command.`);
+                return true;
+            }
+
+            this.sendKeyboardInterrupt();
+            await new Promise(resolve => setTimeout(resolve, 500));
+        }
+
+        return false;
+    }
+
     /**
      * Executes a given command and bundles many different behaviors and options.
      * 
@@ -276,27 +327,12 @@ export class ManimShell {
             return;
         }
 
-        if (this.lockDuringStartup) {
-            Window.showWarningMessage("Manim is currently starting. Please wait a moment.");
-            return;
-        }
-
         if (errorOnNoActiveShell) {
             this.errorOnNoActiveShell();
         }
 
-        if (this.isExecutingCommand) {
-            // MacOS specific behavior
-            if (this.shouldLockDuringCommandExecution && !forceExecute) {
-                Window.showWarningMessage(
-                    `Simultaneous Manim commands are not currently supported on MacOS. `
-                    + `Please wait for the current operations to finish before initiating `
-                    + `a new command.`);
-                return;
-            }
-
-            this.sendKeyboardInterrupt();
-            await new Promise(resolve => setTimeout(resolve, 500));
+        if (await this.isLocked(forceExecute)) {
+            return;
         }
 
         this.isExecutingCommand = true;
@@ -503,7 +539,7 @@ export class ManimShell {
      * A shell that was previously used to run Manim, but has exited from the
      * Manim session (IPython environment), is considered inactive.
      */
-    private hasActiveShell(): boolean {
+    public hasActiveShell(): boolean {
         const hasActiveShell =
             this.activeShell !== null && this.activeShell.exitStatus === undefined;
         Logger.debug(`👩‍💻 Has active shell?: ${hasActiveShell}`);
@@ -688,14 +724,27 @@ export class ManimShell {
 
                     let ipythonMatches = data.match(IPYTHON_CELL_START_REGEX);
                     if (ipythonMatches) {
-                        // Terminal data might include multiple IPython statements,
-                        // so take the highest cell number found.
+                        // Terminal data might include multiple IPython statements
                         const cellNumbers = ipythonMatches.map(
                             match => parseInt(match.match(/\d+/)![0]));
-                        const maxCellNumber = Math.max(...cellNumbers);
-                        this.iPythonCellCount = maxCellNumber;
-                        Logger.debug(`📦 IPython cell ${maxCellNumber} detected`);
-                        this.eventEmitter.emit(ManimShellEvent.IPYTHON_CELL_FINISHED);
+
+                        if (this.waitForRestartedIPythonInstance) {
+                            const cellNumber = Math.min(...cellNumbers);
+                            Logger.debug("📦 While waiting for restarted IPython instance:"
+                                + ` cell ${cellNumber} detected`);
+                            if (cellNumber === 1) {
+                                Logger.debug("🔄 Restarted IPython instance detected");
+                                this.iPythonCellCount = 1;
+                                this.waitForRestartedIPythonInstance = false;
+                                this.eventEmitter.emit(ManimShellEvent.IPYTHON_CELL_FINISHED);
+                            }
+                        } else {
+                            // more frequent case
+                            const cellNumber = Math.max(...cellNumbers);
+                            this.iPythonCellCount = cellNumber;
+                            Logger.debug(`📦 IPython cell ${cellNumber} detected`);
+                            this.eventEmitter.emit(ManimShellEvent.IPYTHON_CELL_FINISHED);
+                        }
                     }
 
                     if (this.isExecutingCommand && data.match(IPYTHON_MULTILINE_START_REGEX)) {

src/previewCode.ts

@@ -2,11 +2,84 @@ import * as vscode from 'vscode';
 import { window } from 'vscode';
 import { ManimShell } from './manimShell';
 import { EventEmitter } from 'events';
-import { Logger } from './logger';
+import { ManimCellRanges } from './pythonParsing';
+import { Logger, Window } from './logger';
 
 // \x0C: is Ctrl + L, which clears the terminal screen
 const PREVIEW_COMMAND = `\x0Ccheckpoint_paste()`;
 
+function parsePreviewCellArgs(cellCode?: string, startLine?: number) {
+    let startLineParsed: number | undefined = startLine;
+
+    // User has executed the command via command pallette
+    if (cellCode === undefined) {
+        const editor = window.activeTextEditor;
+        if (!editor) {
+            Window.showErrorMessage(
+                'No opened file found. Place your cursor in a Manim cell.');
+            return;
+        }
+        const document = editor.document;
+
+        // Get the code of the cell where the cursor is placed
+        const cursorLine = editor.selection.active.line;
+        const range = ManimCellRanges.getCellRangeAtLine(document, cursorLine);
+        if (!range) {
+            Window.showErrorMessage('Place your cursor in a Manim cell.');
+            return;
+        }
+        cellCode = document.getText(range);
+        startLineParsed = range.start.line;
+    }
+
+    if (startLineParsed === undefined) {
+        Window.showErrorMessage(
+            'Internal error: Line number not found in `parsePreviewCellArgs()`.');
+        return;
+    }
+
+    return { cellCodeParsed: cellCode, startLineParsed };
+}
+
+/**
+ * Previews all code inside of a Manim cell.
+ * 
+ * A Manim cell starts with `##`.
+ * 
+ * This can be invoked by either:
+ * - clicking the code lens (the button above the cell) -> this cell is previewed
+ * - command pallette -> the 1 cell where the cursor is is previewed
+ * 
+ * If Manim isn't running, it will be automatically started
+ * (at the start of the cell which will be previewed: on its starting ## line),
+ * and then this cell is previewed.
+ */
+export async function previewManimCell(cellCode?: string, startLine?: number) {
+    const res = parsePreviewCellArgs(cellCode, startLine);
+    if (!res) {
+        return;
+    }
+    const { cellCodeParsed, startLineParsed } = res;
+
+    await previewCode(cellCodeParsed, startLineParsed);
+}
+
+export async function reloadAndPreviewManimCell(cellCode?: string, startLine?: number) {
+    const res = parsePreviewCellArgs(cellCode, startLine);
+    if (!res) {
+        return;
+    }
+    const { cellCodeParsed, startLineParsed } = res;
+
+    if (ManimShell.instance.hasActiveShell()) {
+        const reloadCmd = `reload(${startLineParsed + 1})`;
+        await ManimShell.instance.nextTimeWaitForRestartedIPythonInstance();
+        await ManimShell.instance.executeCommandErrorOnNoActiveSession(reloadCmd, true);
+    }
+    await previewManimCell(cellCodeParsed, startLineParsed);
+}
+
+
 /**
  * Interactively previews the given Manim code by means of the
  * `checkpoint_paste()` method from Manim.