From 9416c48f6bb4dc17ae5b24f1e0c92ab5d93b5044 Mon Sep 17 00:00:00 2001 From: Savannah Ostrowski Date: Wed, 4 Mar 2026 08:42:31 -0800 Subject: [PATCH 1/6] =?UTF-8?q?=F0=9F=90=9B=20Fix=20discovery=20for=20subc?= =?UTF-8?q?lassed=20routers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/analyzer.ts | 6 ++- src/core/extractors.ts | 28 ++++++++++- src/test/core/extractors.test.ts | 50 ++++++++++++++++++++ src/test/core/routerResolver.test.ts | 31 ++++++++++++ src/test/fixtures/custom-subclass/main.py | 5 ++ src/test/fixtures/custom-subclass/routers.py | 18 +++++++ src/test/testUtils.ts | 4 ++ 7 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 src/test/fixtures/custom-subclass/main.py create mode 100644 src/test/fixtures/custom-subclass/routers.py diff --git a/src/core/analyzer.ts b/src/core/analyzer.ts index 3390b15..ba5b784 100644 --- a/src/core/analyzer.ts +++ b/src/core/analyzer.ts @@ -5,6 +5,7 @@ import type { Tree } from "web-tree-sitter" import { logError } from "../utils/logger" import { + collectAPIRouterSubclasses, collectStringVariables, decoratorExtractor, findNodesByType, @@ -44,7 +45,10 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis { // Get all router assignments const assignments = findNodesByType(rootNode, "assignment") - const routers = assignments.map(routerExtractor).filter(notNull) + const apiRouterSubclasses = collectAPIRouterSubclasses(rootNode) + const routers = assignments + .map((node) => routerExtractor(node, apiRouterSubclasses)) + .filter(notNull) // Get all include_router and mount calls const callNodes = findNodesByType(rootNode, "call") diff --git a/src/core/extractors.ts b/src/core/extractors.ts index ccd127b..9a4f8ea 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -49,6 +49,23 @@ function stripDocstring(raw: string): string { return dedented.join("\n").trim() } +export function collectAPIRouterSubclasses(rootNode: Node): Set { + const subclasses = new Set() + for (const cls of findNodesByType(rootNode, "class_definition")) { + const nameNode = cls.childForFieldName("name") + const superclassesNode = cls.childForFieldName("superclasses") + if (!nameNode || !superclassesNode) continue + + const extendsAPIRouter = superclassesNode.namedChildren.some( + (s) => s.text === "APIRouter" || s.text === "fastapi.APIRouter", + ) + if (extendsAPIRouter) { + subclasses.add(nameNode.text) + } + } + return subclasses +} + function collectNodesByType(node: Node, type: string, results: Node[]): void { if (node.type === type) { results.push(node) @@ -237,7 +254,10 @@ function extractTags(listNode: Node): string[] { .filter((v): v is string => v !== null) } -export function routerExtractor(node: Node): RouterInfo | null { +export function routerExtractor( + node: Node, + apiRouterSubclasses?: Set, +): RouterInfo | null { if (node.type !== "assignment") { return null } @@ -250,7 +270,11 @@ export function routerExtractor(node: Node): RouterInfo | null { const funcName = valueNode.childForFieldName("function")?.text let type: RouterType - if (funcName === "APIRouter" || funcName === "fastapi.APIRouter") { + if ( + funcName === "APIRouter" || + funcName === "fastapi.APIRouter" || + (funcName !== undefined && apiRouterSubclasses?.has(funcName)) + ) { type = "APIRouter" } else if (funcName === "FastAPI" || funcName === "fastapi.FastAPI") { type = "FastAPI" diff --git a/src/test/core/extractors.test.ts b/src/test/core/extractors.test.ts index a449f2b..70c89cc 100644 --- a/src/test/core/extractors.test.ts +++ b/src/test/core/extractors.test.ts @@ -1,5 +1,6 @@ import * as assert from "node:assert" import { + collectAPIRouterSubclasses, decoratorExtractor, extractPathFromNode, extractStringValue, @@ -465,6 +466,55 @@ def list_users(): assert.strictEqual(result.type, "APIRouter") assert.strictEqual(result.prefix, "/api") }) + + test("returns null for custom subclass without subclasses set", () => { + const code = "admin_router = AdminAPIRouter(prefix='/admin')" + const tree = parse(code) + const assignments = findNodesByType(tree.rootNode, "assignment") + const result = routerExtractor(assignments[0]) + + assert.strictEqual(result, null) + }) + + test("recognizes custom APIRouter subclass when subclasses set provided", () => { + const code = ` +class AdminAPIRouter(APIRouter): + pass + +admin_router = AdminAPIRouter(prefix="/admin") +` + const tree = parse(code) + const subclasses = collectAPIRouterSubclasses(tree.rootNode) + assert.ok(subclasses.has("AdminAPIRouter")) + + const assignments = findNodesByType(tree.rootNode, "assignment") + const result = routerExtractor(assignments[0], subclasses) + + assert.ok(result) + assert.strictEqual(result.variableName, "admin_router") + assert.strictEqual(result.type, "APIRouter") + assert.strictEqual(result.prefix, "/admin") + }) + + test("does not treat FastAPI subclass as APIRouter", () => { + const code = ` +class MyApp(FastAPI): + pass + +app = MyApp() +` + const tree = parse(code) + const subclasses = collectAPIRouterSubclasses(tree.rootNode) + assert.ok( + !subclasses.has("MyApp"), + "FastAPI subclass should not be in APIRouter subclasses", + ) + + const assignments = findNodesByType(tree.rootNode, "assignment") + const result = routerExtractor(assignments[0], subclasses) + + assert.strictEqual(result, null) + }) }) suite("importExtractor", () => { diff --git a/src/test/core/routerResolver.test.ts b/src/test/core/routerResolver.test.ts index b37d4c0..8b25f3d 100644 --- a/src/test/core/routerResolver.test.ts +++ b/src/test/core/routerResolver.test.ts @@ -570,5 +570,36 @@ suite("routerResolver", () => { assert.strictEqual(result, null) }) + + test("resolves custom APIRouter subclass as child router", async () => { + const result = await buildRouterGraph( + fixtures.customSubclass.mainPy, + parser, + fixtures.customSubclass.root, + nodeFileSystem, + ) + + assert.ok(result) + assert.strictEqual(result.type, "FastAPI") + assert.strictEqual(result.variableName, "app") + + assert.strictEqual( + result.children.length, + 1, + "Should have one child router", + ) + + const adminRouter = result.children[0].router + assert.strictEqual(adminRouter.type, "APIRouter") + assert.strictEqual(adminRouter.prefix, "/admin") + assert.strictEqual(adminRouter.routes.length, 2) + + const paths = adminRouter.routes.map((r) => r.path) + assert.ok(paths.includes("/users")) + + const methods = adminRouter.routes.map((r) => r.method.toLowerCase()) + assert.ok(methods.includes("get")) + assert.ok(methods.includes("post")) + }) }) }) diff --git a/src/test/fixtures/custom-subclass/main.py b/src/test/fixtures/custom-subclass/main.py new file mode 100644 index 0000000..2093db3 --- /dev/null +++ b/src/test/fixtures/custom-subclass/main.py @@ -0,0 +1,5 @@ +from fastapi import FastAPI +from routers import admin_router + +app = FastAPI() +app.include_router(admin_router) diff --git a/src/test/fixtures/custom-subclass/routers.py b/src/test/fixtures/custom-subclass/routers.py new file mode 100644 index 0000000..b8217df --- /dev/null +++ b/src/test/fixtures/custom-subclass/routers.py @@ -0,0 +1,18 @@ +from fastapi import APIRouter + + +class AdminAPIRouter(APIRouter): + pass + + +admin_router = AdminAPIRouter(prefix="/admin") + + +@admin_router.get("/users") +def list_users(): + return [] + + +@admin_router.post("/users") +def create_user(): + return {} diff --git a/src/test/testUtils.ts b/src/test/testUtils.ts index 02751c2..854ac03 100644 --- a/src/test/testUtils.ts +++ b/src/test/testUtils.ts @@ -87,6 +87,10 @@ export const fixtures = { root: uri(join(fixturesPath, "factory-func")), mainPy: uri(join(fixturesPath, "factory-func", "main.py")), }, + customSubclass: { + root: uri(join(fixturesPath, "custom-subclass")), + mainPy: uri(join(fixturesPath, "custom-subclass", "main.py")), + }, } /** From 2621cbd9b8b9ab3d5f9329559d479d8e55a41141 Mon Sep 17 00:00:00 2001 From: Savannah Ostrowski Date: Wed, 4 Mar 2026 08:49:00 -0800 Subject: [PATCH 2/6] =?UTF-8?q?=F0=9F=90=9B=20Aliases?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/analyzer.ts | 5 ++- src/core/extractors.ts | 32 ++++++++++++++++- src/test/core/extractors.test.ts | 48 +++++++++++++++++++++++++ src/test/core/routerResolver.test.ts | 24 +++++++++++++ src/test/fixtures/aliased-class/main.py | 18 ++++++++++ src/test/testUtils.ts | 4 +++ 6 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 src/test/fixtures/aliased-class/main.py diff --git a/src/core/analyzer.ts b/src/core/analyzer.ts index ba5b784..15b9c88 100644 --- a/src/core/analyzer.ts +++ b/src/core/analyzer.ts @@ -6,6 +6,7 @@ import type { Tree } from "web-tree-sitter" import { logError } from "../utils/logger" import { collectAPIRouterSubclasses, + collectFastAPIAliases, collectStringVariables, decoratorExtractor, findNodesByType, @@ -46,8 +47,10 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis { // Get all router assignments const assignments = findNodesByType(rootNode, "assignment") const apiRouterSubclasses = collectAPIRouterSubclasses(rootNode) + const { fastAPIAliases, apiRouterAliases } = collectFastAPIAliases(rootNode) + for (const alias of apiRouterAliases) apiRouterSubclasses.add(alias) const routers = assignments - .map((node) => routerExtractor(node, apiRouterSubclasses)) + .map((node) => routerExtractor(node, apiRouterSubclasses, fastAPIAliases)) .filter(notNull) // Get all include_router and mount calls diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 9a4f8ea..0973fd4 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -66,6 +66,31 @@ export function collectAPIRouterSubclasses(rootNode: Node): Set { return subclasses } +/** + * Collects aliases for FastAPI and APIRouter from import statements. + * Handles: from fastapi import FastAPI as FA, APIRouter as AR + */ +export function collectFastAPIAliases(rootNode: Node): { + fastAPIAliases: Set + apiRouterAliases: Set +} { + const fastAPIAliases = new Set() + const apiRouterAliases = new Set() + + for (const node of findNodesByType(rootNode, "import_from_statement")) { + const info = importExtractor(node) + if (!info || info.modulePath !== "fastapi") continue + + for (const ni of info.namedImports) { + if (ni.alias === null) continue + if (ni.name === "FastAPI") fastAPIAliases.add(ni.alias) + if (ni.name === "APIRouter") apiRouterAliases.add(ni.alias) + } + } + + return { fastAPIAliases, apiRouterAliases } +} + function collectNodesByType(node: Node, type: string, results: Node[]): void { if (node.type === type) { results.push(node) @@ -257,6 +282,7 @@ function extractTags(listNode: Node): string[] { export function routerExtractor( node: Node, apiRouterSubclasses?: Set, + fastAPIAliases?: Set, ): RouterInfo | null { if (node.type !== "assignment") { return null @@ -276,7 +302,11 @@ export function routerExtractor( (funcName !== undefined && apiRouterSubclasses?.has(funcName)) ) { type = "APIRouter" - } else if (funcName === "FastAPI" || funcName === "fastapi.FastAPI") { + } else if ( + funcName === "FastAPI" || + funcName === "fastapi.FastAPI" || + (funcName !== undefined && fastAPIAliases?.has(funcName)) + ) { type = "FastAPI" } else { return null diff --git a/src/test/core/extractors.test.ts b/src/test/core/extractors.test.ts index 70c89cc..3c820a2 100644 --- a/src/test/core/extractors.test.ts +++ b/src/test/core/extractors.test.ts @@ -1,6 +1,7 @@ import * as assert from "node:assert" import { collectAPIRouterSubclasses, + collectFastAPIAliases, decoratorExtractor, extractPathFromNode, extractStringValue, @@ -515,6 +516,53 @@ app = MyApp() assert.strictEqual(result, null) }) + + test("recognizes aliased FastAPI import (FastAPI as FA)", () => { + const code = ` +from fastapi import FastAPI as FA + +app = FA() +` + const tree = parse(code) + const { fastAPIAliases } = collectFastAPIAliases(tree.rootNode) + assert.ok(fastAPIAliases.has("FA")) + + const assignments = findNodesByType(tree.rootNode, "assignment") + const result = routerExtractor(assignments[0], undefined, fastAPIAliases) + + assert.ok(result) + assert.strictEqual(result.variableName, "app") + assert.strictEqual(result.type, "FastAPI") + }) + + test("recognizes aliased APIRouter import (APIRouter as AR)", () => { + const code = ` +from fastapi import APIRouter as AR + +router = AR(prefix="/items") +` + const tree = parse(code) + const { apiRouterAliases } = collectFastAPIAliases(tree.rootNode) + assert.ok(apiRouterAliases.has("AR")) + + const assignments = findNodesByType(tree.rootNode, "assignment") + const result = routerExtractor(assignments[0], apiRouterAliases) + + assert.ok(result) + assert.strictEqual(result.variableName, "router") + assert.strictEqual(result.type, "APIRouter") + assert.strictEqual(result.prefix, "/items") + }) + + test("collectFastAPIAliases ignores non-aliased imports", () => { + const code = "from fastapi import FastAPI, APIRouter" + const tree = parse(code) + const { fastAPIAliases, apiRouterAliases } = collectFastAPIAliases( + tree.rootNode, + ) + assert.strictEqual(fastAPIAliases.size, 0) + assert.strictEqual(apiRouterAliases.size, 0) + }) }) suite("importExtractor", () => { diff --git a/src/test/core/routerResolver.test.ts b/src/test/core/routerResolver.test.ts index 8b25f3d..4422d29 100644 --- a/src/test/core/routerResolver.test.ts +++ b/src/test/core/routerResolver.test.ts @@ -601,5 +601,29 @@ suite("routerResolver", () => { assert.ok(methods.includes("get")) assert.ok(methods.includes("post")) }) + + test("resolves aliased FastAPI and APIRouter class imports", async () => { + const result = await buildRouterGraph( + fixtures.aliasedClass.mainPy, + parser, + fixtures.aliasedClass.root, + nodeFileSystem, + ) + + assert.ok(result) + assert.strictEqual(result.type, "FastAPI") + assert.strictEqual(result.variableName, "app") + + assert.strictEqual( + result.children.length, + 1, + "Should have one child router", + ) + + const usersRouter = result.children[0].router + assert.strictEqual(usersRouter.type, "APIRouter") + assert.strictEqual(usersRouter.prefix, "/users") + assert.strictEqual(usersRouter.routes.length, 2) + }) }) }) diff --git a/src/test/fixtures/aliased-class/main.py b/src/test/fixtures/aliased-class/main.py new file mode 100644 index 0000000..616a5b2 --- /dev/null +++ b/src/test/fixtures/aliased-class/main.py @@ -0,0 +1,18 @@ +from fastapi import FastAPI as FA +from fastapi import APIRouter as AR + +app = FA() +router = AR(prefix="/users") + + +@router.get("/") +def list_users(): + return [] + + +@router.post("/") +def create_user(): + return {} + + +app.include_router(router) diff --git a/src/test/testUtils.ts b/src/test/testUtils.ts index 854ac03..f4a8a99 100644 --- a/src/test/testUtils.ts +++ b/src/test/testUtils.ts @@ -91,6 +91,10 @@ export const fixtures = { root: uri(join(fixturesPath, "custom-subclass")), mainPy: uri(join(fixturesPath, "custom-subclass", "main.py")), }, + aliasedClass: { + root: uri(join(fixturesPath, "aliased-class")), + mainPy: uri(join(fixturesPath, "aliased-class", "main.py")), + }, } /** From 346aa66a96a6cdb0ea1b07a347a55b31ba2d6f83 Mon Sep 17 00:00:00 2001 From: Savannah Ostrowski Date: Wed, 4 Mar 2026 09:03:37 -0800 Subject: [PATCH 3/6] =?UTF-8?q?=F0=9F=90=9B=20Handle=20module=20aliase?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/analyzer.ts | 9 +- src/core/extractors.ts | 110 ++++++++++++------- src/test/core/extractors.test.ts | 132 +++++++++++++++++++---- src/test/core/routerResolver.test.ts | 20 ++++ src/test/fixtures/aliased-module/main.py | 17 +++ src/test/testUtils.ts | 4 + 6 files changed, 224 insertions(+), 68 deletions(-) create mode 100644 src/test/fixtures/aliased-module/main.py diff --git a/src/core/analyzer.ts b/src/core/analyzer.ts index 15b9c88..754593b 100644 --- a/src/core/analyzer.ts +++ b/src/core/analyzer.ts @@ -5,8 +5,7 @@ import type { Tree } from "web-tree-sitter" import { logError } from "../utils/logger" import { - collectAPIRouterSubclasses, - collectFastAPIAliases, + collectRecognizedNames, collectStringVariables, decoratorExtractor, findNodesByType, @@ -46,11 +45,9 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis { // Get all router assignments const assignments = findNodesByType(rootNode, "assignment") - const apiRouterSubclasses = collectAPIRouterSubclasses(rootNode) - const { fastAPIAliases, apiRouterAliases } = collectFastAPIAliases(rootNode) - for (const alias of apiRouterAliases) apiRouterSubclasses.add(alias) + const { fastAPINames, apiRouterNames } = collectRecognizedNames(rootNode) const routers = assignments - .map((node) => routerExtractor(node, apiRouterSubclasses, fastAPIAliases)) + .map((node) => routerExtractor(node, apiRouterNames, fastAPINames)) .filter(notNull) // Get all include_router and mount calls diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 0973fd4..42c8168 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -48,47 +48,66 @@ function stripDocstring(raw: string): string { const dedented = lines.map((l, i) => (i === 0 ? l : l.slice(minIndent))) return dedented.join("\n").trim() } - -export function collectAPIRouterSubclasses(rootNode: Node): Set { - const subclasses = new Set() - for (const cls of findNodesByType(rootNode, "class_definition")) { - const nameNode = cls.childForFieldName("name") - const superclassesNode = cls.childForFieldName("superclasses") - if (!nameNode || !superclassesNode) continue - - const extendsAPIRouter = superclassesNode.namedChildren.some( - (s) => s.text === "APIRouter" || s.text === "fastapi.APIRouter", - ) - if (extendsAPIRouter) { - subclasses.add(nameNode.text) - } - } - return subclasses -} - -/** - * Collects aliases for FastAPI and APIRouter from import statements. - * Handles: from fastapi import FastAPI as FA, APIRouter as AR +/* + * Extracts recognized FastAPI and APIRouter names from imports and class definitions + * This allows the rest of the extractors to handle user-defined aliases and subclasses. + * + * For example, if the code has: + * from fastapi import FastAPI as MyApp + * from fastapi import APIRouter as MyRouter + * class CustomRouter(MyRouter): ... + * + * Then this function will return: + * fastAPINames = Set { "FastAPI", "fastapi.FastAPI", "MyApp" } + * apiRouterNames = Set { "APIRouter", "fastapi.APIRouter", "MyRouter", "CustomRouter" } + * + * This allows decoratorExtractor and routerExtractor to recognize routes and routers + * defined using these aliases and subclasses without needing complex logic in those functions. */ -export function collectFastAPIAliases(rootNode: Node): { - fastAPIAliases: Set - apiRouterAliases: Set +export function collectRecognizedNames(rootNode: Node): { + fastAPINames: Set + apiRouterNames: Set } { - const fastAPIAliases = new Set() - const apiRouterAliases = new Set() + const fastAPINames = new Set(["FastAPI", "fastapi.FastAPI"]) + const apiRouterNames = new Set(["APIRouter", "fastapi.APIRouter"]) + // Add aliases from "from fastapi import X as Y" imports for (const node of findNodesByType(rootNode, "import_from_statement")) { const info = importExtractor(node) if (!info || info.modulePath !== "fastapi") continue + for (const named of info.namedImports) { + if (named.alias === null) continue + if (named.name === "FastAPI") fastAPINames.add(named.alias) + else if (named.name === "APIRouter") apiRouterNames.add(named.alias) + } + } + + // Add module aliases from "import fastapi as f" → recognizes f.FastAPI, f.APIRouter + for (const node of findNodesByType(rootNode, "import_statement")) { + const info = importExtractor(node) + if (!info) continue + for (const named of info.namedImports) { + if (named.alias === null) continue + if (named.name === "fastapi") { + fastAPINames.add(`${named.alias}.FastAPI`) + apiRouterNames.add(`${named.alias}.APIRouter`) + } + } + } - for (const ni of info.namedImports) { - if (ni.alias === null) continue - if (ni.name === "FastAPI") fastAPIAliases.add(ni.alias) - if (ni.name === "APIRouter") apiRouterAliases.add(ni.alias) + // Add subclasses, checking against the already-accumulated alias sets so + // "class MyRouter(AR)" works when AR is an alias for APIRouter + for (const cls of findNodesByType(rootNode, "class_definition")) { + const nameNode = cls.childForFieldName("name") + const superclassesNode = cls.childForFieldName("superclasses") + if (!nameNode || !superclassesNode) continue + for (const parent of superclassesNode.namedChildren) { + if (apiRouterNames.has(parent.text)) apiRouterNames.add(nameNode.text) + else if (fastAPINames.has(parent.text)) fastAPINames.add(nameNode.text) } } - return { fastAPIAliases, apiRouterAliases } + return { fastAPINames, apiRouterNames } } function collectNodesByType(node: Node, type: string, results: Node[]): void { @@ -281,8 +300,8 @@ function extractTags(listNode: Node): string[] { export function routerExtractor( node: Node, - apiRouterSubclasses?: Set, - fastAPIAliases?: Set, + apiRouterNames?: Set, + fastAPINames?: Set, ): RouterInfo | null { if (node.type !== "assignment") { return null @@ -299,13 +318,13 @@ export function routerExtractor( if ( funcName === "APIRouter" || funcName === "fastapi.APIRouter" || - (funcName !== undefined && apiRouterSubclasses?.has(funcName)) + (funcName !== undefined && apiRouterNames?.has(funcName)) ) { type = "APIRouter" } else if ( funcName === "FastAPI" || funcName === "fastapi.FastAPI" || - (funcName !== undefined && fastAPIAliases?.has(funcName)) + (funcName !== undefined && fastAPINames?.has(funcName)) ) { type = "FastAPI" } else { @@ -380,13 +399,26 @@ export function importExtractor(node: Node): ImportInfo | null { const namedImports: ImportedName[] = [] if (node.type === "import_statement") { + // Handle aliased imports: "import fastapi as f" + for (const aliased of findNodesByType(node, "aliased_import")) { + 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 }) + } + } + // Non-aliased: "import fastapi" const nameNodes = findNodesByType(node, "dotted_name") for (const nameNode of nameNodes) { - const firstName = nameNode.text.split(".")[0] - names.push(firstName) - namedImports.push({ name: firstName, alias: null }) + if (!hasAncestor(nameNode, "aliased_import")) { + const firstName = nameNode.text.split(".")[0] + names.push(firstName) + namedImports.push({ name: firstName, alias: null }) + } } - const modulePath = nameNodes[0]?.text ?? "" + const modulePath = namedImports[0]?.name ?? "" return { modulePath, names, diff --git a/src/test/core/extractors.test.ts b/src/test/core/extractors.test.ts index 3c820a2..4cf79a5 100644 --- a/src/test/core/extractors.test.ts +++ b/src/test/core/extractors.test.ts @@ -1,7 +1,6 @@ import * as assert from "node:assert" import { - collectAPIRouterSubclasses, - collectFastAPIAliases, + collectRecognizedNames, decoratorExtractor, extractPathFromNode, extractStringValue, @@ -477,7 +476,7 @@ def list_users(): assert.strictEqual(result, null) }) - test("recognizes custom APIRouter subclass when subclasses set provided", () => { + test("recognizes custom APIRouter subclass", () => { const code = ` class AdminAPIRouter(APIRouter): pass @@ -485,11 +484,11 @@ class AdminAPIRouter(APIRouter): admin_router = AdminAPIRouter(prefix="/admin") ` const tree = parse(code) - const subclasses = collectAPIRouterSubclasses(tree.rootNode) - assert.ok(subclasses.has("AdminAPIRouter")) + const { apiRouterNames } = collectRecognizedNames(tree.rootNode) + assert.ok(apiRouterNames.has("AdminAPIRouter")) const assignments = findNodesByType(tree.rootNode, "assignment") - const result = routerExtractor(assignments[0], subclasses) + const result = routerExtractor(assignments[0], apiRouterNames) assert.ok(result) assert.strictEqual(result.variableName, "admin_router") @@ -497,7 +496,7 @@ admin_router = AdminAPIRouter(prefix="/admin") assert.strictEqual(result.prefix, "/admin") }) - test("does not treat FastAPI subclass as APIRouter", () => { + test("recognizes FastAPI subclass", () => { const code = ` class MyApp(FastAPI): pass @@ -505,16 +504,22 @@ class MyApp(FastAPI): app = MyApp() ` const tree = parse(code) - const subclasses = collectAPIRouterSubclasses(tree.rootNode) - assert.ok( - !subclasses.has("MyApp"), - "FastAPI subclass should not be in APIRouter subclasses", + const { fastAPINames, apiRouterNames } = collectRecognizedNames( + tree.rootNode, ) + assert.ok(fastAPINames.has("MyApp")) + assert.ok(!apiRouterNames.has("MyApp")) const assignments = findNodesByType(tree.rootNode, "assignment") - const result = routerExtractor(assignments[0], subclasses) + const result = routerExtractor( + assignments[0], + apiRouterNames, + fastAPINames, + ) - assert.strictEqual(result, null) + assert.ok(result) + assert.strictEqual(result.variableName, "app") + assert.strictEqual(result.type, "FastAPI") }) test("recognizes aliased FastAPI import (FastAPI as FA)", () => { @@ -524,11 +529,17 @@ from fastapi import FastAPI as FA app = FA() ` const tree = parse(code) - const { fastAPIAliases } = collectFastAPIAliases(tree.rootNode) - assert.ok(fastAPIAliases.has("FA")) + const { fastAPINames, apiRouterNames } = collectRecognizedNames( + tree.rootNode, + ) + assert.ok(fastAPINames.has("FA")) const assignments = findNodesByType(tree.rootNode, "assignment") - const result = routerExtractor(assignments[0], undefined, fastAPIAliases) + const result = routerExtractor( + assignments[0], + apiRouterNames, + fastAPINames, + ) assert.ok(result) assert.strictEqual(result.variableName, "app") @@ -542,11 +553,11 @@ from fastapi import APIRouter as AR router = AR(prefix="/items") ` const tree = parse(code) - const { apiRouterAliases } = collectFastAPIAliases(tree.rootNode) - assert.ok(apiRouterAliases.has("AR")) + const { apiRouterNames } = collectRecognizedNames(tree.rootNode) + assert.ok(apiRouterNames.has("AR")) const assignments = findNodesByType(tree.rootNode, "assignment") - const result = routerExtractor(assignments[0], apiRouterAliases) + const result = routerExtractor(assignments[0], apiRouterNames) assert.ok(result) assert.strictEqual(result.variableName, "router") @@ -554,14 +565,71 @@ router = AR(prefix="/items") assert.strictEqual(result.prefix, "/items") }) - test("collectFastAPIAliases ignores non-aliased imports", () => { + test("recognizes subclass of aliased APIRouter (class MyRouter(AR))", () => { + const code = ` +from fastapi import APIRouter as AR + +class MyRouter(AR): + pass + +router = MyRouter(prefix="/items") +` + const tree = parse(code) + const { apiRouterNames } = collectRecognizedNames(tree.rootNode) + assert.ok(apiRouterNames.has("AR")) + assert.ok(apiRouterNames.has("MyRouter")) + + const assignments = findNodesByType(tree.rootNode, "assignment") + const result = routerExtractor(assignments[0], apiRouterNames) + + assert.ok(result) + assert.strictEqual(result.type, "APIRouter") + }) + + test("collectRecognizedNames ignores non-aliased imports", () => { const code = "from fastapi import FastAPI, APIRouter" const tree = parse(code) - const { fastAPIAliases, apiRouterAliases } = collectFastAPIAliases( + const { fastAPINames, apiRouterNames } = collectRecognizedNames( + tree.rootNode, + ) + // Only the defaults — no extras from non-aliased imports + assert.strictEqual(fastAPINames.size, 2) // "FastAPI", "fastapi.FastAPI" + assert.strictEqual(apiRouterNames.size, 2) // "APIRouter", "fastapi.APIRouter" + }) + + test("recognizes module alias (import fastapi as f)", () => { + const code = ` +import fastapi as f + +app = f.FastAPI() +router = f.APIRouter(prefix="/items") +` + const tree = parse(code) + const { fastAPINames, apiRouterNames } = collectRecognizedNames( tree.rootNode, ) - assert.strictEqual(fastAPIAliases.size, 0) - assert.strictEqual(apiRouterAliases.size, 0) + assert.ok(fastAPINames.has("f.FastAPI")) + assert.ok(apiRouterNames.has("f.APIRouter")) + + const assignments = findNodesByType(tree.rootNode, "assignment") + const appResult = routerExtractor( + assignments[0], + apiRouterNames, + fastAPINames, + ) + assert.ok(appResult) + assert.strictEqual(appResult.variableName, "app") + assert.strictEqual(appResult.type, "FastAPI") + + const routerResult = routerExtractor( + assignments[1], + apiRouterNames, + fastAPINames, + ) + assert.ok(routerResult) + assert.strictEqual(routerResult.variableName, "router") + assert.strictEqual(routerResult.type, "APIRouter") + assert.strictEqual(routerResult.prefix, "/items") }) }) @@ -575,6 +643,24 @@ router = AR(prefix="/items") assert.ok(result) assert.strictEqual(result.modulePath, "fastapi") assert.deepStrictEqual(result.names, ["fastapi"]) + assert.deepStrictEqual(result.namedImports, [ + { name: "fastapi", alias: null }, + ]) + assert.strictEqual(result.isRelative, false) + }) + + test("extracts aliased module import (import fastapi as f)", () => { + const code = "import fastapi as f" + const tree = parse(code) + const imports = findNodesByType(tree.rootNode, "import_statement") + const result = importExtractor(imports[0]) + + assert.ok(result) + assert.strictEqual(result.modulePath, "fastapi") + assert.deepStrictEqual(result.names, ["f"]) + assert.deepStrictEqual(result.namedImports, [ + { name: "fastapi", alias: "f" }, + ]) assert.strictEqual(result.isRelative, false) }) diff --git a/src/test/core/routerResolver.test.ts b/src/test/core/routerResolver.test.ts index 4422d29..3e92675 100644 --- a/src/test/core/routerResolver.test.ts +++ b/src/test/core/routerResolver.test.ts @@ -625,5 +625,25 @@ suite("routerResolver", () => { assert.strictEqual(usersRouter.prefix, "/users") assert.strictEqual(usersRouter.routes.length, 2) }) + + test("resolves module-aliased fastapi import (import fastapi as f)", async () => { + const result = await buildRouterGraph( + fixtures.aliasedModule.mainPy, + parser, + fixtures.aliasedModule.root, + nodeFileSystem, + ) + + assert.ok(result) + assert.strictEqual(result.type, "FastAPI") + assert.strictEqual(result.variableName, "app") + + assert.strictEqual(result.children.length, 1) + + const usersRouter = result.children[0].router + assert.strictEqual(usersRouter.type, "APIRouter") + assert.strictEqual(usersRouter.prefix, "/users") + assert.strictEqual(usersRouter.routes.length, 2) + }) }) }) diff --git a/src/test/fixtures/aliased-module/main.py b/src/test/fixtures/aliased-module/main.py new file mode 100644 index 0000000..f10be75 --- /dev/null +++ b/src/test/fixtures/aliased-module/main.py @@ -0,0 +1,17 @@ +import fastapi as f + +app = f.FastAPI() +router = f.APIRouter(prefix="/users") + + +@router.get("/") +def list_users(): + return [] + + +@router.post("/") +def create_user(): + return {} + + +app.include_router(router) diff --git a/src/test/testUtils.ts b/src/test/testUtils.ts index f4a8a99..8724b59 100644 --- a/src/test/testUtils.ts +++ b/src/test/testUtils.ts @@ -95,6 +95,10 @@ export const fixtures = { root: uri(join(fixturesPath, "aliased-class")), mainPy: uri(join(fixturesPath, "aliased-class", "main.py")), }, + aliasedModule: { + root: uri(join(fixturesPath, "aliased-module")), + mainPy: uri(join(fixturesPath, "aliased-module", "main.py")), + }, } /** From beb9b54e9ba8d8155732f3b684e5b2b2db94e410 Mon Sep 17 00:00:00 2001 From: Savannah Ostrowski Date: Wed, 4 Mar 2026 09:33:33 -0800 Subject: [PATCH 4/6] =?UTF-8?q?=F0=9F=90=9B=20Fix=20named=20args?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/extractors.ts | 8 +++++--- src/test/core/extractors.test.ts | 14 ++++++++++++++ src/test/providers/testCodeLensProvider.test.ts | 16 ++++++++++++++++ src/vscode/testCodeLensProvider.ts | 12 ++++++++++-- 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 42c8168..9ddb4b2 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -399,26 +399,28 @@ export function importExtractor(node: Node): ImportInfo | null { const namedImports: ImportedName[] = [] if (node.type === "import_statement") { + let modulePath = "" // Handle aliased imports: "import fastapi as f" for (const aliased of findNodesByType(node, "aliased_import")) { const nameNode = aliased.childForFieldName("name") const aliasNode = aliased.childForFieldName("alias") if (nameNode) { + if (!modulePath) modulePath = nameNode.text // preserve full dotted path const alias = aliasNode?.text ?? null names.push(alias ?? nameNode.text) namedImports.push({ name: nameNode.text, alias }) } } - // Non-aliased: "import fastapi" + // Non-aliased: "import fastapi" or "import fastapi.routing" const nameNodes = findNodesByType(node, "dotted_name") for (const nameNode of nameNodes) { if (!hasAncestor(nameNode, "aliased_import")) { + if (!modulePath) modulePath = nameNode.text // preserve full dotted path const firstName = nameNode.text.split(".")[0] names.push(firstName) namedImports.push({ name: firstName, alias: null }) } } - const modulePath = namedImports[0]?.name ?? "" return { modulePath, names, @@ -465,7 +467,7 @@ export function importExtractor(node: Node): ImportInfo | null { * app.get("/users", response_model=List[User]) → position 0 = string node "/users" * app.get(path="/users", response_model=List[User]) → keyword "path" = string node "/users" */ -function resolveArgNode( +export function resolveArgNode( args: Node[], position: number, keywordName: string, diff --git a/src/test/core/extractors.test.ts b/src/test/core/extractors.test.ts index 4cf79a5..5b1c2db 100644 --- a/src/test/core/extractors.test.ts +++ b/src/test/core/extractors.test.ts @@ -649,6 +649,20 @@ router = f.APIRouter(prefix="/items") assert.strictEqual(result.isRelative, false) }) + test("preserves full dotted modulePath for import fastapi.routing", () => { + const code = "import fastapi.routing" + const tree = parse(code) + const imports = findNodesByType(tree.rootNode, "import_statement") + const result = importExtractor(imports[0]) + + assert.ok(result) + assert.strictEqual(result.modulePath, "fastapi.routing") + assert.deepStrictEqual(result.names, ["fastapi"]) + assert.deepStrictEqual(result.namedImports, [ + { name: "fastapi", alias: null }, + ]) + }) + test("extracts aliased module import (import fastapi as f)", () => { const code = "import fastapi as f" const tree = parse(code) diff --git a/src/test/providers/testCodeLensProvider.test.ts b/src/test/providers/testCodeLensProvider.test.ts index adefffb..5d067d5 100644 --- a/src/test/providers/testCodeLensProvider.test.ts +++ b/src/test/providers/testCodeLensProvider.test.ts @@ -287,6 +287,22 @@ def test_something(): assert.strictEqual(lenses.length, 0) }) + test("creates CodeLens for url= keyword argument", async () => { + const app = createMockApp([createRoute("GET", "/users")]) + provider.setApps([app]) + + const doc = await vscode.workspace.openTextDocument({ + content: ` +def test_get_users(): + response = client.get(url="/users") +`, + language: "python", + }) + const lenses = provider.provideCodeLenses(doc) + assert.strictEqual(lenses.length, 1) + assert.ok(lenses[0].command?.title.includes("/users")) + }) + test("ignores calls with no arguments", async () => { const app = createMockApp([createRoute("GET", "/users")]) provider.setApps([app]) diff --git a/src/vscode/testCodeLensProvider.ts b/src/vscode/testCodeLensProvider.ts index a7c89c7..947629d 100644 --- a/src/vscode/testCodeLensProvider.ts +++ b/src/vscode/testCodeLensProvider.ts @@ -14,7 +14,11 @@ import { Uri, } from "vscode" import type { Node } from "web-tree-sitter" -import { extractPathFromNode, findNodesByType } from "../core/extractors" +import { + extractPathFromNode, + findNodesByType, + resolveArgNode, +} from "../core/extractors" import { ROUTE_METHODS } from "../core/internal" import type { Parser } from "../core/parser" import { @@ -134,7 +138,11 @@ export class TestCodeLensProvider implements CodeLensProvider { continue } - const pathArg = args[0] + const pathArg = resolveArgNode(args, 0, "url") + + if (!pathArg) { + continue + } // extractPathFromNode always returns a non-empty string for valid AST nodes const path = extractPathFromNode(pathArg) From 5891cc55b81541def221f3051794d68d3c705378 Mon Sep 17 00:00:00 2001 From: Savannah Ostrowski Date: Wed, 4 Mar 2026 09:44:25 -0800 Subject: [PATCH 5/6] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Simplify=20checks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/extractors.ts | 132 ++++++++++++++--------------- src/vscode/testCodeLensProvider.ts | 1 - 2 files changed, 65 insertions(+), 68 deletions(-) diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 9ddb4b2..4632e54 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -48,67 +48,6 @@ function stripDocstring(raw: string): string { const dedented = lines.map((l, i) => (i === 0 ? l : l.slice(minIndent))) return dedented.join("\n").trim() } -/* - * Extracts recognized FastAPI and APIRouter names from imports and class definitions - * This allows the rest of the extractors to handle user-defined aliases and subclasses. - * - * For example, if the code has: - * from fastapi import FastAPI as MyApp - * from fastapi import APIRouter as MyRouter - * class CustomRouter(MyRouter): ... - * - * Then this function will return: - * fastAPINames = Set { "FastAPI", "fastapi.FastAPI", "MyApp" } - * apiRouterNames = Set { "APIRouter", "fastapi.APIRouter", "MyRouter", "CustomRouter" } - * - * This allows decoratorExtractor and routerExtractor to recognize routes and routers - * defined using these aliases and subclasses without needing complex logic in those functions. - */ -export function collectRecognizedNames(rootNode: Node): { - fastAPINames: Set - apiRouterNames: Set -} { - const fastAPINames = new Set(["FastAPI", "fastapi.FastAPI"]) - const apiRouterNames = new Set(["APIRouter", "fastapi.APIRouter"]) - - // Add aliases from "from fastapi import X as Y" imports - for (const node of findNodesByType(rootNode, "import_from_statement")) { - const info = importExtractor(node) - if (!info || info.modulePath !== "fastapi") continue - for (const named of info.namedImports) { - if (named.alias === null) continue - if (named.name === "FastAPI") fastAPINames.add(named.alias) - else if (named.name === "APIRouter") apiRouterNames.add(named.alias) - } - } - - // Add module aliases from "import fastapi as f" → recognizes f.FastAPI, f.APIRouter - for (const node of findNodesByType(rootNode, "import_statement")) { - const info = importExtractor(node) - if (!info) continue - for (const named of info.namedImports) { - if (named.alias === null) continue - if (named.name === "fastapi") { - fastAPINames.add(`${named.alias}.FastAPI`) - apiRouterNames.add(`${named.alias}.APIRouter`) - } - } - } - - // Add subclasses, checking against the already-accumulated alias sets so - // "class MyRouter(AR)" works when AR is an alias for APIRouter - for (const cls of findNodesByType(rootNode, "class_definition")) { - const nameNode = cls.childForFieldName("name") - const superclassesNode = cls.childForFieldName("superclasses") - if (!nameNode || !superclassesNode) continue - for (const parent of superclassesNode.namedChildren) { - if (apiRouterNames.has(parent.text)) apiRouterNames.add(nameNode.text) - else if (fastAPINames.has(parent.text)) fastAPINames.add(nameNode.text) - } - } - - return { fastAPINames, apiRouterNames } -} function collectNodesByType(node: Node, type: string, results: Node[]): void { if (node.type === type) { @@ -316,15 +255,15 @@ export function routerExtractor( const funcName = valueNode.childForFieldName("function")?.text let type: RouterType if ( - funcName === "APIRouter" || - funcName === "fastapi.APIRouter" || - (funcName !== undefined && apiRouterNames?.has(funcName)) + funcName !== undefined && + (apiRouterNames?.has(funcName) ?? + (funcName === "APIRouter" || funcName === "fastapi.APIRouter")) ) { type = "APIRouter" } else if ( - funcName === "FastAPI" || - funcName === "fastapi.FastAPI" || - (funcName !== undefined && fastAPINames?.has(funcName)) + funcName !== undefined && + (fastAPINames?.has(funcName) ?? + (funcName === "FastAPI" || funcName === "fastapi.FastAPI")) ) { type = "FastAPI" } else { @@ -460,6 +399,65 @@ export function importExtractor(node: Node): ImportInfo | null { return { modulePath, names, namedImports, isRelative, relativeDots } } +/** + * Extracts recognized FastAPI and APIRouter names from imports and class definitions. + * This allows routerExtractor to handle user-defined aliases and subclasses. + * + * For example, if the code has: + * from fastapi import FastAPI as MyApp + * from fastapi import APIRouter as MyRouter + * class CustomRouter(MyRouter): ... + * + * Then this function will return: + * fastAPINames = Set { "FastAPI", "fastapi.FastAPI", "MyApp" } + * apiRouterNames = Set { "APIRouter", "fastapi.APIRouter", "MyRouter", "CustomRouter" } + */ +export function collectRecognizedNames(rootNode: Node): { + fastAPINames: Set + apiRouterNames: Set +} { + const fastAPINames = new Set(["FastAPI", "fastapi.FastAPI"]) + const apiRouterNames = new Set(["APIRouter", "fastapi.APIRouter"]) + + // Add aliases from "from fastapi import X as Y" imports + for (const node of findNodesByType(rootNode, "import_from_statement")) { + const info = importExtractor(node) + if (!info || info.modulePath !== "fastapi") continue + for (const named of info.namedImports) { + if (named.alias === null) continue + if (named.name === "FastAPI") fastAPINames.add(named.alias) + else if (named.name === "APIRouter") apiRouterNames.add(named.alias) + } + } + + // Add module aliases from "import fastapi as f" → recognizes f.FastAPI, f.APIRouter + for (const node of findNodesByType(rootNode, "import_statement")) { + const info = importExtractor(node) + if (!info) continue + for (const named of info.namedImports) { + if (named.alias === null) continue + if (named.name === "fastapi") { + fastAPINames.add(`${named.alias}.FastAPI`) + apiRouterNames.add(`${named.alias}.APIRouter`) + } + } + } + + // Add subclasses, checking against the already-accumulated alias sets so + // "class MyRouter(AR)" works when AR is an alias for APIRouter + for (const cls of findNodesByType(rootNode, "class_definition")) { + const nameNode = cls.childForFieldName("name") + const superclassesNode = cls.childForFieldName("superclasses") + if (!nameNode || !superclassesNode) continue + for (const parent of superclassesNode.namedChildren) { + if (apiRouterNames.has(parent.text)) apiRouterNames.add(nameNode.text) + else if (fastAPINames.has(parent.text)) fastAPINames.add(nameNode.text) + } + } + + return { fastAPINames, apiRouterNames } +} + /** * Resolves a function argument value node by positional index or keyword name. * diff --git a/src/vscode/testCodeLensProvider.ts b/src/vscode/testCodeLensProvider.ts index 947629d..c793a59 100644 --- a/src/vscode/testCodeLensProvider.ts +++ b/src/vscode/testCodeLensProvider.ts @@ -143,7 +143,6 @@ export class TestCodeLensProvider implements CodeLensProvider { if (!pathArg) { continue } - // extractPathFromNode always returns a non-empty string for valid AST nodes const path = extractPathFromNode(pathArg) calls.push({ From 9ab7caa4ed8b0b70feb5b897ecc755feffec3160 Mon Sep 17 00:00:00 2001 From: Savannah Ostrowski Date: Wed, 4 Mar 2026 09:51:17 -0800 Subject: [PATCH 6/6] =?UTF-8?q?=F0=9F=8E=A8=20Add=20module=20level=20guard?= =?UTF-8?q?=20for=20string=20variables?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/extractors.ts | 11 ++++++++ src/test/core/extractors.test.ts | 48 ++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 4632e54..3421472 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -65,16 +65,27 @@ function collectNodesByType(node: Node, type: string, results: Node[]): void { * Collects string variable assignments from the AST for path resolution. * Handles simple assignments like `WEBHOOK_PATH = "/webhook"`. * + * Only module-level assignments are collected — function/class-local variables + * are skipped to prevent shadowing module-level constants with the same name. + * * Examples: * WEBHOOK_PATH = "/webhook" -> Map { "WEBHOOK_PATH" => "/webhook" } * BASE = "/api" -> Map { "BASE" => "/api" } * settings.PREFIX = "/api" -> (skipped, not a simple identifier) + * def f(): BASE = "/local" -> (skipped, inside function) */ export function collectStringVariables(rootNode: Node): Map { const variables = new Map() const assignmentNodes = findNodesByType(rootNode, "assignment") for (const assign of assignmentNodes) { + if ( + hasAncestor(assign, "function_definition") || + hasAncestor(assign, "class_definition") + ) { + continue + } + const left = assign.childForFieldName("left") const right = assign.childForFieldName("right") if ( diff --git a/src/test/core/extractors.test.ts b/src/test/core/extractors.test.ts index 5b1c2db..1cc63db 100644 --- a/src/test/core/extractors.test.ts +++ b/src/test/core/extractors.test.ts @@ -1,6 +1,7 @@ import * as assert from "node:assert" import { collectRecognizedNames, + collectStringVariables, decoratorExtractor, extractPathFromNode, extractStringValue, @@ -863,6 +864,53 @@ router = f.APIRouter(prefix="/items") }) }) + suite("collectStringVariables", () => { + test("collects module-level string assignments", () => { + const code = ` +PREFIX = "/api" +VERSION = "/v1" +` + const tree = parse(code) + const vars = collectStringVariables(tree.rootNode) + assert.strictEqual(vars.get("PREFIX"), "/api") + assert.strictEqual(vars.get("VERSION"), "/v1") + }) + + test("ignores function-local variables", () => { + const code = ` +PREFIX = "/api" + +def handler(): + PREFIX = "/local" +` + const tree = parse(code) + const vars = collectStringVariables(tree.rootNode) + assert.strictEqual(vars.get("PREFIX"), "/api") + }) + + test("ignores class-level variables", () => { + const code = ` +PREFIX = "/api" + +class Config: + PREFIX = "/class-level" +` + const tree = parse(code) + const vars = collectStringVariables(tree.rootNode) + assert.strictEqual(vars.get("PREFIX"), "/api") + }) + + test("ignores non-string assignments", () => { + const code = ` +COUNT = 42 +FLAG = True +` + const tree = parse(code) + const vars = collectStringVariables(tree.rootNode) + assert.strictEqual(vars.size, 0) + }) + }) + suite("extractStringValue", () => { test("returns null for non-string node", () => { const code = "x = 42"