diff --git a/src/core/analyzer.ts b/src/core/analyzer.ts index 2b6b4fe..9e1f927 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 { + collectStringVariables, decoratorExtractor, findNodesByType, importExtractor, @@ -20,6 +21,16 @@ function notNull(value: T | null): value is T { return value !== null } +function resolveVariables( + path: string, + variables: Map, +): string { + return path.replace( + /\{([^}]+)\}/g, + (match, name) => variables.get(name) ?? match, + ) +} + /** Analyze a syntax tree and extract FastAPI-related information */ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis { const rootNode = tree.rootNode @@ -44,7 +55,29 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis { .map(importExtractor) .filter(notNull) - return { filePath, routes, routers, includeRouters, mounts, imports } + const stringVariables = collectStringVariables(rootNode) + + for (const route of routes) { + route.path = resolveVariables(route.path, stringVariables) + } + for (const router of routers) { + router.prefix = resolveVariables(router.prefix, stringVariables) + } + for (const ir of includeRouters) { + ir.prefix = resolveVariables(ir.prefix, stringVariables) + } + for (const mount of mounts) { + mount.path = resolveVariables(mount.path, stringVariables) + } + + return { + filePath, + routes, + routers, + includeRouters, + mounts, + imports, + } } /** Analyze a file given its URI string and a parser instance */ diff --git a/src/core/extractors.ts b/src/core/extractors.ts index 5e0fb34..87c3b43 100644 --- a/src/core/extractors.ts +++ b/src/core/extractors.ts @@ -33,6 +33,39 @@ function collectNodesByType(node: Node, type: string, results: Node[]): void { } } +/** + * Collects string variable assignments from the AST for path resolution. + * Only resolves simple assignments (e.g. `WEBHOOK_PATH = "/webhook"`). + * + * Examples: + * WEBHOOK_PATH = "/webhook" -> Map { "WEBHOOK_PATH" => "/webhook" } + * BASE = "/api" -> Map { "BASE" => "/api" } + * settings.PREFIX = "/api" -> (skipped, not a simple identifier) + */ +export function collectStringVariables(rootNode: Node): Map { + const variables = new Map() + const assignmentNodes = findNodesByType(rootNode, "assignment") + + for (const assign of assignmentNodes) { + const left = assign.childForFieldName("left") + const right = assign.childForFieldName("right") + if ( + left && + right && + left.type === "identifier" && + right.type === "string" + ) { + const varName = left.text + const value = extractStringValue(right) + if (value !== null) { + variables.set(varName, value) + } + } + } + + return variables +} + /** * Extracts the string value from a string AST node, handling quotes and f-string prefix. * Returns null if the node is not a string. diff --git a/src/test/core/analyzer.test.ts b/src/test/core/analyzer.test.ts index 46c2876..16ad751 100644 --- a/src/test/core/analyzer.test.ts +++ b/src/test/core/analyzer.test.ts @@ -100,6 +100,88 @@ import os assert.strictEqual(routesImport.isRelative, true) }) + test("resolves same-file string variables in route paths", () => { + const code = ` +from fastapi import FastAPI + +app = FastAPI() + +WEBHOOK_PATH = "/webhook" + +@app.post(WEBHOOK_PATH) +def some_webhook(): + pass +` + const tree = parse(code) + const result = analyzeTree(tree, "/test/file.py") + + assert.strictEqual(result.routes.length, 1) + assert.strictEqual(result.routes[0].path, "/webhook") + }) + + test("resolves variable used in path concatenation", () => { + const code = ` +from fastapi import FastAPI + +app = FastAPI() + +BASE = "/api" + +@app.get(BASE + "/users") +def list_users(): + pass +` + const tree = parse(code) + const result = analyzeTree(tree, "/test/file.py") + + assert.strictEqual(result.routes.length, 1) + assert.strictEqual(result.routes[0].path, "/api/users") + }) + + test("leaves unresolvable variables wrapped", () => { + const code = ` +from fastapi import FastAPI + +app = FastAPI() + +@app.get(settings.API_PREFIX) +def handler(): + pass +` + const tree = parse(code) + const result = analyzeTree(tree, "/test/file.py") + + assert.strictEqual(result.routes.length, 1) + assert.strictEqual(result.routes[0].path, "{settings.API_PREFIX}") + }) + + test("resolves variable in router prefix", () => { + const code = ` +from fastapi import APIRouter + +PREFIX = "/users" +router = APIRouter(prefix=PREFIX) +` + const tree = parse(code) + const result = analyzeTree(tree, "/test/file.py") + + const apiRouter = result.routers.find((r) => r.type === "APIRouter") + assert.ok(apiRouter) + assert.strictEqual(apiRouter.prefix, "/users") + }) + + test("resolves variable in include_router prefix", () => { + const code = ` +USERS_PREFIX = "/users" +app.include_router(users.router, prefix=USERS_PREFIX) +` + const tree = parse(code) + const result = analyzeTree(tree, "/test/file.py") + + assert.strictEqual(result.includeRouters.length, 1) + assert.strictEqual(result.includeRouters[0].prefix, "/users") + }) + test("sets filePath correctly", () => { const code = "x = 1" const tree = parse(code) diff --git a/src/test/providers/pathOperationTreeProvider.test.ts b/src/test/providers/pathOperationTreeProvider.test.ts index a044fce..9158bb6 100644 --- a/src/test/providers/pathOperationTreeProvider.test.ts +++ b/src/test/providers/pathOperationTreeProvider.test.ts @@ -334,6 +334,34 @@ suite("PathOperationTreeProvider", () => { } }) + test("route tooltip strips dynamic prefix", () => { + const app = makeApp("app", "main.py") + app.routes = [ + { + method: "GET", + path: "{settings.API_V1_STR}/users/{user_id}", + functionName: "get_user", + location: { filePath: "users.py", line: 10, column: 0 }, + }, + ] + const p = new PathOperationTreeProvider([app]) + const appItem = p.getChildren()[0] + const route = p.getChildren(appItem).find((c) => c.type === "route")! + const treeItem = p.getTreeItem(route) + const tooltipValue = + typeof treeItem.tooltip === "string" + ? treeItem.tooltip + : (treeItem.tooltip as { value: string }).value + assert.ok( + tooltipValue.includes("GET /users/{user_id}"), + "Tooltip should show stripped path", + ) + assert.ok( + !tooltipValue.includes("settings.API_V1_STR"), + "Tooltip should not include dynamic prefix", + ) + }) + test("getChildren returns message when no apps", () => { const emptyProvider = new PathOperationTreeProvider([]) const roots = emptyProvider.getChildren() diff --git a/src/vscode/pathOperationTreeProvider.ts b/src/vscode/pathOperationTreeProvider.ts index 134d465..b06575a 100644 --- a/src/vscode/pathOperationTreeProvider.ts +++ b/src/vscode/pathOperationTreeProvider.ts @@ -268,8 +268,9 @@ export class PathOperationTreeProvider routeItem.description = element.route.functionName routeItem.iconPath = new ThemeIcon(METHOD_ICONS[element.route.method]) routeItem.contextValue = "route" + const tooltipPath = stripLeadingDynamicSegments(element.route.path) routeItem.tooltip = new MarkdownString( - `${element.route.method} ${element.route.path}\n\nFunction: ${element.route.functionName}\nFile: ${element.route.location.filePath}:${element.route.location.line}`, + `${element.route.method} ${tooltipPath}\n\nFunction: ${element.route.functionName}\nFile: ${element.route.location.filePath}:${element.route.location.line}`, ) routeItem.command = { command: "fastapi-vscode.goToPathOperation",