feat: streamline Python file processing by removing unused helper functions and optimizing import extraction

0-don · 0-don · commit 228bf3875119 · 2025-12-13T00:55:07.000+01:00
diff --git a/package.json b/package.json
@@ -2,7 +2,7 @@
   "name": "code-collector",
   "displayName": "Code Collector",
   "description": "Gather code context for AI based on imports",
-  "version": "0.0.61",
+  "version": "0.0.62",
   "publisher": "0-don",
   "icon": "assets/icon.png",
   "repository": {
diff --git a/src/commands.ts b/src/commands.ts
@@ -252,14 +252,13 @@ export class CommandHandler {
 
     const allContexts: FileContext[] = [];
     const processed = new Set<string>();
-    const pythonFiles = new Set<string>();
     const workspaceRoot = getWorkspaceRoot();
 
     this.output.log(
       `Processing ${filesToProcess.length} initial files for imports...`,
     );
 
-    for (const filePath of filesToProcess.filter((f) => !f.endsWith(".py"))) {
+    for (const filePath of filesToProcess) {
       if (token?.isCancellationRequested) {
         return;
       }
@@ -268,23 +267,9 @@ export class CommandHandler {
         allContexts,
         processed,
         workspaceRoot,
-        pythonFiles,
       );
     }
 
-    filesToProcess
-      .filter((f) => f.endsWith(".py"))
-      .forEach((f) => pythonFiles.add(f));
-    if (token?.isCancellationRequested) {
-      return;
-    }
-
-    await this.contextCollector.processPythonFiles(
-      pythonFiles,
-      allContexts,
-      processed,
-      workspaceRoot,
-    );
     if (token?.isCancellationRequested) {
       return;
     }
diff --git a/src/lib/core.ts b/src/lib/core.ts
@@ -1,10 +1,8 @@
-// src/lib/core.ts
 import * as fs from "fs";
 import * as micromatch from "micromatch";
 import * as path from "path";
 import { parserRegistry } from "../parsers";
 import { resolverRegistry } from "../resolvers";
-import { PythonResolver } from "../resolvers/python-resolver";
 import { FileContext } from "../types";
 import { findProjectRoot, isTextFile } from "../utils";
 import { getIgnorePatterns } from "./config";
@@ -22,15 +20,13 @@ export class ContextCollector {
 
     this.output.log(`Using ${ignorePatterns.length} ignore patterns`);
 
-    // Recursively discover files while respecting ignore patterns
     const filteredFiles = await this.discoverFiles(
       workspaceRoot,
       workspaceRoot,
       ignorePatterns,
     );
     this.output.log(`Discovered ${filteredFiles.length} files after filtering`);
 
-    // Process files and create contexts
     const contexts: FileContext[] = [];
     for (let i = 0; i < filteredFiles.length; i++) {
       if (progressCallback && !progressCallback(i + 1, filteredFiles.length)) {
@@ -68,7 +64,6 @@ export class ContextCollector {
         const relativePath = path.relative(workspaceRoot, fullPath);
 
         if (entry.isDirectory()) {
-          // For directories, check ignore patterns against both full path and directory name
           const directoryName = entry.name;
           const isIgnored =
             micromatch.isMatch(relativePath, ignorePatterns, { dot: true }) ||
@@ -78,7 +73,6 @@ export class ContextCollector {
             micromatch.isMatch(directoryName, ignorePatterns, { dot: true });
 
           if (!isIgnored) {
-            // Only recurse into directories that aren't ignored
             const subFiles = await this.discoverFiles(
               fullPath,
               workspaceRoot,
@@ -87,7 +81,6 @@ export class ContextCollector {
             files.push(...subFiles);
           }
         } else if (entry.isFile()) {
-          // For files, check ignore patterns against both full relative path and just filename
           const filename = path.basename(fullPath);
           const isIgnored =
             micromatch.isMatch(relativePath, ignorePatterns, { dot: true }) ||
@@ -120,20 +113,13 @@ export class ContextCollector {
     contexts: FileContext[],
     processed: Set<string>,
     workspaceRoot: string,
-    pythonFiles: Set<string>,
   ): Promise<void> {
     const normalizedPath = path.resolve(filePath);
 
     if (processed.has(normalizedPath) || !fs.existsSync(normalizedPath)) {
       return;
     }
 
-    // Collect Python files for batch processing
-    if (filePath.endsWith(".py")) {
-      pythonFiles.add(normalizedPath);
-      return;
-    }
-
     processed.add(normalizedPath);
 
     try {
@@ -150,7 +136,6 @@ export class ContextCollector {
           this.output.log(`${relativePath}: ${imports.length} imports`);
         }
 
-        // Use project-specific root for import resolution instead of workspace root
         const projectRoot = this.getProjectRootForFile(normalizedPath);
 
         for (const importInfo of imports) {
@@ -166,7 +151,6 @@ export class ContextCollector {
               contexts,
               processed,
               workspaceRoot,
-              pythonFiles,
             );
           }
         }
@@ -175,74 +159,4 @@ export class ContextCollector {
       this.output.error(`Failed to process: ${normalizedPath}`, error);
     }
   }
-
-  async processPythonFiles(
-    pythonFiles: Set<string>,
-    contexts: FileContext[],
-    processed: Set<string>,
-    workspaceRoot: string,
-  ): Promise<void> {
-    if (pythonFiles.size === 0) {
-      return;
-    }
-
-    this.output.log(
-      `Processing ${pythonFiles.size} Python files with helper...`,
-    );
-
-    const resolver = resolverRegistry.getResolver("dummy.py") as PythonResolver;
-    const ignorePatterns = getIgnorePatterns();
-
-    try {
-      const allPythonFiles = await resolver.resolveAllImports(
-        Array.from(pythonFiles),
-        ignorePatterns,
-      );
-      this.output.log(
-        `Python helper found ${allPythonFiles.length} total Python files`,
-      );
-
-      for (const pythonFile of allPythonFiles) {
-        const normalizedPath = path.resolve(pythonFile);
-
-        if (!processed.has(normalizedPath) && fs.existsSync(normalizedPath)) {
-          processed.add(normalizedPath);
-
-          try {
-            const content = fs.readFileSync(normalizedPath, "utf8");
-            const relativePath = path.relative(workspaceRoot, normalizedPath);
-            contexts.push({ path: normalizedPath, content, relativePath });
-          } catch (error) {
-            this.output.error(
-              `Failed to read Python file: ${normalizedPath}`,
-              error,
-            );
-          }
-        }
-      }
-    } catch (error) {
-      this.output.error(
-        `Python helper failed, processing files individually`,
-        error,
-      );
-
-      // Fallback: add Python files without import resolution
-      for (const pythonFile of pythonFiles) {
-        if (!processed.has(pythonFile)) {
-          processed.add(pythonFile);
-
-          try {
-            const content = fs.readFileSync(pythonFile, "utf8");
-            const relativePath = path.relative(workspaceRoot, pythonFile);
-            contexts.push({ path: pythonFile, content, relativePath });
-          } catch (error) {
-            this.output.error(
-              `Failed to read Python file: ${pythonFile}`,
-              error,
-            );
-          }
-        }
-      }
-    }
-  }
 }
diff --git a/src/parsers/python-parser.ts b/src/parsers/python-parser.ts
@@ -1,25 +1,103 @@
+import { execFile } from "child_process";
+import * as fs from "fs";
+import * as os from "os";
+import * as path from "path";
+import * as util from "util";
 import { ImportInfo, ParserConfig } from "../types";
 import { BaseParser } from "./base-parser";
 
+const execFileP = util.promisify(execFile);
+
+const HELPER_PYTHON_CONTENT = `import ast
+import json
+import sys
+import os
+
+def extract_imports_from_file(file_path):
+    """Extract all imports from a Python file"""
+    if not os.path.isfile(file_path) or not file_path.endswith('.py'):
+        return []
+    
+    imports = []
+    
+    try:
+        with open(file_path, 'r', encoding='utf-8') as f:
+            tree = ast.parse(f.read())
+        
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    imports.append({
+                        'module': alias.name,
+                        'type': 'import',
+                        'line': node.lineno
+                    })
+            
+            elif isinstance(node, ast.ImportFrom):
+                if node.module:
+                    imports.append({
+                        'module': node.module,
+                        'type': 'from',
+                        'line': node.lineno,
+                        'level': node.level
+                    })
+    
+    except (SyntaxError, UnicodeDecodeError):
+        pass
+    
+    return imports
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(json.dumps([]))
+        sys.exit(0)
+    
+    try:
+        file_path = sys.argv[1]
+        imports = extract_imports_from_file(file_path)
+        print(json.dumps(imports))
+    except Exception:
+        print(json.dumps([]))
+`;
+
 export class PythonParser extends BaseParser {
   config: ParserConfig = {
     extensions: [".py"],
     name: "Python",
   };
 
+  private helperPath: string | null = null;
+
   async parseImports(content: string, filePath: string): Promise<ImportInfo[]> {
+    const helperPath = await this.ensureHelperExists();
+    if (!helperPath) {
+      return this.parseImportsFallback(content);
+    }
+
+    try {
+      const { stdout } = await execFileP("python", [helperPath, filePath]);
+      const result = JSON.parse(stdout) || [];
+      return result.map((imp: any) => ({
+        module: imp.module,
+        type: imp.type,
+        line: imp.line,
+      }));
+    } catch (error) {
+      return this.parseImportsFallback(content);
+    }
+  }
+
+  private parseImportsFallback(content: string): ImportInfo[] {
     const imports: ImportInfo[] = [];
     const lines = content.split("\n");
 
     for (let i = 0; i < lines.length; i++) {
       const line = lines[i].trim();
 
-      // Skip empty lines and comments
       if (!line || line.startsWith("#")) {
         continue;
       }
 
-      // Standard import: import module
       const importMatch = line.match(/^import\s+([^\s#]+)/);
       if (importMatch) {
         imports.push({
@@ -30,7 +108,6 @@ export class PythonParser extends BaseParser {
         continue;
       }
 
-      // From import: from module import ...
       const fromMatch = line.match(/^from\s+([^\s#]+)\s+import/);
       if (fromMatch) {
         imports.push({
@@ -41,7 +118,6 @@ export class PythonParser extends BaseParser {
         continue;
       }
 
-      // Stop at first non-import statement (basic version)
       if (
         line &&
         !line.startsWith("from ") &&
@@ -57,4 +133,19 @@ export class PythonParser extends BaseParser {
 
     return imports;
   }
+
+  private async ensureHelperExists(): Promise<string | null> {
+    if (this.helperPath && fs.existsSync(this.helperPath)) {
+      return this.helperPath;
+    }
+
+    try {
+      const tempDir = os.tmpdir();
+      this.helperPath = path.join(tempDir, "code-collector-python-helper.py");
+      fs.writeFileSync(this.helperPath, HELPER_PYTHON_CONTENT, "utf8");
+      return this.helperPath;
+    } catch (error) {
+      return null;
+    }
+  }
 }
diff --git a/src/resolvers/python-resolver.ts b/src/resolvers/python-resolver.ts
