SONARPY-2076: Fix highlighting of single line Jupyter Notebook (#1935)

Author: joke1196

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

@@ -78,9 +78,10 @@ public Optional<String> parseLanguage() throws IOException {
 
   public GeneratedIPythonFile parseNotebook() throws IOException {
     String content = inputFile.wrappedFile().contents();
+    boolean isCompressed = content.lines().count() <= 1;
     JsonFactory factory = new JsonFactory();
     try (JsonParser jParser = factory.createParser(content)) {
-      return parseCells(jParser).map(notebookData -> {
+      return parseCells(jParser, isCompressed).map(notebookData -> {
         // Account for EOF token
         JsonLocation location = jParser.currentTokenLocation();
         notebookData.addDefaultLocation(lastPythonLine, location.getLineNr(), location.getColumnNr());
@@ -90,27 +91,27 @@ public GeneratedIPythonFile parseNotebook() throws IOException {
 
   }
 
-  private Optional<NotebookParsingData> parseCells(JsonParser parser) throws IOException {
+  private Optional<NotebookParsingData> parseCells(JsonParser parser, boolean isCompressed) throws IOException {
     while (!parser.isClosed()) {
       parser.nextToken();
       String fieldName = parser.currentName();
       if ("cells".equals(fieldName)) {
         // consume array start token
         parser.nextToken();
-        Optional<NotebookParsingData> data = parseCellArray(parser);
+        Optional<NotebookParsingData> data = parseCellArray(parser, isCompressed);
         parser.close();
         return data;
       }
     }
     return Optional.empty();
   }
 
-  private Optional<NotebookParsingData> parseCellArray(JsonParser jParser) throws IOException {
+  private Optional<NotebookParsingData> parseCellArray(JsonParser jParser, boolean isCompressed) throws IOException {
     List<NotebookParsingData> cellsData = new ArrayList<>();
 
     while (jParser.nextToken() != JsonToken.END_ARRAY) {
       if (jParser.currentToken() == JsonToken.START_OBJECT) {
-        processCodeCell(cellsData, jParser);
+        processCodeCell(cellsData, jParser, isCompressed);
       }
     }
     Optional<NotebookParsingData> aggregatedNotebookData = cellsData.stream().reduce(NotebookParsingData::combine);
@@ -134,7 +135,7 @@ private static boolean processCodeCellType(JsonParser jParser) throws IOExceptio
     return false;
   }
 
-  private void processCodeCell(List<NotebookParsingData> accumulator, JsonParser jParser) throws IOException {
+  private void processCodeCell(List<NotebookParsingData> accumulator, JsonParser jParser, boolean isCompressed) throws IOException {
     boolean isCodeCell = false;
     Optional<NotebookParsingData> notebookData = Optional.empty();
     while (jParser.nextToken() != JsonToken.END_OBJECT) {
@@ -154,7 +155,7 @@ private void processCodeCell(List<NotebookParsingData> accumulator, JsonParser j
         }
         switch (jParser.currentToken()) {
           case START_ARRAY:
-            notebookData = Optional.of(parseSourceArray(startLine, jParser));
+            notebookData = Optional.of(parseSourceArray(startLine, jParser, isCompressed));
             break;
           case VALUE_STRING:
             notebookData = Optional.of(parseSourceMultilineString(startLine, jParser));
@@ -172,7 +173,7 @@ private void processCodeCell(List<NotebookParsingData> accumulator, JsonParser j
     }
   }
 
-  private static NotebookParsingData parseSourceArray(int startLine, JsonParser jParser) throws IOException {
+  private static NotebookParsingData parseSourceArray(int startLine, JsonParser jParser, boolean isCompressed) throws IOException {
     NotebookParsingData cellData = NotebookParsingData.fromLine(startLine);
     JsonLocation tokenLocation = jParser.currentTokenLocation();
     // In case of an empty cell, we don't add an extra line
@@ -181,7 +182,7 @@ private static NotebookParsingData parseSourceArray(int startLine, JsonParser jP
       String sourceLine = jParser.getValueAsString();
       var newTokenLocation = jParser.currentTokenLocation();
       var countEscapedChar = countEscapeCharacters(sourceLine, newTokenLocation.getColumnNr());
-      cellData.addLineToSource(sourceLine, newTokenLocation.getLineNr(), newTokenLocation.getColumnNr(), countEscapedChar);
+      cellData.addLineToSource(sourceLine, newTokenLocation.getLineNr(), newTokenLocation.getColumnNr(), countEscapedChar, isCompressed);
       lastSourceLine = sourceLine;
       tokenLocation = newTokenLocation;
     }

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

@@ -70,8 +70,8 @@ public void appendToSource(String str) {
     aggregatedSource.append(str);
   }
 
-  public void addLineToSource(String sourceLine, int lineNr, int columnNr, Map<Integer, Integer> colOffset) {
-    addLineToSource(sourceLine, new IPythonLocation(lineNr, columnNr, colOffset));
+  public void addLineToSource(String sourceLine, int lineNr, int columnNr, Map<Integer, Integer> colOffset, boolean isCompressed) {
+    addLineToSource(sourceLine, new IPythonLocation(lineNr, columnNr, colOffset, isCompressed));
   }
 
   private void appendLine(String line) {

sonar-python-plugin/src/test/java/org/sonar/plugins/python/PythonHighlighterTest.java

@@ -41,7 +41,9 @@ class PythonHighlighterTest {
 
   private File file;
   private File notebookFile;
+  private File notebookFileSingleLine;
   private DefaultInputFile notebookInputFile;
+  private DefaultInputFile notebookInputFileSingleLine;
   private String dir = "src/test/resources/org/sonar/plugins/python";
 
   @BeforeEach
@@ -56,9 +58,16 @@ void scanFile() {
     notebookInputFile = TestInputFileBuilder.create("moduleKey", notebookFile.getName())
       .initMetadata(TestUtils.fileContent(notebookFile, StandardCharsets.UTF_8))
       .build();
+
+    notebookFileSingleLine = new File(dir, "/notebookHighlighterSingleLine.ipynb");
+    notebookInputFileSingleLine = TestInputFileBuilder.create("moduleKey", notebookFileSingleLine.getName())
+      .initMetadata(TestUtils.fileContent(notebookFileSingleLine, StandardCharsets.UTF_8))
+      .build();
+
     context = SensorContextTester.create(new File(dir));
     context.fileSystem().add(inputFile);
     context.fileSystem().add(notebookInputFile);
+    context.fileSystem().add(notebookInputFileSingleLine);
 
     PythonHighlighter pythonHighlighter = new PythonHighlighter(context, new PythonInputFileImpl(inputFile));
     TestPythonVisitorRunner.scanFile(file, pythonHighlighter);
@@ -226,6 +235,26 @@ void highlightingNotebooks() {
     checkOnRange(13, 9, 2, notebookFile, TypeOfText.CONSTANT);
   }
 
+  @Test
+  void highlightingNotebooksSingleLine() {
+    String pythonContent = "def foo():\n    pass\na = 2 # comment\n#SONAR_PYTHON_NOTEBOOK_CELL_DELIMITER";
+    var locations = Map.of(
+      1, new IPythonLocation(1, 93, Map.of(-1, 0), true),
+      2, new IPythonLocation(1, 108, Map.of(-1, 0), true),
+      3, new IPythonLocation(1, 121, Map.of(-1, 0), true),
+      4, new IPythonLocation(1, 93, Map.of(-1, 0), true)); //EOF Token
+    PythonHighlighter pythonHighlighter = new PythonHighlighter(context, new GeneratedIPythonFile(notebookInputFileSingleLine, pythonContent, locations));
+    TestPythonVisitorRunner.scanNotebookFile(notebookFileSingleLine, locations, pythonContent, pythonHighlighter);
+    // def
+    checkOnRange(1, 93, 3, notebookFileSingleLine, TypeOfText.KEYWORD);
+    // pass
+    checkOnRange(1, 112, 4, notebookFileSingleLine, TypeOfText.KEYWORD);
+    // 2
+    checkOnRange(1, 125, 1, notebookFileSingleLine, TypeOfText.CONSTANT);
+    // # comment
+    checkOnRange(1, 127 , 9, notebookFileSingleLine, TypeOfText.COMMENT);
+  }
+
   /**
    * Checks the highlighting of a range of columns. The first column of a line has index 0.
    * The range is the columns of the token.

sonar-python-plugin/src/test/resources/org/sonar/plugins/python/notebookHighlighterSingleLine.ipynb

@@ -0,0 +1 @@
+{"cells":[{"cell_type":"code","execution_count": null,"metadata":{},"outputs":[],"source":["def foo():\n","    pass\n","a = 2 # comment"]}],"metadata": {   "language_info": {    "name": "python"   }  },  "nbformat": 4,  "nbformat_minor": 2 }