♻️ Refactor endpoint tree provider

Author: savannahostrowski

src/appDiscovery.ts

@@ -172,22 +172,27 @@ export async function discoverFastAPIApps(
 
       if (routerNode) {
         const app = routerNodeToAppDefinition(routerNode, folder.uri.fsPath)
-        const totalRoutes = collectRoutes([app]).length
-        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: collectRoutes(folderApps).length,
+      routes_count: folderRoutes.length,
       routers_count: countRouters(folderApps),
     })
   }

src/core/index.ts

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

src/core/treeUtils.ts

@@ -1,79 +1,45 @@
-/**
- * Utility functions for traversing router trees.
- * Consolidates common recursive tree operations used across the codebase.
- */
-
 import type { AppDefinition, RouteDefinition, RouterDefinition } from "./types"
 
-/**
- * Traverses all routers in a tree, calling the visitor function for each.
- * Returns early if visitor returns a non-undefined value.
- */
-function traverseRouters<T>(
-  routers: RouterDefinition[],
-  visitor: (router: RouterDefinition) => T | undefined,
-): T | undefined {
-  for (const router of routers) {
-    const result = visitor(router)
-    if (result !== undefined) {
-      return result
-    }
-    const childResult = traverseRouters(router.children, visitor)
-    if (childResult !== undefined) {
-      return childResult
-    }
-  }
-  return undefined
-}
-
-/**
- * Finds a router matching the predicate across all apps.
- */
+/** 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 = traverseRouters(app.routers, (router) =>
-      predicate(router) ? router : undefined,
-    )
+    const found = findIn(app.routers)
     if (found) return found
   }
   return undefined
 }
 
-/**
- * Collects all routes from a list of routers, including nested children.
- */
-function collectRoutesFromRouters(
-  routers: RouterDefinition[],
-): RouteDefinition[] {
-  return routers.flatMap((router) => [
-    ...router.routes,
-    ...collectRoutesFromRouters(router.children),
-  ])
+/** 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)])
 }
 
-/**
- * Collects all routes from all apps, including nested router routes.
- */
-export function collectRoutes(apps: AppDefinition[]): RouteDefinition[] {
-  return apps.flatMap((app) => [
-    ...app.routes,
-    ...collectRoutesFromRouters(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.
- */
+/** Counts total routers across all apps. */
 export function countRouters(apps: AppDefinition[]): number {
-  let count = 0
-  for (const app of apps) {
-    traverseRouters(app.routers, () => {
-      count++
-      return undefined
-    })
+  function count(routers: RouterDefinition[]): number {
+    return routers.reduce((sum, r) => sum + 1 + count(r.children), 0)
   }
-  return count
+  return apps.reduce((sum, app) => sum + count(app.routers), 0)
 }

src/extension.ts

@@ -138,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
@@ -148,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)
   }
@@ -214,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)
       },
     ),