two small fixes

- parsing issue with span name
- _main_ was not handled correctly by span search leading to performance issues

Author: doppleware

src/services/documentInfoProvider.ts

@@ -48,37 +48,55 @@ export class DocumentInfoProvider implements vscode.Disposable
     public async searchForSpan(instrumentationInfo: InstrumentationInfo): Promise<SpanInfo|undefined>{
         const codeFileHint = instrumentationInfo.instrumentationName;
         if (codeFileHint){
-            let codeHintFiles = codeFileHint.split(' ');
-            let head = codeHintFiles[0];
-            let folder = await vscode.workspace.workspaceFolders?.find(w => w.name === head);
-            let tail= codeHintFiles;
 
-            if (folder) {
-                tail = codeHintFiles.slice(1);
-            }
+            if (codeFileHint==='__main__'){
 
-            if (codeHintFiles.length>=1){
-                const files = await vscode.workspace.findFiles(`**/${tail.join('/')}.*`);
+                let doc = vscode.window.activeTextEditor?.document
+                if (doc){
+                    const docInfo = await this.getDocumentInfo(doc);
+                    if (docInfo){
+                        let spanInfos = docInfo.spans.filter(span => span.name===instrumentationInfo.spanName);
+                        return spanInfos.firstOrDefault(x=>x!== undefined);
 
-                const spansPromises = files.map(async file =>{
-                    try{
-                        const doc = await vscode.workspace.openTextDocument(file);
-                        const docInfo = await this.getDocumentInfo(doc);
-                        if(docInfo){
-                            return docInfo.spans.filter(span => span.name===instrumentationInfo.spanName);
-                        }
-                    }
-                    catch(error){
-                        Logger.warn(`Searching for span "${instrumentationInfo.spanName}" skipped ${file.fsPath}`, error);
                     }
-                    return [];
-                });
 
-                const spnaInfos = (await Promise.all(spansPromises)).flat();
-                if (spnaInfos.length===1){
-                    return spnaInfos[0];
+                }
+
+            }
+            else{
+
+                let codeHintFiles = codeFileHint.split(' ');
+                let head = codeHintFiles[0];
+                let folder = await vscode.workspace.workspaceFolders?.find(w => w.name === head);
+                let tail= codeHintFiles;
+    
+                if (folder) {
+                    tail = codeHintFiles.slice(1);
+                }
+    
+                if (codeHintFiles.length>=1){
+                    const files = await vscode.workspace.findFiles(`**/${tail.join('/')}.*`);
+    
+                    const spansPromises = files.map(async file =>{
+                        try{
+                            const doc = await vscode.workspace.openTextDocument(file);
+                            const docInfo = await this.getDocumentInfo(doc);
+                            if(docInfo){
+                                return docInfo.spans.filter(span => span.name===instrumentationInfo.spanName);
+                            }
+                        }
+                        catch(error){
+                            Logger.warn(`Searching for span "${instrumentationInfo.spanName}" skipped ${file.fsPath}`, error);
+                        }
+                        return [];
+                    });
+    
+                    const spnaInfos = (await Promise.all(spansPromises)).flat().firstOrDefault(x=>x!== undefined);
+                    return spnaInfos;
                 }
             }
+
+ 
         }
     }
 

src/services/languages/python/spanExtractor.ts

@@ -90,6 +90,8 @@ export class PythonSpanExtractor implements ISpanExtractor {
         return results;
     }
 
+
+
     private async extractNameTypeTrace(filePath: string) : Promise<string> {
 
         const pythonFileSuffix = ".py";
@@ -107,7 +109,9 @@ export class PythonSpanExtractor implements ISpanExtractor {
             if (relativePath.endsWith(pythonFileSuffix)){
                 relativePath=relativePath.substring(0,relativePath.length-pythonFileSuffix.length);
             }
-            relativePath=relativePath.replace("/",".").replace("\\",".");
+
+            relativePath=this.replaceAll(relativePath,"/",".");
+            relativePath=this.replaceAll(relativePath,"\\",".");
 
             return relativePath;
         }
@@ -116,6 +120,14 @@ export class PythonSpanExtractor implements ISpanExtractor {
 
     }
 
+    private escapeRegExp(s: string) {
+        return s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string
+    }
+    private replaceAll(str:string, find:string, replace:string) {
+        return str.replace(new RegExp(this.escapeRegExp(find), 'g'), replace);
+    }
+    
+
     private isCallToStartSpan(token: Token) {
         return token.type === TokenType.method && token.text === 'start_as_current_span';
     }