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
39 changes: 38 additions & 1 deletion src/core/extractors.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,34 @@ export function findNodesByType(node: Node, type: string): Node[] {
return results
}

function stripDocstring(raw: string): string {
let content: string
if (
(raw.startsWith('"""') && raw.endsWith('"""')) ||
(raw.startsWith("'''") && raw.endsWith("'''"))
) {
content = raw.slice(3, -3)
} else {
content = raw.slice(1, -1)
}

// Dedent: strip common leading whitespace (like Python's textwrap.dedent)
const lines = content.split("\n")
// First line is either empty or unindented (follows opening quotes), so skip it
const indentedLines = lines.slice(1).filter((l) => l.trim().length > 0)
if (indentedLines.length === 0) {
return content.trim()
}

// Find minimum indentation of all non-empty lines (except first) so we can
// remove it from all lines, preserving relative indentation
const minIndent = Math.min(
...indentedLines.map((l) => l.length - l.trimStart().length),
)
const dedented = lines.map((l, i) => (i === 0 ? l : l.slice(minIndent)))
return dedented.join("\n").trim()
}

function collectNodesByType(node: Node, type: string, results: Node[]): void {
if (node.type === type) {
results.push(node)
Expand Down Expand Up @@ -179,14 +207,23 @@ export function decoratorExtractor(node: Node): RouteInfo | null {
// Grammar guarantees: decorated_definition always has a definition field with a name
const functionDefNode = node.childForFieldName("definition")!
const functionName = functionDefNode.childForFieldName("name")?.text ?? ""

const functionBody = functionDefNode.childForFieldName("body")
const firstStatement = functionBody?.namedChildren[0]
let docstring: string | undefined
if (firstStatement?.type === "expression_statement") {
const expr = firstStatement.firstNamedChild
if (expr?.type === "string") {
docstring = stripDocstring(expr.text)
}
}
return {
owner: objectNode.text,
method: resolvedMethod,
path,
function: functionName,
line: node.startPosition.row + 1,
column: node.startPosition.column,
docstring,
}
}

Expand Down
2 changes: 2 additions & 0 deletions src/core/internal.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ export interface RouteInfo {
function: string
line: number
column: number
docstring?: string
}

export type RouterType = "APIRouter" | "FastAPI" | "Unknown"
Expand Down Expand Up @@ -99,6 +100,7 @@ export interface RouterNode {
function: string
line: number
column: number
docstring?: string
}[]
children: { router: RouterNode; prefix: string; tags: string[] }[]
}
Expand Down
3 changes: 3 additions & 0 deletions src/core/routerResolver.ts
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ async function buildRouterGraphInternal(
function: r.function,
line: r.line,
column: r.column,
docstring: r.docstring,
})),
children: [],
}
Expand Down Expand Up @@ -232,6 +233,7 @@ async function resolveRouterReference(
function: r.function,
line: r.line,
column: r.column,
docstring: r.docstring,
})),
children: [],
}
Expand Down Expand Up @@ -348,6 +350,7 @@ async function resolveRouterReference(
function: r.function,
line: r.line,
column: r.column,
docstring: r.docstring,
})),
children: [],
}
Expand Down
1 change: 1 addition & 0 deletions src/core/transformer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ function toRouteDefinition(
method: normalizeMethod(route.method),
path: prefix + route.path,
functionName: route.function,
docstring: route.docstring,
location: {
filePath,
line: route.line,
Expand Down
1 change: 1 addition & 0 deletions src/core/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ export interface RouteDefinition {
path: string
functionName: string
location: SourceLocation
docstring?: string
}

export interface RouterDefinition {
Expand Down
78 changes: 78 additions & 0 deletions src/test/core/extractors.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,84 @@ def handler():
assert.strictEqual(result.line, 2) // 1-indexed
assert.strictEqual(result.column, 0)
})

test("extracts single-line docstring", () => {
const code = `
@router.get("/users")
def list_users():
"""List all users."""
pass
`
const tree = parse(code)
const decoratedDefs = findNodesByType(
tree.rootNode,
"decorated_definition",
)
const result = decoratorExtractor(decoratedDefs[0])

assert.ok(result)
assert.strictEqual(result.docstring, "List all users.")
})

test("extracts multi-line docstring and dedents", () => {
const code = `
@router.get("/users")
def list_users():
"""
List all users.

Returns a list of user objects.
"""
pass
`
const tree = parse(code)
const decoratedDefs = findNodesByType(
tree.rootNode,
"decorated_definition",
)
const result = decoratorExtractor(decoratedDefs[0])

assert.ok(result)
assert.strictEqual(
result.docstring,
"List all users.\n\nReturns a list of user objects.",
)
})

test("extracts single-quote docstring", () => {
const code = `
@router.get("/users")
def list_users():
'''List all users.'''
pass
`
const tree = parse(code)
const decoratedDefs = findNodesByType(
tree.rootNode,
"decorated_definition",
)
const result = decoratorExtractor(decoratedDefs[0])

assert.ok(result)
assert.strictEqual(result.docstring, "List all users.")
})

test("returns undefined docstring when none present", () => {
const code = `
@router.get("/users")
def list_users():
pass
`
const tree = parse(code)
const decoratedDefs = findNodesByType(
tree.rootNode,
"decorated_definition",
)
const result = decoratorExtractor(decoratedDefs[0])

assert.ok(result)
assert.strictEqual(result.docstring, undefined)
})
})

suite("routerExtractor", () => {
Expand Down
2 changes: 1 addition & 1 deletion src/test/providers/pathOperationTreeProvider.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,7 @@ suite("PathOperationTreeProvider", () => {
? treeItem.tooltip
: (treeItem.tooltip as { value: string }).value
assert.ok(
tooltipValue.includes("GET /users/{user_id}"),
tooltipValue.includes("/users/{user_id}"),
"Tooltip should show stripped path",
)
assert.ok(
Expand Down
9 changes: 8 additions & 1 deletion src/vscode/pathOperationTreeProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import {
type TreeDataProvider,
TreeItem,
TreeItemCollapsibleState,
Uri,
} from "vscode"
import { stripLeadingDynamicSegments } from "../core/pathUtils"
import { countRoutesInRouter, findRouter } from "../core/treeUtils"
Expand Down Expand Up @@ -269,8 +270,14 @@ export class PathOperationTreeProvider
routeItem.iconPath = new ThemeIcon(METHOD_ICONS[element.route.method])
routeItem.contextValue = "route"
const tooltipPath = stripLeadingDynamicSegments(element.route.path)
const docstringSection = element.route.docstring
? `\n\n---\n\n${element.route.docstring}`
: ""
routeItem.tooltip = new MarkdownString(
`${element.route.method} ${tooltipPath}\n\nFunction: ${element.route.functionName}\nFile: ${element.route.location.filePath}:${element.route.location.line}`,
`**${element.route.method}** \`${tooltipPath}\`\n\n` +
`**Function:** \`${element.route.functionName}\`\n\n` +
`**File:** ${Uri.parse(element.route.location.filePath).fsPath}:${element.route.location.line}` +
docstringSection,
)
routeItem.command = {
command: "fastapi-vscode.goToPathOperation",
Expand Down
Loading