Merge pull request #94 from digma-ai/feature/fastapi-alternatesyntax-support

added fastapi router syntax

Author: doppleware

src/services/languages/python/fastapiEndpointExtractor.ts

@@ -16,14 +16,48 @@ export class FastapiEndpointExtractor implements IEndpointExtractor
         // Ensure fastapi module was imported
         if(!tokens.any(t => t.text == 'fastapi' && t.type == TokenType.module))
             return [];
+        
+        let prefix = '';
 
         // Search for "@app.get" decorators
-        const results: EndpointInfo[] = []
+        const results: EndpointInfo[] = [];
         for(let i=0; i<tokens.length-1; i++)
         {
             const appToken = tokens[i];
             const methodToken = tokens[i+1];
-            if (appToken.text != 'app' || appToken.type != TokenType.variable || methodToken.type != TokenType.method)
+
+            //Check if this is a router variable
+            if ( appToken.type === TokenType.variable 
+                    && methodToken.type === TokenType.class && methodToken.text==='APIRouter' ){
+               
+                const routerLine = document.getText(new vscode.Range(
+                    appToken.range.start, 
+                    new vscode.Position(methodToken.range.end.line, 1000)));
+                
+                //extract the prefix
+                let match = new RegExp(`^${appToken.text}\\s*=\\s*\\APIRouter\\s*\\((.*=.*,\\s*)*prefix=\\s*["'](.*?)["'](ֿֿֿ\\s*,.*=.*)*\\)*\\)`).exec(routerLine);
+                if (match){
+                    prefix = match[2];
+                }
+
+                //Future: use the references to extract the 'include' statement to the router if it contains a prefix
+                
+                // const position  =new vscode.Position(appToken.range.start.line,appToken.range.start.character );
+
+                // let references : vscode.Location[] = await vscode.commands.executeCommand("vscode.executeReferenceProvider", document.uri,position);
+                // references = references.filter(x=>x.uri.fsPath!==document.uri.fsPath);
+                // if (references.length>0){
+                //     let referencedDocument = vscode.workspace.openTextDocument(references[0].uri);
+                //     let text = (await referencedDocument).getText(references[0].range);
+                //     console.log(text);
+
+                // }
+
+
+            }
+                
+
+            if ((appToken.text != 'app' && appToken.text != 'router') || appToken.type != TokenType.variable || methodToken.type != TokenType.method)
                 continue;
 
             const method = methodToken.text;
@@ -33,21 +67,37 @@ export class FastapiEndpointExtractor implements IEndpointExtractor
             const lineText = document.getText(new vscode.Range(
                 appToken.range.start, 
                 new vscode.Position(methodToken.range.end.line, 1000)));
-            const match = new RegExp(`^app\\.${method}\\(["'](.*?)["']`).exec(lineText);
+            
+            let index = 2;
+            let match = new RegExp(`^(app|router)\\.${method}\\(["'](\/.*?)["']`).exec(lineText);
+            if (!match){
+                //Different regex for optional params (named)
+                match  =  new RegExp(`^(app|router)\\.${method}\\((.*=.*,\\s*)*path=\\s*["'](\\/.*?)["'](ֿֿֿ\\s*,.*=.*)*\\)`).exec(lineText);
+                index =3;
+            }
+
             if (!match)
                 continue;
 
             const relevantFunc = symbolInfo.firstOrDefault(s => s.range.contains(methodToken.range))
             if (!relevantFunc)
                 continue;
 
-            const path = match[1];
-            results.push(new EndpointInfo(
-                vscode.workspace.getWorkspaceFolder(document.uri)!.name + '$_$' + method + ' ' + match[1],
-                method, 
-                path,
-                relevantFunc.range,
-                document.uri));
+            const path = match[index];
+            let folder = vscode.workspace.getWorkspaceFolder(document.uri);
+            let folderPrefix = folder?.uri.path.split('/').slice(0,-1).join('/');
+            let relevantPath = document.uri.path.substring(folderPrefix!.length);
+            let pathParts = relevantPath.split('/').filter(x=>x);
+            for (let j=0;j<pathParts.length-1;j++){
+                let possibleRoot = pathParts[j];
+                results.push(new EndpointInfo(
+                    possibleRoot + '$_$' + method + ' ' + prefix + path,
+                    method, 
+                    path,
+                    relevantFunc.range,
+                    document.uri));
+            }
+
         }
         return results;
     }