diff --git a/src/core/routerResolver.ts b/src/core/routerResolver.ts index 91cd45e..c5f2d98 100644 --- a/src/core/routerResolver.ts +++ b/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) } diff --git a/src/test/core/routerResolver.test.ts b/src/test/core/routerResolver.test.ts index 3512c51..1c9a8a0 100644 --- a/src/test/core/routerResolver.test.ts +++ b/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") + } + }) }) }) diff --git a/src/test/fixtures/multi-router-same-file/main.py b/src/test/fixtures/multi-router-same-file/main.py new file mode 100644 index 0000000..1c34e61 --- /dev/null +++ b/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) diff --git a/src/test/fixtures/multi-router-same-file/routers.py b/src/test/fixtures/multi-router-same-file/routers.py new file mode 100644 index 0000000..1643231 --- /dev/null +++ b/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 diff --git a/src/test/testUtils.ts b/src/test/testUtils.ts index a146b0a..7dd22de 100644 --- a/src/test/testUtils.ts +++ b/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")),