🐛 Fix discovery when multiple routers are imported from the same file (#129)

Author: savannahostrowski

src/core/routerResolver.ts

@@ -122,12 +122,13 @@ async function buildRouterGraphInternal(
   }
 
   // Prevent infinite recursion on circular imports
-  if (visited.has(entryFileUri)) {
-    log(`Skipping already visited file: "${entryFileUri}"`)
+  const visitedKey = `${entryFileUri}#${targetVariable ?? ""}`
+  if (visited.has(visitedKey)) {
+    log(`Skipping already visited: "${visitedKey}"`)
     return null
   }
 
-  visited.add(entryFileUri)
+  visited.add(visitedKey)
 
   // Helper to analyze a file with the filesystem
   const analyzeFileFn = (uri: string) => analyzeFile(uri, parser, fs)
@@ -361,11 +362,13 @@ async function resolveRouterReference(
       (r) => r.variableName === attributeName,
     )
     if (targetRouter) {
+      const visitedKey = `${importedFileUri}#${attributeName}`
+
       // Mark as visited to prevent infinite recursion
-      if (visited.has(importedFileUri)) {
+      if (visited.has(visitedKey)) {
         return null
       }
-      visited.add(importedFileUri)
+      visited.add(visitedKey)
 
       // Get routes belonging to this router
       const routerRoutes = importedAnalysis.routes.filter(
@@ -387,8 +390,9 @@ async function resolveRouterReference(
 
       return routerNode
     }
-    // If not found as a router, fall through to try building from file
   }
 
-  return buildRouterGraphInternal(importedFileUri, ctx)
+  // Resolve by variable name for named imports, or fall back to full-file discovery
+  const targetVarName = namedImport ? originalName : undefined
+  return buildRouterGraphInternal(importedFileUri, ctx, targetVarName)
 }

src/test/core/routerResolver.test.ts

@@ -678,5 +678,37 @@ suite("routerResolver", () => {
       assert.strictEqual(usersRouter.prefix, "/users")
       assert.strictEqual(usersRouter.routes.length, 2)
     })
+
+    test("resolves multiple routers imported from same file", async () => {
+      const result = await buildRouterGraph(
+        fixtures.multiRouterSameFile.mainPy,
+        parser,
+        fixtures.multiRouterSameFile.root,
+        nodeFileSystem,
+      )
+
+      assert.ok(result)
+      assert.strictEqual(result.type, "FastAPI")
+      assert.strictEqual(
+        result.children.length,
+        2,
+        "Should resolve both routers from same file",
+      )
+
+      const prefixes = result.children.map((c) => c.router.prefix)
+      assert.ok(
+        prefixes.includes("/v1"),
+        "Should include router1 with /v1 prefix",
+      )
+      assert.ok(
+        prefixes.includes("/v2"),
+        "Should include router2 with /v2 prefix",
+      )
+
+      for (const child of result.children) {
+        assert.strictEqual(child.router.routes.length, 1)
+        assert.strictEqual(child.router.routes[0].path, "/items")
+      }
+    })
   })
 })

src/test/fixtures/multi-router-same-file/main.py

@@ -0,0 +1,6 @@
+from fastapi import FastAPI
+from .routers import router1, router2
+
+app = FastAPI()
+app.include_router(router1)
+app.include_router(router2)

src/test/fixtures/multi-router-same-file/routers.py

@@ -0,0 +1,12 @@
+from fastapi import APIRouter
+
+router1 = APIRouter(prefix="/v1")
+router2 = APIRouter(prefix="/v2")
+
+@router1.get("/items")
+def get_items_v1():
+    pass
+
+@router2.get("/items")
+def get_items_v2():
+    pass

src/test/testUtils.ts

@@ -65,6 +65,10 @@ export const fixtures = {
     root: uri(join(fixturesPath, "multi-app")),
     mainPy: uri(join(fixturesPath, "multi-app", "main.py")),
   },
+  multiRouterSameFile: {
+    root: uri(join(fixturesPath, "multi-router-same-file")),
+    mainPy: uri(join(fixturesPath, "multi-router-same-file", "main.py")),
+  },
   namespace: {
     root: uri(join(fixturesPath, "namespace")),
     mainPy: uri(join(fixturesPath, "namespace", "app", "main.py")),