Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 25 additions & 17 deletions src/core/extractors.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
}
}

Expand Down
23 changes: 22 additions & 1 deletion src/core/importResolver.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

/**
Expand Down
17 changes: 17 additions & 0 deletions src/test/core/extractors.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
38 changes: 38 additions & 0 deletions src/test/core/importResolver.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
41 changes: 41 additions & 0 deletions src/test/core/routerResolver.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions src/test/fixtures/src-layout/pyproject.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
[project]
name = "src-layout-app"
6 changes: 6 additions & 0 deletions src/test/fixtures/src-layout/src/app/api/__init__.py
Original file line number Diff line number Diff line change
@@ -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"])
13 changes: 13 additions & 0 deletions src/test/fixtures/src-layout/src/app/api/users.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
from fastapi import APIRouter

router = APIRouter()


@router.get("/")
async def list_users():
return []


@router.post("/")
async def create_user():
return {}
13 changes: 13 additions & 0 deletions src/test/fixtures/src-layout/src/app/main.py
Original file line number Diff line number Diff line change
@@ -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)
4 changes: 4 additions & 0 deletions src/test/testUtils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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")),
Expand Down
Loading