support aliased FastAPI/APIRouter imports

Author: Ashif4354

src/core/analyzer.ts

@@ -5,6 +5,7 @@
 import type { Tree } from "web-tree-sitter"
 import { logError } from "../utils/logger"
 import {
+  aliasExtractor,
   collectStringVariables,
   decoratorExtractor,
   findNodesByType,
@@ -42,9 +43,13 @@ export function analyzeTree(tree: Tree, filePath: string): FileAnalysis {
   const decoratedDefs = findNodesByType(rootNode, "decorated_definition")
   const routes = decoratedDefs.map(decoratorExtractor).filter(notNull)
 
+  const aliases = aliasExtractor(rootNode)
+
   // Get all router assignments
   const assignments = findNodesByType(rootNode, "assignment")
-  const routers = assignments.map(routerExtractor).filter(notNull)
+  const routers = assignments
+    .map((node) => routerExtractor(node, aliases))
+    .filter(notNull)
 
   // Get all include_router and mount calls
   const callNodes = findNodesByType(rootNode, "call")

src/core/extractors.ts

@@ -237,7 +237,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,
+  aliasMap?: Map<string, "FastAPI" | "APIRouter">,
+): RouterInfo | null {
   if (node.type !== "assignment") {
     return null
   }
@@ -254,6 +257,8 @@ export function routerExtractor(node: Node): RouterInfo | null {
     type = "APIRouter"
   } else if (funcName === "FastAPI" || funcName === "fastapi.FastAPI") {
     type = "FastAPI"
+  } else if (funcName && aliasMap?.has(funcName)) {
+    type = aliasMap.get(funcName)!
   } else {
     return null
   }
@@ -476,3 +481,36 @@ export function mountExtractor(node: Node): MountInfo | null {
     app: appNode?.text ?? "",
   }
 }
+
+export function aliasExtractor(
+  node: Node,
+): Map<string, "FastAPI" | "APIRouter"> {
+  const importAliases = findNodesByType(node, "aliased_import")
+
+  const aliasMap = new Map<string, "FastAPI" | "APIRouter">()
+
+  for (const node of importAliases) {
+    const originalNode = node.childForFieldName("name")
+    const aliasNode = node.childForFieldName("alias")
+
+    if (
+      originalNode?.text === "FastAPI" ||
+      originalNode?.text === "fastapi.FastAPI"
+    ) {
+      const alias = aliasNode?.text
+      if (alias) {
+        aliasMap.set(alias, "FastAPI")
+      }
+    } else if (
+      originalNode?.text === "APIRouter" ||
+      originalNode?.text === "fastapi.APIRouter"
+    ) {
+      const alias = aliasNode?.text
+      if (alias) {
+        aliasMap.set(alias, "APIRouter")
+      }
+    }
+  }
+
+  return aliasMap
+}