Skip 'if __name__=='__main__' blocks in snapshots

vscode_extensions/run-until-cursor/src/extension.ts

@@ -51,14 +51,14 @@ export function activate(context: vscode.ExtensionContext) {
             if (cursorLine >= startLine && cursorLine < endLine) {
                 executeRange = new vscode.Range(startLine, 0, cursorLine + 1, 0);
             }
-            // If outside `run()` (either above or below), execute **the entire script**
+            // If outside `run()` (either above or below), execute **the entire script** (without `if __name__ == "__main__":`)
             else {
-                executeRange = new vscode.Range(0, 0, document.lineCount, 0);
+                executeRange = filterOutMainBlock(document);
                 shouldMoveCursor = true;
             }
         } else {
-            // If `run()` is not found, execute everything as a fallback
-            executeRange = new vscode.Range(0, 0, document.lineCount, 0);
+            // If `run()` is not found, execute everything (without `if __name__ == "__main__":`)
+            executeRange = filterOutMainBlock(document);
         }
 
         // Step 3: Apply selection and execute immediately (with an adjustable delay)
@@ -82,3 +82,47 @@ export function activate(context: vscode.ExtensionContext) {
 }
 
 export function deactivate() {}
+
+/**
+ * Filters out `if __name__ == "__main__":` and its indented block.
+ */
+function filterOutMainBlock(document: vscode.TextDocument): vscode.Range {
+    let mainStart: number | null = null;
+    let mainEnd: number | null = null;
+
+    // Flexible regex for detecting variations of `if __name__ == "__main__":`
+    const mainRegex = /^\s*if\s*["']__main__["']\s*==\s*__name__\s*:|^\s*if\s*__name__\s*==\s*["']__main__["']\s*:/;
+
+    // Step 1: Find `if __name__ == "__main__":`
+    for (let i = 0; i < document.lineCount; i++) {
+        if (mainRegex.test(document.lineAt(i).text)) {
+            mainStart = i;
+            break;
+        }
+    }
+
+    if (mainStart !== null) {
+        // Step 2: Find where `if __name__ == "__main__":` block ends (first non-indented line)
+        for (let i = mainStart + 1; i < document.lineCount; i++) {
+            const lineText = document.lineAt(i).text;
+
+            // A **non-indented line** (not empty & not a comment) signals the end of the block
+            if (lineText.trim().length > 0 && !lineText.startsWith(" ") && !lineText.startsWith("\t")) {
+                mainEnd = i;
+                break;
+            }
+        }
+
+        // If no non-indented line was found, assume the block goes to the last line
+        if (mainEnd === null) {
+            mainEnd = document.lineCount;
+        }
+    }
+
+    // If we found `if __name__ == "__main__":`, remove that block
+    if (mainStart !== null && mainEnd !== null) {
+        return new vscode.Range(0, 0, mainStart, 0); // Select everything **before** `if __name__ == "__main__"`
+    } else {
+        return new vscode.Range(0, 0, document.lineCount, 0); // No `main` block, execute everything
+    }
+}