SONARPY-1990: Map JSON locations to python locations in IPythonTreeMaker (#1884)

Author: joke1196

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

@@ -0,0 +1,25 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2024 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.python;
+
+import java.util.Map;
+
+public record IPythonLocation(int line, int column, Map<Integer, Integer> colOffset) {
+}

python-frontend/src/main/java/org/sonar/python/tree/IPythonTreeMaker.java

@@ -24,8 +24,10 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import javax.annotation.Nullable;
 import org.sonar.plugins.python.api.tree.CellMagicStatement;
 import org.sonar.plugins.python.api.tree.DynamicObjectInfoStatement;
 import org.sonar.plugins.python.api.tree.Expression;
@@ -37,17 +39,38 @@
 import org.sonar.plugins.python.api.tree.Token;
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.python.DocstringExtractor;
+import org.sonar.python.IPythonLocation;
 import org.sonar.python.api.IPythonGrammar;
 import org.sonar.python.api.PythonGrammar;
 
 public class IPythonTreeMaker extends PythonTreeMaker {
 
+  private Map<Integer, IPythonLocation> offsetMap;
+
+  public IPythonTreeMaker() {
+    this(null);
+  }
+  public IPythonTreeMaker(@Nullable Map<Integer, IPythonLocation> offsetMap) {
+    this.offsetMap = offsetMap;
+  }
+
+  @Override
+  protected Token toPyToken(@Nullable com.sonar.sslr.api.Token token) {
+    if (token == null) {
+      return null;
+    }
+    if (this.offsetMap != null) {
+      return TokenEnricher.enrichToken(token, offsetMap);
+    }
+    return new TokenImpl(token);
+  }
+
   @Override
   public FileInput fileInput(AstNode astNode) {
     StatementList statementList = astNode.getChildren(IPythonGrammar.CELL, IPythonGrammar.MAGIC_CELL)
       .stream()
       .flatMap(this::getStatementsFromCell)
-      .collect(Collectors.collectingAndThen(Collectors.toList(), l -> l.isEmpty()? null : new StatementListImpl(l)));
+      .collect(Collectors.collectingAndThen(Collectors.toList(), l -> l.isEmpty() ? null : new StatementListImpl(l)));
     Token endOfFile = toPyToken(astNode.getFirstChild(GenericTokenType.EOF).getToken());
     FileInputImpl pyFileInputTree = new FileInputImpl(statementList, endOfFile, DocstringExtractor.extractDocstring(statementList));
     setParents(pyFileInputTree);
@@ -59,7 +82,7 @@ private Stream<Statement> getStatementsFromCell(AstNode cell) {
     if (cell.is(IPythonGrammar.CELL)) {
       return getStatements(cell).stream().map(this::statement);
     } else {
-      return Stream.of(cell.getFirstChild(IPythonGrammar.CELL_MAGIC_STATEMENT)).map(IPythonTreeMaker::cellMagicStatement);
+      return Stream.of(cell.getFirstChild(IPythonGrammar.CELL_MAGIC_STATEMENT)).map(this::cellMagicStatement);
     }
   }
 
@@ -85,12 +108,12 @@ protected Expression annotatedRhs(AstNode annotatedRhs) {
     return super.annotatedRhs(annotatedRhs);
   }
 
-  private static CellMagicStatement cellMagicStatement(AstNode astNode) {
+  private  CellMagicStatement cellMagicStatement(AstNode astNode) {
     var tokens = astNode.getChildren()
       .stream()
       .map(AstNode::getTokens)
       .flatMap(Collection::stream)
-      .map(IPythonTreeMaker::toPyToken)
+      .map(this::toPyToken)
       .toList();
     return new CellMagicStatementImpl(tokens);
   }
@@ -121,7 +144,7 @@ protected DynamicObjectInfoStatement dynamicObjectInfoStatement(AstNode astNode)
         isQuestionMarksBefore = false;
         child.getTokens()
           .stream()
-          .map(PythonTreeMaker::toPyToken)
+          .map(this::toPyToken)
           .forEach(children::add);
       }
     }
@@ -137,9 +160,8 @@ protected LineMagic lineMagic(AstNode astNode) {
       .skip(2)
       .map(AstNode::getTokens)
       .flatMap(Collection::stream)
-      .map(IPythonTreeMaker::toPyToken)
+      .map(this::toPyToken)
       .toList();
     return new LineMagicImpl(percent, name, tokens);
   }
-
 }