Add support for namespace package

Author: savannahostrowski

src/core/extractors.ts

@@ -190,9 +190,10 @@ export function routerExtractor(node: Node): RouterInfo | null {
 
   if (valueNode.type === "call") {
     const functionNameNode = valueNode.childForFieldName("function")
-    if (functionNameNode && functionNameNode.text === "APIRouter") {
+    const funcName = functionNameNode?.text
+    if (funcName === "APIRouter" || funcName === "fastapi.APIRouter") {
       type = "APIRouter"
-    } else if (functionNameNode && functionNameNode.text === "FastAPI") {
+    } else if (funcName === "FastAPI" || funcName === "fastapi.FastAPI") {
       type = "FastAPI"
     }
 

src/core/importResolver.ts

@@ -1,20 +1,56 @@
+/**
+ * Python import resolution for static analysis.
+ *
+ * Resolves Python import statements to file paths without executing code.
+ * Handles relative imports, absolute imports, namespace packages (PEP 420),
+ * and re-exports from __init__.py files.
+ *
+ * Resolution order (matching Python):
+ * 1. module.py (direct file)
+ * 2. module/__init__.py (package)
+ * 3. module/ without __init__.py (namespace package)
+ */
+
 import { existsSync } from "node:fs"
-import { dirname } from "node:path"
+import { dirname, join } from "node:path"
 import { analyzeFile } from "./analyzer"
 import type { ImportInfo } from "./internal"
 import type { Parser } from "./parser"
 
+/**
+ * Cache for file existence checks to avoid repeated filesystem calls.
+ * Maps file path -> exists (true/false).
+ */
+const fileExistsCache = new Map<string, boolean>()
+
+function cachedExistsSync(path: string): boolean {
+  const cached = fileExistsCache.get(path)
+  if (cached !== undefined) {
+    return cached
+  }
+  const exists = existsSync(path)
+  fileExistsCache.set(path, exists)
+  return exists
+}
+
+/** Clears the file existence cache. Call when files may have changed. */
+export function clearImportCache(): void {
+  fileExistsCache.clear()
+}
+
 /**
  * Resolves a module path to its Python file.
  * Checks for direct .py file first, then package __init__.py
  * (matching Python's import resolution order).
  */
 function resolvePythonModule(basePath: string): string | null {
-  if (existsSync(`${basePath}.py`)) {
-    return `${basePath}.py`
+  const pyPath = `${basePath}.py`
+  if (cachedExistsSync(pyPath)) {
+    return pyPath
   }
-  if (existsSync(`${basePath}/__init__.py`)) {
-    return `${basePath}/__init__.py`
+  const initPath = join(basePath, "__init__.py")
+  if (cachedExistsSync(initPath)) {
+    return initPath
   }
   return null
 }
@@ -41,30 +77,50 @@ function findImportByExportedName(
 }
 
 /**
- * Base resolution of a module import to its file path.
- * Handles both relative and absolute imports.
+ * Converts a Python module path to a filesystem directory path.
+ *
+ * Examples (modulePath, relativeDots → result):
+ *   Absolute: ("app.api.routes", 0) from projectRoot="/project" → "/project/app/api/routes"
+ *   Relative: ("routes", 1) from "/project/app/api/main.py" → "/project/app/api/routes"
+ *   Relative: ("routes", 2) from "/project/app/api/main.py" → "/project/app/routes"
  */
-export function resolveImport(
+function modulePathToDir(
   importInfo: Pick<ImportInfo, "modulePath" | "isRelative" | "relativeDots">,
   currentFilePath: string,
   projectRoot: string,
-): string | null {
-  let resolvedPath: string
-
+): string {
+  let baseDir: string
   if (importInfo.isRelative) {
     // For relative imports, go up 'relativeDots' directories from current file
-    let currentDir = dirname(currentFilePath)
+    baseDir = dirname(currentFilePath)
     for (let i = 1; i < importInfo.relativeDots; i++) {
-      currentDir = dirname(currentDir)
+      baseDir = dirname(baseDir)
     }
-    resolvedPath = importInfo.modulePath
-      ? `${currentDir}/${importInfo.modulePath.replace(/\./g, "/")}`
-      : currentDir
-    // Absolute import
   } else {
-    resolvedPath = `${projectRoot}/${importInfo.modulePath.replace(/\./g, "/")}`
+    baseDir = projectRoot
+  }
+
+  if (importInfo.modulePath) {
+    return join(baseDir, ...importInfo.modulePath.split("."))
   }
+  return baseDir
+}
 
+/**
+ * Resolves a module import to its file path.
+ *
+ * Examples:
+ *   "from app.api import routes" → "/project/app/api/routes.py" or "/project/app/api/routes/__init__.py"
+ *   "from .routes import users" → "/project/app/api/routes.py" or "/project/app/api/routes/__init__.py"
+ *
+ * Returns null if the module doesn't exist (may be a namespace package).
+ */
+export function resolveImport(
+  importInfo: Pick<ImportInfo, "modulePath" | "isRelative" | "relativeDots">,
+  currentFilePath: string,
+  projectRoot: string,
+): string | null {
+  const resolvedPath = modulePathToDir(importInfo, currentFilePath, projectRoot)
   return resolvePythonModule(resolvedPath)
 }
 
@@ -83,22 +139,24 @@ export function resolveNamedImport(
   parser?: Parser,
 ): string | null {
   const basePath = resolveImport(importInfo, currentFilePath, projectRoot)
-  if (!basePath) {
-    return null
-  }
 
-  const baseDir = dirname(basePath)
+  // Calculate base directory for named import resolution.
+  // For namespace packages (directories without __init__.py), basePath will be null,
+  // so we compute the directory path directly from the module path.
+  const baseDir = basePath
+    ? dirname(basePath)
+    : modulePathToDir(importInfo, currentFilePath, projectRoot)
 
   for (const name of importInfo.names) {
     // Try direct file: from .routes import users -> routes/users.py
-    const namedPath = `${baseDir}/${name.replace(/\./g, "/")}`
+    const namedPath = join(baseDir, ...name.split("."))
     const resolved = resolvePythonModule(namedPath)
     if (resolved) {
       return resolved
     }
 
     // Try re-exports: from .routes import users where routes/__init__.py re-exports users
-    if (basePath.endsWith("__init__.py") && parser) {
+    if (basePath?.endsWith("__init__.py") && parser) {
       const analysis = analyzeFile(basePath, parser)
       const imp = analysis && findImportByExportedName(analysis.imports, name)
       if (imp) {

src/core/pathUtils.ts

@@ -1,3 +1,6 @@
+import { existsSync } from "node:fs"
+import { dirname, join, relative, sep } from "node:path"
+
 /**
  * Strips leading dynamic segments (like {settings.API_V1_STR}) from a path.
  * These are runtime variables, not URL path parameters.
@@ -12,6 +15,44 @@ export function stripLeadingDynamicSegments(path: string): string {
   return path.replace(/^(\{[^}]+\})+/, "") || "/"
 }
 
+/**
+ * Checks if a path is within or equal to a base directory.
+ * Uses relative path calculation to avoid false positives from string prefix matching.
+ */
+export function isWithinDirectory(filePath: string, baseDir: string): boolean {
+  const rel = relative(baseDir, filePath)
+  // If relative path starts with "..", the path is outside baseDir
+  return !rel.startsWith("..") && !rel.startsWith(sep)
+}
+
+/**
+ * Finds the Python project root by walking up from the entry file
+ * until we find a directory without __init__.py (or hit the workspace root).
+ * This is the directory from which absolute imports are resolved.
+ */
+export function findProjectRoot(
+  entryPath: string,
+  workspaceRoot: string,
+): string {
+  let dir = dirname(entryPath)
+
+  // If the entry file's directory doesn't have __init__.py, it's a top-level script
+  if (!existsSync(join(dir, "__init__.py"))) {
+    return dir
+  }
+
+  // Walk up until we find a directory whose parent doesn't have __init__.py
+  while (isWithinDirectory(dir, workspaceRoot) && dir !== workspaceRoot) {
+    const parent = dirname(dir)
+    if (!existsSync(join(parent, "__init__.py"))) {
+      return parent
+    }
+    dir = parent
+  }
+
+  return workspaceRoot
+}
+
 /**
  * Gets the first N segments of a path.
  *

src/extension.ts

@@ -2,9 +2,11 @@
  * VSCode extension entry point for FastAPI endpoint discovery.
  */
 
+import { sep } from "node:path"
 import * as vscode from "vscode"
+import { clearImportCache } from "./core/importResolver"
 import { Parser } from "./core/parser"
-import { stripLeadingDynamicSegments } from "./core/pathUtils"
+import { findProjectRoot, stripLeadingDynamicSegments } from "./core/pathUtils"
 import { buildRouterGraph } from "./core/routerResolver"
 import { routerNodeToAppDefinition } from "./core/transformer"
 import type { AppDefinition, SourceLocation } from "./core/types"
@@ -21,25 +23,38 @@ async function discoverFastAPIApps(parser: Parser): Promise<AppDefinition[]> {
     return apps
   }
 
-  const defaultPatterns = [
-    "main.py",
-    "app/main.py",
-    "api/main.py",
-    "src/main.py",
-    "backend/app/main.py",
-  ]
-
   for (const folder of workspaceFolders) {
     const config = vscode.workspace.getConfiguration("fastapi", folder.uri)
     const customEntryPoint = config.get<string>("entryPoint")
-    const patterns = customEntryPoint ? [customEntryPoint] : defaultPatterns
-
-    for (const pattern of patterns) {
-      // Handle both relative patterns and absolute paths
-      const entryPath = pattern.startsWith("/")
-        ? pattern
-        : vscode.Uri.joinPath(folder.uri, pattern).fsPath
-      const projectRoot = entryPath.split("/").slice(0, -2).join("/")
+
+    let candidates: string[] = []
+
+    if (customEntryPoint) {
+      // Use custom entry point if specified
+      const entryPath = customEntryPoint.startsWith("/")
+        ? customEntryPoint
+        : vscode.Uri.joinPath(folder.uri, customEntryPoint).fsPath
+      candidates = [entryPath]
+    } else {
+      // Scan for main.py and __init__.py files (likely FastAPI entry points)
+      const mainFiles = await vscode.workspace.findFiles(
+        new vscode.RelativePattern(folder, "**/main.py"),
+        undefined,
+        20,
+      )
+      const initFiles = await vscode.workspace.findFiles(
+        new vscode.RelativePattern(folder, "**/__init__.py"),
+        undefined,
+        20,
+      )
+      // Prefer main.py, then __init__.py, sorted by path depth (shallower first)
+      candidates = [...mainFiles, ...initFiles]
+        .map((uri) => uri.fsPath)
+        .sort((a, b) => a.split(sep).length - b.split(sep).length)
+    }
+
+    for (const entryPath of candidates) {
+      const projectRoot = findProjectRoot(entryPath, folder.uri.fsPath)
       const routerNode = buildRouterGraph(entryPath, parser, projectRoot)
 
       if (routerNode) {
@@ -95,6 +110,7 @@ export async function activate(context: vscode.ExtensionContext) {
     vscode.commands.registerCommand(
       "fastapi-vscode.refreshEndpoints",
       async () => {
+        clearImportCache()
         const newApps = await discoverFastAPIApps(parserService)
         endpointProvider.setApps(newApps)
       },

src/test/extractors.test.ts

@@ -279,6 +279,29 @@ def handler():
 
       assert.strictEqual(result, null)
     })
+
+    test("extracts qualified fastapi.FastAPI() call", () => {
+      const code = "app = fastapi.FastAPI()"
+      const tree = parse(code)
+      const assignments = findNodesByType(tree.rootNode, "assignment")
+      const result = routerExtractor(assignments[0])
+
+      assert.ok(result)
+      assert.strictEqual(result.variableName, "app")
+      assert.strictEqual(result.type, "FastAPI")
+    })
+
+    test("extracts qualified fastapi.APIRouter() call", () => {
+      const code = "router = fastapi.APIRouter(prefix='/api')"
+      const tree = parse(code)
+      const assignments = findNodesByType(tree.rootNode, "assignment")
+      const result = routerExtractor(assignments[0])
+
+      assert.ok(result)
+      assert.strictEqual(result.variableName, "router")
+      assert.strictEqual(result.type, "APIRouter")
+      assert.strictEqual(result.prefix, "/api")
+    })
   })
 
   suite("importExtractor", () => {

src/test/fixtures/python/app/api/namespace_routes/api_routes.py

@@ -0,0 +1,7 @@
+from fastapi import APIRouter
+
+router = APIRouter()
+
+@router.get("/items")
+def list_items():
+    pass

src/test/importResolver.test.ts

@@ -183,5 +183,47 @@ suite("importResolver", () => {
         result.endsWith("routes/__init__.py") || result.endsWith("routes.py"),
       )
     })
+
+    test("resolves relative named import from namespace package (no __init__.py)", () => {
+      const currentFile = path.join(fixturesPath, "app", "api", "main.py")
+      const projectRoot = fixturesPath
+
+      // namespace_routes has no __init__.py, but api_routes.py exists
+      const result = resolveNamedImport(
+        {
+          modulePath: "namespace_routes",
+          names: ["api_routes"],
+          isRelative: true,
+          relativeDots: 1,
+        },
+        currentFile,
+        projectRoot,
+        parser,
+      )
+
+      assert.ok(result)
+      assert.ok(result.endsWith("api_routes.py"))
+    })
+
+    test("resolves absolute named import from namespace package (no __init__.py)", () => {
+      const currentFile = path.join(fixturesPath, "main.py")
+      const projectRoot = fixturesPath
+
+      // app.api.namespace_routes has no __init__.py, but api_routes.py exists
+      const result = resolveNamedImport(
+        {
+          modulePath: "app.api.namespace_routes",
+          names: ["api_routes"],
+          isRelative: false,
+          relativeDots: 0,
+        },
+        currentFile,
+        projectRoot,
+        parser,
+      )
+
+      assert.ok(result)
+      assert.ok(result.endsWith("api_routes.py"))
+    })
   })
 })