SONARPY-2033: Adapt Highlighter for Jupyter Notebooks (#1899)

Author: joke1196

python-frontend/src/main/java/org/sonar/python/TestPythonVisitorRunner.java

@@ -25,6 +25,7 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.util.List;
+import java.util.Map;
 import javax.annotation.Nullable;
 import org.sonar.plugins.python.api.PythonCheck;
 import org.sonar.plugins.python.api.PythonFile;
@@ -52,6 +53,14 @@ public static PythonVisitorContext scanFile(File file, PythonCheck... visitors)
     return context;
   }
 
+  public static PythonVisitorContext scanNotebookFile(File file, Map<Integer, IPythonLocation> locations, String content, PythonCheck... visitors) {
+    PythonVisitorContext context = createNotebookContext(file, locations, content, "", ProjectLevelSymbolTable.empty(), CacheContextImpl.dummyCache());
+    for (PythonCheck visitor : visitors) {
+      visitor.scanFile(context);
+    }
+    return context;
+  }
+
   public static PythonVisitorContext createContext(File file) {
     return createContext(file, null);
   }
@@ -67,6 +76,13 @@ public static PythonVisitorContext createContext(File file, @Nullable File worki
     return new PythonVisitorContext(rootTree, pythonFile, workingDirectory, packageName, projectLevelSymbolTable, cacheContext);
   }
 
+  public static PythonVisitorContext createNotebookContext(File file, Map<Integer, IPythonLocation> locations, String content, String packageName,
+    ProjectLevelSymbolTable projectLevelSymbolTable, CacheContext cacheContext) {
+    TestPythonFile pythonFile = new TestPythonFile(file);
+    FileInput rootTree = parseNotebookFile(locations, content);
+    return new PythonVisitorContext(rootTree, pythonFile, null, packageName, projectLevelSymbolTable, cacheContext);
+  }
+
   public static ProjectLevelSymbolTable globalSymbols(List<File> files, File baseDir) {
     ProjectLevelSymbolTable projectLevelSymbolTable = new ProjectLevelSymbolTable();
     for (File file : files) {
@@ -81,9 +97,16 @@ public static ProjectLevelSymbolTable globalSymbols(List<File> files, File baseD
     return projectLevelSymbolTable;
   }
 
+  private static FileInput parseNotebookFile(Map<Integer, IPythonLocation> locations, String content) {
+    var parser = PythonParser.createIPythonParser();
+    var treeMaker = new IPythonTreeMaker(locations);
+    var astNode = parser.parse(content);
+    return treeMaker.fileInput(astNode);
+  }
+
   private static FileInput parseFile(TestPythonFile file) {
     var parser = file.isIPython() ? PythonParser.createIPythonParser() : PythonParser.create();
-    var treeMaker = file.isIPython() ? new IPythonTreeMaker() : new PythonTreeMaker();
+    var treeMaker = file.isIPython() ? new IPythonTreeMaker(Map.of()) : new PythonTreeMaker();
 
     var astNode = parser.parse(file.content());
     return treeMaker.fileInput(astNode);

python-frontend/src/test/java/org/sonar/python/TestPythonVisitorRunnerTest.java

@@ -23,6 +23,7 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.util.List;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 import org.sonar.plugins.python.api.PythonCheck;
 import org.sonar.plugins.python.api.PythonVisitorContext;
@@ -75,4 +76,12 @@ void globalSymbolsIPython() {
       .extracting(Symbol::name)
       .containsExactlyInAnyOrder("hello", "A");
   }
+
+  @Test
+  void scanNotebookFile() throws IOException {
+    File tmpFile = Files.createTempFile("foo", ".ipynb").toFile();
+    var check = (PythonCheck) visitorContect -> {};
+    PythonVisitorContext context = TestPythonVisitorRunner.scanNotebookFile(tmpFile, Map.of(), "", check);
+    assertThat(context.pythonFile().uri()).isEqualTo(tmpFile.toURI());
+  }
 }

sonar-python-plugin/src/main/java/org/sonar/plugins/python/PythonHighlighter.java

@@ -89,13 +89,14 @@
 public class PythonHighlighter extends PythonSubscriptionCheck {
 
   private NewHighlighting newHighlighting;
-
   private Set<Token> docStringTokens;
+  private PythonInputFile inputFile;
 
   public PythonHighlighter(SensorContext context, PythonInputFile inputFile) {
     docStringTokens = new HashSet<>();
     newHighlighting = context.newHighlighting();
     newHighlighting.onFile(inputFile.wrappedFile());
+    this.inputFile = inputFile;
   }
 
   @Override
@@ -137,8 +138,10 @@ private void visitToken(Token token) {
 
     }
 
-    for (Trivia trivia : token.trivia()) {
-      highlight(trivia.token(), TypeOfText.COMMENT);
+    if (inputFile.kind() == PythonInputFile.Kind.PYTHON) {
+      for (Trivia trivia : token.trivia()) {
+        highlight(trivia.token(), TypeOfText.COMMENT);
+      }
     }
   }