🐛 Discover routes from imported custom APIRouter instances

src/core/analyzer.ts

@@ -5,10 +5,10 @@
 import type { Tree } from "web-tree-sitter"
 import { logError } from "../utils/logger"
 import {
+  callAssignmentExtractor,
   collectRecognizedNames,
   collectStringVariables,
   decoratorExtractor,
-  factoryCallExtractor,
   getNodesByType,
   importExtractor,
   includeRouterExtractor,
@@ -44,12 +44,11 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis {
   const decoratedDefs = nodesByType.get("decorated_definition") ?? []
   const routes = decoratedDefs.flatMap(decoratorExtractor)
 
-  // Get all router assignments
+  // Get module-level call assignments and recognized router assignments
   const assignments = nodesByType.get("assignment") ?? []
   const { fastAPINames, apiRouterNames } = collectRecognizedNames(nodesByType)
-  const knownConstructors = new Set([...fastAPINames, ...apiRouterNames])
-  const factoryCalls = assignments
-    .map((node) => factoryCallExtractor(node, knownConstructors))
+  const callAssignments = assignments
+    .map(callAssignmentExtractor)
     .filter(notNull)
   const routers = assignments
     .map((node) => routerExtractor(node, apiRouterNames, fastAPINames))
@@ -69,17 +68,11 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis {
 
   const stringVariables = collectStringVariables(nodesByType)
 
-  for (const route of routes) {
-    route.path = resolveVariables(route.path, stringVariables)
+  for (const item of [...routes, ...mounts]) {
+    item.path = resolveVariables(item.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)
+  for (const item of [...routers, ...callAssignments, ...includeRouters]) {
+    item.prefix = resolveVariables(item.prefix, stringVariables)
   }
 
   return {
@@ -89,7 +82,7 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis {
     includeRouters,
     mounts,
     imports,
-    factoryCalls,
+    callAssignments,
   }
 }
 

src/core/extractors.ts

@@ -4,7 +4,7 @@
 
 import type { Node } from "web-tree-sitter"
 import type {
-  FactoryCallInfo,
+  CallAssignmentInfo,
   ImportedName,
   ImportInfo,
   IncludeRouterInfo,
@@ -265,6 +265,30 @@ function extractTags(listNode: Node): string[] {
     .filter((v): v is string => v !== null)
 }
 
+function extractRouterCallMetadata(callNode: Node): {
+  prefix: string
+  tags: string[]
+} {
+  let prefix = ""
+  let tags: string[] = []
+  const argumentsNode = callNode.childForFieldName("arguments")
+  for (const child of argumentsNode?.namedChildren ?? []) {
+    if (child.type !== "keyword_argument") {
+      continue
+    }
+    const argName = child.childForFieldName("name")?.text
+    const argValue = child.childForFieldName("value")
+
+    if (argName === "prefix" && argValue) {
+      prefix = extractPathFromNode(argValue)
+    } else if (argName === "tags" && argValue?.type === "list") {
+      tags = extractTags(argValue)
+    }
+  }
+
+  return { prefix, tags }
+}
+
 export function routerExtractor(
   node: Node,
   apiRouterNames?: Set<string>,
@@ -298,22 +322,7 @@ export function routerExtractor(
     return null
   }
 
-  let prefix = ""
-  let tags: string[] = []
-  const argumentsNode = valueNode.childForFieldName("arguments")
-  for (const child of argumentsNode?.namedChildren ?? []) {
-    if (child.type !== "keyword_argument") {
-      continue
-    }
-    const argName = child.childForFieldName("name")?.text
-    const argValue = child.childForFieldName("value")
-
-    if (argName === "prefix" && argValue) {
-      prefix = extractPathFromNode(argValue)
-    } else if (argName === "tags" && argValue?.type === "list") {
-      tags = extractTags(argValue)
-    }
-  }
+  const { prefix, tags } = extractRouterCallMetadata(valueNode)
 
   return {
     variableName: variableNameNode.text,
@@ -601,10 +610,7 @@ export function mountExtractor(node: Node): MountInfo | null {
   }
 }
 
-export function factoryCallExtractor(
-  node: Node,
-  knownConstructors: Set<string>,
-): FactoryCallInfo | null {
+export function callAssignmentExtractor(node: Node): CallAssignmentInfo | null {
   if (node.type !== "assignment") {
     return null
   }
@@ -620,11 +626,6 @@ export function factoryCallExtractor(
     return null
   }
 
-  const functionName = functionNode.text
-  if (knownConstructors.has(functionName)) {
-    return null
-  }
-
   // Skip function and class-local variables to avoid false positives
   if (
     hasAncestor(node, "function_definition") ||
@@ -633,9 +634,15 @@ export function factoryCallExtractor(
     return null
   }
 
+  const { prefix, tags } = extractRouterCallMetadata(valueNode)
+
   return {
     variableName: variableNameNode.text,
-    functionName: functionName,
+    callee: functionNode.text,
+    prefix,
+    tags,
+    line: node.startPosition.row + 1,
+    column: node.startPosition.column,
   }
 }
 

src/core/internal.ts

@@ -78,9 +78,13 @@ export interface MountInfo {
   app: string
 }
 
-export interface FactoryCallInfo {
+export interface CallAssignmentInfo {
   variableName: string
-  functionName: string
+  callee: string
+  prefix: string
+  tags: string[]
+  line: number
+  column: number
 }
 
 export interface FileAnalysis {
@@ -90,7 +94,7 @@ export interface FileAnalysis {
   includeRouters: IncludeRouterInfo[]
   mounts: MountInfo[]
   imports: ImportInfo[]
-  factoryCalls: FactoryCallInfo[]
+  callAssignments: CallAssignmentInfo[]
 }
 
 export interface RouterNode {

src/core/routerResolver.ts

@@ -19,6 +19,17 @@ interface ResolutionContext {
   visited: Set<string>
 }
 
+interface ResolveReferenceOptions {
+  reference: string
+  analysis: FileAnalysis
+  currentFileUri: string
+  ctx: ResolutionContext
+}
+
+interface InternalResolveReferenceOptions extends ResolveReferenceOptions {
+  kind: "includedRouter" | "mountedApp"
+}
+
 /**
  * Finds the main FastAPI app or APIRouter in the list of routers.
  * If targetVariable is specified, only returns the router with that variable name.
@@ -37,6 +48,27 @@ function findAppRouter(
   )
 }
 
+function inferIncludedRouter(
+  analysis: FileAnalysis,
+  variableName: string,
+): RouterInfo | undefined {
+  const assignment = analysis.callAssignments.find(
+    (candidate) => candidate.variableName === variableName,
+  )
+  if (!assignment) {
+    return undefined
+  }
+
+  return {
+    variableName: assignment.variableName,
+    type: "APIRouter",
+    prefix: assignment.prefix,
+    tags: assignment.tags,
+    line: assignment.line,
+    column: assignment.column,
+  }
+}
+
 function createRouterNode(
   router: RouterInfo,
   routes: RouteInfo[],
@@ -55,6 +87,14 @@ function createRouterNode(
   }
 }
 
+function resolveIncludedRouter(options: ResolveReferenceOptions) {
+  return resolveReference({ ...options, kind: "includedRouter" })
+}
+
+function resolveMountedApp(options: ResolveReferenceOptions) {
+  return resolveReference({ ...options, kind: "mountedApp" })
+}
+
 async function processIncludeRouters(
   analysis: FileAnalysis,
   ownerRouter: RouterNode,
@@ -68,12 +108,12 @@ async function processIncludeRouters(
     log(
       `Resolving include_router: ${include.router} (prefix: ${include.prefix || "none"})`,
     )
-    const childRouter = await resolveRouterReference(
-      include.router,
+    const childRouter = await resolveIncludedRouter({
+      reference: include.router,
       analysis,
       currentFileUri,
       ctx,
-    )
+    })
     if (childRouter) {
       // Merge tags from include_router call with the router's own tags
       if (include.tags.length > 0) {
@@ -198,18 +238,18 @@ async function buildRouterGraphInternal(
   // `app = FastAPI()` and `app.include_router(...)` inside `create_app` are visible
   // when analyzing the factory file directly.
   if (!appRouter && targetVariable) {
-    const factoryCall = analysis.factoryCalls.find(
-      (fc) => fc.variableName === targetVariable,
+    const callAssignment = analysis.callAssignments.find(
+      (assignment) => assignment.variableName === targetVariable,
     )
-    if (factoryCall) {
+    if (callAssignment) {
       const matchingImport = analysis.imports.find((imp) =>
-        imp.names.includes(factoryCall.functionName),
+        imp.names.includes(callAssignment.callee),
       )
       if (matchingImport) {
         const namedImport = matchingImport.namedImports.find(
-          (ni) => (ni.alias ?? ni.name) === factoryCall.functionName,
+          (ni) => (ni.alias ?? ni.name) === callAssignment.callee,
         )
-        const originalName = namedImport?.name ?? factoryCall.functionName
+        const originalName = namedImport?.name ?? callAssignment.callee
         const factoryFileUri = await resolveNamedImport(
           {
             modulePath: matchingImport.modulePath,
@@ -252,12 +292,12 @@ async function buildRouterGraphInternal(
 
   // Process mount() calls for subapps
   for (const mount of analysis.mounts) {
-    const childRouter = await resolveRouterReference(
-      mount.app,
+    const childRouter = await resolveMountedApp({
+      reference: mount.app,
       analysis,
-      resolvedEntryUri,
+      currentFileUri: resolvedEntryUri,
       ctx,
-    )
+    })
     if (childRouter) {
       rootRouter.children.push({
         router: childRouter,
@@ -277,12 +317,13 @@ async function buildRouterGraphInternal(
  * Handles both simple references (e.g., "router") and dotted references
  * (e.g., "api_routes.router" where api_routes is an imported module).
  */
-async function resolveRouterReference(
-  reference: string,
-  analysis: FileAnalysis,
-  currentFileUri: string,
-  ctx: ResolutionContext,
-): Promise<RouterNode | null> {
+async function resolveReference({
+  reference,
+  analysis,
+  currentFileUri,
+  ctx,
+  kind,
+}: InternalResolveReferenceOptions): Promise<RouterNode | null> {
   const { projectRootUri, parser, fs, visited } = ctx
   const parts = reference.split(".")
   const moduleName = parts[0]
@@ -358,9 +399,13 @@ async function resolveRouterReference(
     }
 
     // Find the router with the matching variable name
-    const targetRouter = importedAnalysis.routers.find(
-      (r) => r.variableName === attributeName,
-    )
+    const targetRouter =
+      importedAnalysis.routers.find(
+        (router) => router.variableName === attributeName,
+      ) ??
+      (kind === "includedRouter"
+        ? inferIncludedRouter(importedAnalysis, attributeName)
+        : undefined)
     if (targetRouter) {
       const visitedKey = `${importedFileUri}#${attributeName}`
 

src/test/core/analyzer.test.ts

@@ -63,6 +63,31 @@ router = APIRouter(prefix="/api")
       assert.strictEqual(result.routers[1].prefix, "/api")
     })
 
+    test("records module-level direct call assignments as neutral facts", () => {
+      const code = `
+from auth.testing_router import ProtectedRouter
+
+def build_router():
+    local_router = ProtectedRouter()
+
+router = ProtectedRouter(prefix="/protected", tags=["protected"])
+`
+      const tree = parse(code)
+      const result = analyzeTree(tree, "/test/file.py")
+
+      assert.strictEqual(result.routers.length, 0)
+      assert.deepStrictEqual(result.callAssignments, [
+        {
+          variableName: "router",
+          callee: "ProtectedRouter",
+          prefix: "/protected",
+          tags: ["protected"],
+          line: 7,
+          column: 0,
+        },
+      ])
+    })
+
     test("extracts include_router calls", () => {
       const code = `
 app.include_router(users.router, prefix="/users")