diff --git a/backend/parsers/windmill-parser-py-imports/src/lib.rs b/backend/parsers/windmill-parser-py-imports/src/lib.rs
index 92f4346115..4e51be6b6a 100644
--- a/backend/parsers/windmill-parser-py-imports/src/lib.rs
+++ b/backend/parsers/windmill-parser-py-imports/src/lib.rs
@@ -160,10 +160,17 @@ fn parse_code_for_imports(code: &str, path: &str) -> error::Result<Vec<NImport>>
             + "\n";
     }
 
-    let ast = Suite::parse(&code, "main.py").map_err(|e| {
+    // Add a fake main function to ensure the parser can process the code correctly
+    // This is needed because we've split off the real main function above
+    let code_with_fake_main = format!("{}\n\ndef main(): pass", code);
+
+    let ast = Suite::parse(&code_with_fake_main, "main.py").map_err(|e| {
         error::Error::ExecutionErr(format!("Error parsing code for imports: {}", e.to_string()))
     })?;
 
+    // Note: We're still using the original code for finding pins,
+    // as the TextRange values from the parsed AST would be based on code_with_fake_main
+    // but we want to match against the original code
     let find_pin = |range: TextRange, key: String| {
         let hs = code
             .chars()