🐛 Handle module aliase

Author: savannahostrowski

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

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<string> {
-  const subclasses = new Set<string>()
-  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<string>
-  apiRouterAliases: Set<string>
+export function collectRecognizedNames(rootNode: Node): {
+  fastAPINames: Set<string>
+  apiRouterNames: Set<string>
 } {
-  const fastAPIAliases = new Set<string>()
-  const apiRouterAliases = new Set<string>()
+  const fastAPINames = new Set<string>(["FastAPI", "fastapi.FastAPI"])
+  const apiRouterNames = new Set<string>(["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<string>,
-  fastAPIAliases?: Set<string>,
+  apiRouterNames?: Set<string>,
+  fastAPINames?: Set<string>,
 ): 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,

src/test/core/extractors.test.ts

@@ -1,7 +1,6 @@
 import * as assert from "node:assert"
 import {
-  collectAPIRouterSubclasses,
-  collectFastAPIAliases,
+  collectRecognizedNames,
   decoratorExtractor,
   extractPathFromNode,
   extractStringValue,
@@ -477,44 +476,50 @@ 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
 
 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")
       assert.strictEqual(result.type, "APIRouter")
       assert.strictEqual(result.prefix, "/admin")
     })
 
-    test("does not treat FastAPI subclass as APIRouter", () => {
+    test("recognizes FastAPI subclass", () => {
       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 { 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,26 +553,83 @@ 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")
       assert.strictEqual(result.type, "APIRouter")
       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)
     })