♻️ Refactor tree utils and EndpointTreeProvider (#29)

Author: savannahostrowski

src/appDiscovery.ts

@@ -10,15 +10,11 @@ import type { Parser } from "./core/parser"
 import { findProjectRoot, uriPath } from "./core/pathUtils"
 import { buildRouterGraph } from "./core/routerResolver"
 import { routerNodeToAppDefinition } from "./core/transformer"
+import { collectRoutes, countRouters } from "./core/treeUtils"
 import type { AppDefinition } from "./core/types"
 import { vscodeFileSystem } from "./providers/vscodeFileSystem"
 import { log } from "./utils/logger"
-import {
-  countRouters,
-  countRoutes,
-  createTimer,
-  trackEntrypointDetected,
-} from "./utils/telemetry"
+import { createTimer, trackEntrypointDetected } from "./utils/telemetry"
 
 export type { EntryPoint }
 
@@ -176,28 +172,27 @@ export async function discoverFastAPIApps(
 
       if (routerNode) {
         const app = routerNodeToAppDefinition(routerNode, folder.uri.fsPath)
-        // Count all routes: direct routes + routes in all routers (recursively)
-        const countRoutes = (routers: typeof app.routers): number =>
-          routers.reduce(
-            (sum, r) => sum + r.routes.length + countRoutes(r.children),
-            0,
-          )
-        const totalRoutes = app.routes.length + countRoutes(app.routers)
-        log(
-          `Found FastAPI app "${app.name}" with ${totalRoutes} route(s) in ${app.routers.length} router(s)`,
-        )
         folderApps.push(app)
         apps.push(app)
         break // TODO: Only use first successful app per workspace folder, for now
       }
     }
 
+    const folderRoutes = collectRoutes(folderApps)
+
+    if (folderApps.length > 0) {
+      const app = folderApps[0]
+      log(
+        `Found FastAPI app "${app.name}" with ${folderRoutes.length} route(s) in ${app.routers.length} router(s)`,
+      )
+    }
+
     // Track entrypoint detection per workspace folder
     trackEntrypointDetected({
       duration_ms: folderTimer(),
       method: detectionMethod,
       success: folderApps.length > 0,
-      routes_count: countRoutes(folderApps),
+      routes_count: folderRoutes.length,
       routers_count: countRouters(folderApps),
     })
   }

src/core/index.ts

@@ -11,6 +11,12 @@ export { Parser } from "./parser"
 export { findProjectRoot } from "./pathUtils"
 export { buildRouterGraph, type RouterNode } from "./routerResolver"
 export { routerNodeToAppDefinition } from "./transformer"
+export {
+  collectRoutes,
+  countRouters,
+  countRoutesInRouter,
+  findRouter,
+} from "./treeUtils"
 export type {
   AppDefinition,
   HTTPMethod,

src/core/treeUtils.ts

@@ -0,0 +1,45 @@
+import type { AppDefinition, RouteDefinition, RouterDefinition } from "./types"
+
+/** Finds the first router matching a condition across all apps. */
+export function findRouter(
+  apps: AppDefinition[],
+  predicate: (router: RouterDefinition) => boolean,
+): RouterDefinition | undefined {
+  function findIn(routers: RouterDefinition[]): RouterDefinition | undefined {
+    for (const r of routers) {
+      if (predicate(r)) return r
+      const found = findIn(r.children)
+      if (found) return found
+    }
+    return undefined
+  }
+  for (const app of apps) {
+    const found = findIn(app.routers)
+    if (found) return found
+  }
+  return undefined
+}
+
+/** Collects all routes from all apps, including nested router routes. */
+export function collectRoutes(apps: AppDefinition[]): RouteDefinition[] {
+  function fromRouters(routers: RouterDefinition[]): RouteDefinition[] {
+    return routers.flatMap((r) => [...r.routes, ...fromRouters(r.children)])
+  }
+  return apps.flatMap((app) => [...app.routes, ...fromRouters(app.routers)])
+}
+
+/** Counts all routes in a router, including nested routers. */
+export function countRoutesInRouter(router: RouterDefinition): number {
+  return (
+    router.routes.length +
+    router.children.reduce((sum, child) => sum + countRoutesInRouter(child), 0)
+  )
+}
+
+/** Counts total routers across all apps. */
+export function countRouters(apps: AppDefinition[]): number {
+  function count(routers: RouterDefinition[]): number {
+    return routers.reduce((sum, r) => sum + 1 + count(r.children), 0)
+  }
+  return apps.reduce((sum, app) => sum + count(app.routers), 0)
+}

src/extension.ts

@@ -7,6 +7,7 @@ import { discoverFastAPIApps } from "./appDiscovery"
 import { clearImportCache } from "./core/importResolver"
 import { Parser } from "./core/parser"
 import { stripLeadingDynamicSegments } from "./core/pathUtils"
+import { collectRoutes, countRouters } from "./core/treeUtils"
 import type { AppDefinition, SourceLocation } from "./core/types"
 import {
   type EndpointTreeItem,
@@ -16,8 +17,6 @@ import {
 import { TestCodeLensProvider } from "./providers/testCodeLensProvider"
 import { disposeLogger, log } from "./utils/logger"
 import {
-  countRouters,
-  countRoutes,
   createTimer,
   flushSessionSummary,
   getInstalledVersions,
@@ -110,7 +109,7 @@ export async function activate(context: vscode.ExtensionContext) {
   trackActivation({
     duration_ms: elapsed(),
     success,
-    routes_count: countRoutes(apps),
+    routes_count: collectRoutes(apps).length,
     routers_count: countRouters(apps),
     apps_count: apps.length,
     workspace_folder_count: vscode.workspace.workspaceFolders?.length ?? 0,
@@ -139,7 +138,7 @@ export async function activate(context: vscode.ExtensionContext) {
     })
   }
 
-  const endpointProvider = new EndpointTreeProvider(apps, groupApps)
+  const endpointProvider = new EndpointTreeProvider(apps, groupApps(apps))
   const codeLensProvider = new TestCodeLensProvider(parserService, apps)
 
   // File watcher for auto-refresh
@@ -149,7 +148,7 @@ export async function activate(context: vscode.ExtensionContext) {
     refreshTimeout = setTimeout(async () => {
       if (!parserService) return
       const newApps = await discoverFastAPIApps(parserService)
-      endpointProvider.setApps(newApps, groupApps)
+      endpointProvider.setApps(newApps, groupApps(newApps))
       codeLensProvider.setApps(newApps)
     }, 300)
   }
@@ -215,7 +214,7 @@ function registerCommands(
         if (!parserService) return
         clearImportCache()
         const newApps = await discoverFastAPIApps(parserService)
-        endpointProvider.setApps(newApps, groupApps)
+        endpointProvider.setApps(newApps, groupApps(newApps))
         codeLensProvider.setApps(newApps)
       },
     ),
@@ -235,8 +234,7 @@ function registerCommands(
       async () => {
         const workspacePrefix =
           vscode.workspace.workspaceFolders?.[0]?.uri.fsPath ?? ""
-        const items = endpointProvider
-          .getAllRoutes()
+        const items = collectRoutes(endpointProvider.getApps())
           .map((route) => {
             const path = stripLeadingDynamicSegments(route.path)
             return {