diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 081c220..ffc90d1 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -404,25 +404,33 @@ export function importExtractor(node: Node): ImportInfo | null { moduleNode?.text ?? "", ) - // Aliased imports (e.g., "router as users_router") - const aliasedImports = getNodesByType(node).get("aliased_import") ?? [] - for (const aliased of aliasedImports) { - const nameNode = aliased.childForFieldName("name") - const aliasNode = aliased.childForFieldName("alias") - if (nameNode) { - const alias = aliasNode?.text ?? null - names.push(alias ?? nameNode.text) - namedImports.push({ name: nameNode.text, alias }) + // Collect imported names: everything after the "import" keyword. + let afterImport = false + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i) + if (!child) continue + + // Phase 1: scan forward looking for the "import" keyword + if (!afterImport) { + if (child.type === "import") afterImport = true + continue // skip this child either way (module path or the keyword itself) } - } - // Non-aliased imports (skip first dotted_name which is the module path) - const nameNodes = getNodesByType(node).get("dotted_name") ?? [] - for (let i = 1; i < nameNodes.length; i++) { - const nameNode = nameNodes[i] - if (!hasAncestor(nameNode, "aliased_import")) { - names.push(nameNode.text) - namedImports.push({ name: nameNode.text, alias: null }) + // Phase 2: we're past "import", so each child is an imported name + // (commas and other punctuation are silently skipped by the else branch) + if (child.type === "aliased_import") { + // e.g. "router as users_router" + const nameNode = child.childForFieldName("name") + const aliasNode = child.childForFieldName("alias") + if (nameNode) { + const alias = aliasNode?.text ?? null + names.push(alias ?? nameNode.text) + namedImports.push({ name: nameNode.text, alias }) + } + } else if (child.type === "dotted_name") { + // e.g. "users" + names.push(child.text) + namedImports.push({ name: child.text, alias: null }) } } diff --git a/src/core/importResolver.ts b/src/core/importResolver.ts index d8135f0..44d7842 100644 --- a/src/core/importResolver.ts +++ b/src/core/importResolver.ts @@ -122,7 +122,28 @@ export async function resolveImport( projectRootUri, fs, ) - return resolvePythonModule(resolvedUri, fs) + const result = await resolvePythonModule(resolvedUri, fs) + if (result) { + return result + } + + // Fallback for src layout: if the module wasn't found under the + // project root, try under projectRoot/src/. This handles projects where + // pyproject.toml is at the root but source lives under src/. + // Only applies to absolute imports — relative imports already resolve + // relative to the current file, not the project root. + if (!importInfo.isRelative) { + const srcRootUri = fs.joinPath(projectRootUri, "src") + const srcResolvedUri = modulePathToDir( + importInfo, + currentFileUri, + srcRootUri, + fs, + ) + return resolvePythonModule(srcResolvedUri, fs) + } + + return null } /** diff --git a/src/test/core/extractors.test.ts b/src/test/core/extractors.test.ts index 5a46257..2a9c0a9 100644 --- a/src/test/core/extractors.test.ts +++ b/src/test/core/extractors.test.ts @@ -706,6 +706,23 @@ router = f.APIRouter(prefix="/items") assert.strictEqual(result.relativeDots, 1) }) + test("extracts bare relative import (from . import X)", () => { + const code = "from . import users" + const tree = parse(code) + const nodesByType = getNodesByType(tree.rootNode) + const imports = nodesByType.get("import_from_statement") ?? [] + const result = importExtractor(imports[0]) + + assert.ok(result) + assert.strictEqual(result.modulePath, "") + assert.strictEqual(result.isRelative, true) + assert.strictEqual(result.relativeDots, 1) + assert.deepStrictEqual(result.names, ["users"]) + assert.deepStrictEqual(result.namedImports, [ + { name: "users", alias: null }, + ]) + }) + test("extracts relative import with double dot", () => { const code = "from ..api import router" const tree = parse(code) diff --git a/src/test/core/importResolver.test.ts b/src/test/core/importResolver.test.ts index 1786709..2e9fbc3 100644 --- a/src/test/core/importResolver.test.ts +++ b/src/test/core/importResolver.test.ts @@ -99,6 +99,22 @@ suite("importResolver", () => { assert.ok(result.endsWith("users.py")) }) + test("falls back to src/ for absolute imports in src layout", async () => { + // Project root is the pyproject.toml dir, but source is under src/ + const currentFile = fixtures.srcLayout.mainPy + const projectRoot = fixtures.srcLayout.workspaceRoot + + const result = await resolveImport( + { modulePath: "app.api", isRelative: false, relativeDots: 0 }, + currentFile, + projectRoot, + nodeFileSystem, + ) + + assert.ok(result) + assert.ok(result.endsWith("api/__init__.py")) + }) + test("returns null for non-existent module", async () => { const currentFile = nodeFileSystem.joinPath(standardRoot, "main.py") const projectRoot = standardRoot @@ -242,6 +258,28 @@ suite("importResolver", () => { assert.ok(result.endsWith("api_routes.py")) }) + test("resolves named import via src/ fallback for src layout", async () => { + const currentFile = fixtures.srcLayout.mainPy + const projectRoot = fixtures.srcLayout.workspaceRoot + + // "from app.api import api_router" — the actual import from the issue + const result = await resolveNamedImport( + { + modulePath: "app.api", + names: ["api_router"], + isRelative: false, + relativeDots: 0, + }, + currentFile, + projectRoot, + nodeFileSystem, + (uri) => analyzeFile(uri, parser, nodeFileSystem), + ) + + assert.ok(result) + assert.ok(result.endsWith("api/__init__.py")) + }) + test("resolves variable import from .py file (not submodule)", async () => { // This tests "from .neon import router" where router is a variable in neon.py, // NOT a submodule. Should return neon.py, not look for router.py diff --git a/src/test/core/routerResolver.test.ts b/src/test/core/routerResolver.test.ts index 1c9a8a0..7ba7a9d 100644 --- a/src/test/core/routerResolver.test.ts +++ b/src/test/core/routerResolver.test.ts @@ -604,6 +604,47 @@ suite("routerResolver", () => { assert.strictEqual(result, null) }) + test("resolves absolute imports in src layout projects", async () => { + const projectRoot = await findProjectRoot( + fixtures.srcLayout.mainPy, + fixtures.srcLayout.workspaceRoot, + nodeFileSystem, + ) + const result = await buildRouterGraph( + fixtures.srcLayout.mainPy, + parser, + projectRoot, + nodeFileSystem, + ) + + assert.ok(result) + assert.strictEqual(result.type, "FastAPI") + assert.strictEqual(result.routes.length, 1, "Should have the root route") + assert.strictEqual(result.routes[0].path, "/") + assert.strictEqual( + result.children.length, + 1, + "Should have one child router (api_router)", + ) + + const apiRouter = result.children[0].router + assert.strictEqual(apiRouter.type, "APIRouter") + assert.strictEqual( + apiRouter.children.length, + 1, + "api_router should include users router", + ) + assert.strictEqual(apiRouter.children[0].prefix, "/api/v1/users") + + const usersRouter = apiRouter.children[0].router + assert.strictEqual(usersRouter.routes.length, 2) + const methods = usersRouter.routes + .map((r) => r.method.toLowerCase()) + .sort() + assert.deepStrictEqual(methods, ["get", "post"]) + assert.ok(usersRouter.routes.every((r) => r.path === "/")) + }) + test("resolves custom APIRouter subclass as child router", async () => { const result = await buildRouterGraph( fixtures.customSubclass.mainPy, diff --git a/src/test/fixtures/src-layout/pyproject.toml b/src/test/fixtures/src-layout/pyproject.toml new file mode 100644 index 0000000..114cc2b --- /dev/null +++ b/src/test/fixtures/src-layout/pyproject.toml @@ -0,0 +1,2 @@ +[project] +name = "src-layout-app" diff --git a/src/test/fixtures/src-layout/src/app/api/__init__.py b/src/test/fixtures/src-layout/src/app/api/__init__.py new file mode 100644 index 0000000..a5bb55b --- /dev/null +++ b/src/test/fixtures/src-layout/src/app/api/__init__.py @@ -0,0 +1,6 @@ +from fastapi import APIRouter +from . import users + +api_router = APIRouter() + +api_router.include_router(users.router, prefix="/api/v1/users", tags=["users"]) diff --git a/src/test/fixtures/src-layout/src/app/api/users.py b/src/test/fixtures/src-layout/src/app/api/users.py new file mode 100644 index 0000000..ae9febe --- /dev/null +++ b/src/test/fixtures/src-layout/src/app/api/users.py @@ -0,0 +1,13 @@ +from fastapi import APIRouter + +router = APIRouter() + + +@router.get("/") +async def list_users(): + return [] + + +@router.post("/") +async def create_user(): + return {} diff --git a/src/test/fixtures/src-layout/src/app/main.py b/src/test/fixtures/src-layout/src/app/main.py new file mode 100644 index 0000000..15561a2 --- /dev/null +++ b/src/test/fixtures/src-layout/src/app/main.py @@ -0,0 +1,13 @@ +from fastapi import FastAPI + +from app.api import api_router + +app = FastAPI() + + +@app.get("/") +async def root(): + return {"message": "Hello"} + + +app.include_router(api_router) diff --git a/src/test/testUtils.ts b/src/test/testUtils.ts index 7dd22de..a078b4e 100644 --- a/src/test/testUtils.ts +++ b/src/test/testUtils.ts @@ -93,6 +93,10 @@ export const fixtures = { join(fixturesPath, "reexport", "app", "integrations", "__init__.py"), ), }, + srcLayout: { + workspaceRoot: uri(join(fixturesPath, "src-layout")), + mainPy: uri(join(fixturesPath, "src-layout", "src", "app", "main.py")), + }, sameFile: { root: uri(join(fixturesPath, "same-file")), mainPy: uri(join(fixturesPath, "same-file", "main.py")),