🐛 Fix CodeLens route matching and nested router discovery (#20)

Author: savannahostrowski

src/core/extractors.ts

@@ -59,7 +59,7 @@ export function extractStringValue(node: Node): string | null {
  * Extracts a path string from various AST node types.
  * Handles: plain strings, f-strings, concatenation, identifiers.
  */
-function extractPathFromNode(node: Node): string {
+export function extractPathFromNode(node: Node): string {
   switch (node.type) {
     case "string":
       return extractStringValue(node) ?? ""

src/core/importResolver.ts

@@ -157,11 +157,18 @@ export async function resolveNamedImport(
     : modulePathToDir(importInfo, currentFileUri, projectRootUri, fs)
 
   for (const name of importInfo.names) {
-    // Try direct file: from .routes import users -> routes/users.py
-    const namedUri = fs.joinPath(baseDirUri, ...name.split("."))
-    const resolved = await resolvePythonModule(namedUri, fs)
-    if (resolved) {
-      return resolved
+    // Only try submodule resolution if baseUri is a package (__init__.py) or namespace package (null).
+    // For regular .py files, the name is a variable inside the file, not a submodule.
+    // Example: "from .neon import router" where neon.py defines router
+    // should resolve to neon.py, not look for router.py
+    const isPackage = baseUri === null || baseUri.endsWith("__init__.py")
+    if (isPackage) {
+      // Try direct file: from .routes import users -> routes/users.py
+      const namedUri = fs.joinPath(baseDirUri, ...name.split("."))
+      const resolved = await resolvePythonModule(namedUri, fs)
+      if (resolved) {
+        return resolved
+      }
     }
 
     // Try re-exports: from .routes import users where routes/__init__.py re-exports users

src/core/pathUtils.ts

@@ -90,33 +90,52 @@ export function countSegments(path: string): number {
 
 /**
  * Checks if a test path matches an endpoint path pattern.
- * Endpoint paths may contain path parameters like {item_id} which match any segment.
+ * Both paths may contain dynamic segments like {item_id} or {settings.API_V1_STR}
+ * which match any segment.
+ *
+ * Leading dynamic prefixes (like {settings.API_V1_STR}) and query strings are stripped
+ * before comparison.
  *
  * Examples:
  *   pathMatchesEndpoint("/items/123", "/items/{item_id}") -> true
  *   pathMatchesEndpoint("/items/123/details", "/items/{item_id}") -> false
  *   pathMatchesEndpoint("/users/abc/posts/456", "/users/{user_id}/posts/{post_id}") -> true
  *   pathMatchesEndpoint("/items/", "/items/{item_id}") -> false
+ *   pathMatchesEndpoint("{settings.API}/apps/{id}", "/apps/{app_id}") -> true
+ *   pathMatchesEndpoint("{BASE}/users/{id}", "/users/{user_id}") -> true
+ *   pathMatchesEndpoint("/teams/?owner=true", "/teams") -> true (query string stripped)
  */
 export function pathMatchesEndpoint(
   testPath: string,
   endpointPath: string,
 ): boolean {
-  const testSegments = testPath.split("/").filter(Boolean)
-  const endpointSegments = endpointPath.split("/").filter(Boolean)
+  // Strip query string from test path (e.g., "/teams/?owner=true" -> "/teams/")
+  const testPathWithoutQuery = testPath.split("?")[0]
+
+  // Strip leading dynamic segments (e.g., {settings.API_V1_STR}) for comparison
+  const testSegments = stripLeadingDynamicSegments(testPathWithoutQuery)
+    .split("/")
+    .filter(Boolean)
+  const endpointSegments = stripLeadingDynamicSegments(endpointPath)
+    .split("/")
+    .filter(Boolean)
 
   // Segment counts must match
   if (testSegments.length !== endpointSegments.length) {
     return false
   }
 
-  return endpointSegments.every((seg, index) => {
-    // Path parameter (e.g., {item_id}) matches any segment
-    if (seg.startsWith("{") && seg.endsWith("}")) {
+  // Compare each segment positionally
+  return testSegments.every((testSeg, i) => {
+    const endpointSeg = endpointSegments[i]
+    // Dynamic segments (e.g., {id}, {app.id}) match any value
+    const testIsDynamic = testSeg.startsWith("{") && testSeg.endsWith("}")
+    const endpointIsDynamic =
+      endpointSeg.startsWith("{") && endpointSeg.endsWith("}")
+    if (testIsDynamic || endpointIsDynamic) {
       return true
     }
-    // Literal segments must match exactly
-    return seg === testSegments[index]
+    return testSeg === endpointSeg
   })
 }
 

src/core/routerResolver.ts

@@ -218,7 +218,7 @@ async function resolveRouterReference(
   if (localRouter) {
     // Filter routes that belong to this router (decorated with @router.method)
     const routerRoutes = analysis.routes.filter((r) => r.owner === moduleName)
-    return {
+    const routerNode: RouterNode = {
       filePath: currentFileUri,
       variableName: localRouter.variableName,
       type: localRouter.type,
@@ -235,6 +235,39 @@ async function resolveRouterReference(
       })),
       children: [],
     }
+
+    // Process include_router calls owned by this router (nested routers)
+    const routerIncludes = analysis.includeRouters.filter(
+      (inc) => inc.owner === moduleName,
+    )
+    for (const include of routerIncludes) {
+      log(
+        `Resolving nested include_router: ${include.router} (owner: ${moduleName}, prefix: ${include.prefix || "none"})`,
+      )
+      const childRouter = await resolveRouterReference(
+        include.router,
+        analysis,
+        currentFileUri,
+        projectRootUri,
+        parser,
+        fs,
+        visited,
+      )
+      if (childRouter) {
+        if (include.tags.length > 0) {
+          childRouter.tags = [
+            ...new Set([...childRouter.tags, ...include.tags]),
+          ]
+        }
+        routerNode.children.push({
+          router: childRouter,
+          prefix: include.prefix,
+          tags: include.tags,
+        })
+      }
+    }
+
+    return routerNode
   }
 
   // Otherwise, look for an imported router
@@ -250,11 +283,19 @@ async function resolveRouterReference(
   // Helper to analyze a file with the filesystem
   const analyzeFileFn = (uri: string) => analyzeFile(uri, parser, fs)
 
+  // Find the original import name (in case moduleName is an alias)
+  // e.g., "from .api_tokens import router as api_tokens_router"
+  // moduleName = "api_tokens_router", originalName = "router"
+  const namedImport = matchingImport.namedImports.find(
+    (ni) => (ni.alias ?? ni.name) === moduleName,
+  )
+  const originalName = namedImport?.name ?? moduleName
+
   // Resolve the imported module to a file URI
   const importedFileUri = await resolveNamedImport(
     {
       modulePath: matchingImport.modulePath,
-      names: [moduleName],
+      names: [originalName],
       isRelative: matchingImport.isRelative,
       relativeDots: matchingImport.relativeDots,
     },
@@ -293,7 +334,7 @@ async function resolveRouterReference(
       const routerRoutes = importedAnalysis.routes.filter(
         (r) => r.owner === attributeName,
       )
-      return {
+      const routerNode: RouterNode = {
         filePath: importedFileUri,
         variableName: targetRouter.variableName,
         type: targetRouter.type,
@@ -310,6 +351,39 @@ async function resolveRouterReference(
         })),
         children: [],
       }
+
+      // Process include_router calls owned by this router (nested routers)
+      const routerIncludes = importedAnalysis.includeRouters.filter(
+        (inc) => inc.owner === attributeName,
+      )
+      for (const include of routerIncludes) {
+        log(
+          `Resolving nested include_router: ${include.router} (owner: ${attributeName}, prefix: ${include.prefix || "none"})`,
+        )
+        const childRouter = await resolveRouterReference(
+          include.router,
+          importedAnalysis,
+          importedFileUri,
+          projectRootUri,
+          parser,
+          fs,
+          visited,
+        )
+        if (childRouter) {
+          if (include.tags.length > 0) {
+            childRouter.tags = [
+              ...new Set([...childRouter.tags, ...include.tags]),
+            ]
+          }
+          routerNode.children.push({
+            router: childRouter,
+            prefix: include.prefix,
+            tags: include.tags,
+          })
+        }
+      }
+
+      return routerNode
     }
     // If not found as a router, fall through to try building from file
   }

src/providers/testCodeLensProvider.ts

@@ -14,7 +14,7 @@ import {
   Uri,
 } from "vscode"
 import type { Node } from "web-tree-sitter"
-import { extractStringValue, findNodesByType } from "../core/extractors"
+import { extractPathFromNode, findNodesByType } from "../core/extractors"
 import { ROUTE_METHODS } from "../core/internal"
 import type { Parser } from "../core/parser"
 import {
@@ -132,9 +132,8 @@ export class TestCodeLensProvider implements CodeLensProvider {
       }
 
       const pathArg = args[0]
-      // Only handle string literals for now
-      const path = extractStringValue(pathArg)
-      if (path === null) {
+      const path = extractPathFromNode(pathArg)
+      if (!path) {
         continue
       }
 

src/test/core/importResolver.test.ts

@@ -241,5 +241,32 @@ suite("importResolver", () => {
       assert.ok(result)
       assert.ok(result.endsWith("api_routes.py"))
     })
+
+    test("resolves variable import from .py file (not submodule)", async () => {
+      // This tests "from .neon import router" where router is a variable in neon.py,
+      // NOT a submodule. Should return neon.py, not look for router.py
+      const reexportRoot = fixtures.reexport.root
+      const currentFile = nodeFileSystem.joinPath(
+        reexportRoot,
+        "app",
+        "integrations",
+        "router.py",
+      )
+
+      const result = await resolveNamedImport(
+        {
+          modulePath: "neon",
+          names: ["router"],
+          isRelative: true,
+          relativeDots: 1,
+        },
+        currentFile,
+        reexportRoot,
+        nodeFileSystem,
+      )
+
+      assert.ok(result, "Should resolve import")
+      assert.ok(result.endsWith("neon.py"), `Expected neon.py, got ${result}`)
+    })
   })
 })

src/test/core/pathUtils.test.ts

@@ -271,19 +271,49 @@ suite("pathUtils", () => {
       )
     })
 
-    test("matches paths with dynamic prefix", () => {
-      // Dynamic prefixes like {settings.API_V1_STR} match any segment (same as path params)
+    test("matches paths with dynamic prefix in endpoint", () => {
       assert.strictEqual(
         pathMatchesEndpoint(
-          "/v1/items/123",
+          "/items/123",
           "{settings.API_V1_STR}/items/{item_id}",
         ),
         true,
       )
       assert.strictEqual(
         pathMatchesEndpoint("/api/v2/users", "{BASE}/users"),
-        false, // segment count differs
+        false,
+      )
+      assert.strictEqual(pathMatchesEndpoint("/users", "{BASE}/users"), true)
+      assert.strictEqual(pathMatchesEndpoint("/items", "{BASE}/users"), false)
+    })
+
+    test("matches paths with dynamic prefix in test path (f-strings)", () => {
+      assert.strictEqual(
+        pathMatchesEndpoint(
+          "{settings.API_V1_STR}/apps/{app.id}/environment-variables",
+          "/apps/{app_id}/environment-variables",
+        ),
+        true,
+      )
+      assert.strictEqual(
+        pathMatchesEndpoint(
+          "{settings.API}/items/{item_id}",
+          "{BASE}/items/{id}",
+        ),
+        true,
+      )
+      assert.strictEqual(
+        pathMatchesEndpoint("{BASE}/users/{id}", "/items/{item_id}"),
+        false,
+      )
+    })
+
+    test("strips query strings from test path", () => {
+      assert.strictEqual(
+        pathMatchesEndpoint("/teams/?owner=true&order_by=created_at", "/teams"),
+        true,
       )
+      assert.strictEqual(pathMatchesEndpoint("/?page=1", "/"), true)
     })
   })
 })