SONARPY-2008 Infer types for comprehensions (#1866)

guillaume-dequenne · guillaume-dequenne · commit 02bd1daa452d · 2024-08-09T09:14:17.000+02:00
diff --git a/python-checks/src/main/java/org/sonar/python/checks/InfiniteRecursionCheck.java b/python-checks/src/main/java/org/sonar/python/checks/InfiniteRecursionCheck.java
@@ -43,6 +43,7 @@
 import org.sonar.plugins.python.api.tree.ComprehensionExpression;
 import org.sonar.plugins.python.api.tree.ComprehensionIf;
 import org.sonar.plugins.python.api.tree.ConditionalExpression;
+import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.Expression;
 import org.sonar.plugins.python.api.tree.ExpressionList;
 import org.sonar.plugins.python.api.tree.FunctionDef;
@@ -194,7 +195,7 @@ public void visitComprehensionIf(ComprehensionIf tree) {
     }
 
     @Override
-    public void visitDictCompExpression(DictCompExpressionImpl tree) {
+    public void visitDictCompExpression(DictCompExpression tree) {
       // ignore, not broken down in the cfg
     }
 
diff --git a/python-frontend/src/main/java/org/sonar/plugins/python/api/tree/BaseTreeVisitor.java b/python-frontend/src/main/java/org/sonar/plugins/python/api/tree/BaseTreeVisitor.java
@@ -431,7 +431,7 @@ public void visitKeyValuePattern(KeyValuePattern keyValuePattern) {
   }
 
   @Override
-  public void visitDictCompExpression(DictCompExpressionImpl tree) {
+  public void visitDictCompExpression(DictCompExpression tree) {
     scan(tree.keyExpression());
     scan(tree.valueExpression());
     scan(tree.comprehensionFor());
diff --git a/python-frontend/src/main/java/org/sonar/plugins/python/api/tree/TreeVisitor.java b/python-frontend/src/main/java/org/sonar/plugins/python/api/tree/TreeVisitor.java
@@ -157,7 +157,7 @@ public interface TreeVisitor {
 
   void visitKeyValuePattern(KeyValuePattern keyValuePattern);
 
-  void visitDictCompExpression(DictCompExpressionImpl dictCompExpression);
+  void visitDictCompExpression(DictCompExpression dictCompExpression);
 
   void visitCompoundAssignment(CompoundAssignmentStatement compoundAssignmentStatement);
 
diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/SymbolTableBuilder.java b/python-frontend/src/main/java/org/sonar/python/semantic/SymbolTableBuilder.java
@@ -53,6 +53,7 @@
 import org.sonar.plugins.python.api.tree.ComprehensionExpression;
 import org.sonar.plugins.python.api.tree.ComprehensionFor;
 import org.sonar.plugins.python.api.tree.Decorator;
+import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.DottedName;
 import org.sonar.plugins.python.api.tree.ExceptClause;
 import org.sonar.plugins.python.api.tree.Expression;
@@ -294,7 +295,7 @@ public void visitLambda(LambdaExpression pyLambdaExpressionTree) {
     }
 
     @Override
-    public void visitDictCompExpression(DictCompExpressionImpl tree) {
+    public void visitDictCompExpression(DictCompExpression tree) {
       createScope(tree, currentScope());
       enterScope(tree);
       super.visitDictCompExpression(tree);
@@ -662,7 +663,7 @@ public void visitPyListOrSetCompExpression(ComprehensionExpression tree) {
     }
 
     @Override
-    public void visitDictCompExpression(DictCompExpressionImpl tree) {
+    public void visitDictCompExpression(DictCompExpression tree) {
       enterScope(tree);
       scan(tree.keyExpression());
       scan(tree.valueExpression());
diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/ReadUsagesVisitor.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/ReadUsagesVisitor.java
@@ -23,6 +23,7 @@
 import org.sonar.plugins.python.api.tree.ClassDef;
 import org.sonar.plugins.python.api.tree.ComprehensionExpression;
 import org.sonar.plugins.python.api.tree.ComprehensionFor;
+import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.FileInput;
 import org.sonar.plugins.python.api.tree.FunctionDef;
 import org.sonar.plugins.python.api.tree.LambdaExpression;
@@ -91,7 +92,7 @@ public void visitPyListOrSetCompExpression(ComprehensionExpression tree) {
   }
 
   @Override
-  public void visitDictCompExpression(DictCompExpressionImpl tree) {
+  public void visitDictCompExpression(DictCompExpression tree) {
     enterScope(tree);
     scan(tree.keyExpression());
     scan(tree.valueExpression());
diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/WriteUsagesVisitor.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/WriteUsagesVisitor.java
@@ -38,6 +38,7 @@
 import org.sonar.plugins.python.api.tree.ComprehensionExpression;
 import org.sonar.plugins.python.api.tree.ComprehensionFor;
 import org.sonar.plugins.python.api.tree.Decorator;
+import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.DottedName;
 import org.sonar.plugins.python.api.tree.ExceptClause;
 import org.sonar.plugins.python.api.tree.FileInput;
@@ -87,7 +88,7 @@ public void visitLambda(LambdaExpression lambdaExpression) {
   }
 
   @Override
-  public void visitDictCompExpression(DictCompExpressionImpl tree) {
+  public void visitDictCompExpression(DictCompExpression tree) {
     createAndEnterScope(tree, currentScope());
     super.visitDictCompExpression(tree);
     leaveScope();
diff --git a/python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java b/python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java
@@ -33,6 +33,8 @@
 import org.sonar.plugins.python.api.tree.BaseTreeVisitor;
 import org.sonar.plugins.python.api.tree.BinaryExpression;
 import org.sonar.plugins.python.api.tree.ClassDef;
+import org.sonar.plugins.python.api.tree.ComprehensionExpression;
+import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.DictionaryLiteral;
 import org.sonar.plugins.python.api.tree.DottedName;
 import org.sonar.plugins.python.api.tree.Expression;
@@ -60,6 +62,8 @@
 import org.sonar.python.semantic.v2.ProjectLevelTypeTable;
 import org.sonar.python.semantic.v2.SymbolV2;
 import org.sonar.python.semantic.v2.UsageV2;
+import org.sonar.python.tree.ComprehensionExpressionImpl;
+import org.sonar.python.tree.DictCompExpressionImpl;
 import org.sonar.python.tree.DictionaryLiteralImpl;
 import org.sonar.python.tree.ListLiteralImpl;
 import org.sonar.python.tree.NameImpl;
@@ -177,6 +181,26 @@ public void visitListLiteral(ListLiteral listLiteral) {
     ((ListLiteralImpl) listLiteral).typeV2(new ObjectType(listType, attributes, new ArrayList<>()));
   }
 
+  @Override
+  public void visitPyListOrSetCompExpression(ComprehensionExpression comprehensionExpression) {
+    super.visitPyListOrSetCompExpression(comprehensionExpression);
+    var builtins = this.projectLevelTypeTable.getModule();
+    var pythonType = switch (comprehensionExpression.getKind()) {
+      case LIST_COMPREHENSION -> builtins.resolveMember("list").orElse(PythonType.UNKNOWN);
+      case SET_COMPREHENSION -> builtins.resolveMember("set").orElse(PythonType.UNKNOWN);
+      default -> PythonType.UNKNOWN;
+    };
+    ((ComprehensionExpressionImpl) comprehensionExpression).typeV2(new ObjectType(pythonType, new ArrayList<>(), new ArrayList<>()));
+  }
+
+  @Override
+  public void visitDictCompExpression(DictCompExpression dictCompExpression) {
+    super.visitDictCompExpression(dictCompExpression);
+    var builtins = this.projectLevelTypeTable.getModule();
+    var dictType = builtins.resolveMember("dict").orElse(PythonType.UNKNOWN);
+    ((DictCompExpressionImpl) dictCompExpression).typeV2(new ObjectType(dictType, new ArrayList<>(), new ArrayList<>()));
+  }
+
   @Override
   public void visitClassDef(ClassDef classDef) {
     scan(classDef.args());
diff --git a/python-frontend/src/main/java/org/sonar/python/tree/ComprehensionExpressionImpl.java b/python-frontend/src/main/java/org/sonar/python/tree/ComprehensionExpressionImpl.java
@@ -34,6 +34,7 @@
 import org.sonar.plugins.python.api.tree.TreeVisitor;
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.types.InferredTypes;
+import org.sonar.python.types.v2.PythonType;
 
 public class ComprehensionExpressionImpl extends PyTree implements ComprehensionExpression {
 
@@ -43,6 +44,7 @@ public class ComprehensionExpressionImpl extends PyTree implements Comprehension
   private final ComprehensionFor comprehensionFor;
   private final Token closingToken;
   private Set<Symbol> symbols = new HashSet<>();
+  private PythonType pythonType;
 
   public ComprehensionExpressionImpl(Kind kind, @Nullable Token openingToken, Expression resultExpression,
                                      ComprehensionFor compFor, @Nullable Token closingToken) {
@@ -99,4 +101,13 @@ public InferredType type() {
         return InferredTypes.anyType();
     }
   }
+
+  @Override
+  public PythonType typeV2() {
+    return this.pythonType;
+  }
+
+  public void typeV2(PythonType type) {
+    this.pythonType = type;
+  }
 }
diff --git a/python-frontend/src/main/java/org/sonar/python/tree/DictCompExpressionImpl.java b/python-frontend/src/main/java/org/sonar/python/tree/DictCompExpressionImpl.java
@@ -33,6 +33,7 @@
 import org.sonar.plugins.python.api.tree.TreeVisitor;
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.types.InferredTypes;
+import org.sonar.python.types.v2.PythonType;
 
 public class DictCompExpressionImpl extends PyTree implements DictCompExpression {
 
@@ -43,6 +44,7 @@ public class DictCompExpressionImpl extends PyTree implements DictCompExpression
   private final ComprehensionFor comprehensionFor;
   private final Token closingBrace;
   private Set<Symbol> symbols = new HashSet<>();
+  private PythonType pythonType;
 
   public DictCompExpressionImpl(Token openingBrace, Expression keyExpression, Token colon, Expression valueExpression,
                                 ComprehensionFor compFor, Token closingBrace) {
@@ -102,4 +104,13 @@ public void addLocalVariableSymbol(Symbol symbol) {
   public InferredType type() {
     return InferredTypes.DICT;
   }
+
+  @Override
+  public PythonType typeV2() {
+    return this.pythonType;
+  }
+
+  public void typeV2(PythonType type) {
+    this.pythonType = type;
+  }
 }
diff --git a/python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java b/python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java
@@ -62,10 +62,12 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.sonar.python.PythonTestUtils.parse;
+import static org.sonar.python.types.v2.TypesTestUtils.DICT_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.INT_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.LIST_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.NONE_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.PROJECT_LEVEL_TYPE_TABLE;
+import static org.sonar.python.types.v2.TypesTestUtils.SET_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.STR_TYPE;
 
 class TypeInferenceV2Test {
@@ -1679,6 +1681,56 @@ def reassigned_param(a, param):
     Assertions.assertThat(paramType).isSameAs(PythonType.UNKNOWN);
   }
 
+  @Test
+  void list_comprehension() {
+    assertThat(lastExpression(
+      """
+      x = [a for a in foo()]
+      x
+      """
+    ).typeV2().unwrappedType()).isEqualTo(LIST_TYPE);
+  }
+
+  @Test
+  void set_comprehension() {
+    assertThat(lastExpression(
+      """
+      x = {a for a in foo()}
+      x
+      """
+    ).typeV2().unwrappedType()).isEqualTo(SET_TYPE);
+  }
+
+  @Test
+  void dict_comprehension() {
+    assertThat(lastExpression(
+      """
+      x = {num: num**2 for num in numbers()}
+      x
+      """
+    ).typeV2().unwrappedType()).isEqualTo(DICT_TYPE);
+  }
+
+  @Test
+  void comprehension_if() {
+    assertThat(lastExpression(
+      """
+      x = [num for num in numbers if num % 2 == 0]
+      x
+      """
+    ).typeV2().unwrappedType()).isEqualTo(LIST_TYPE);
+  }
+
+  @Test
+  void generator_expression() {
+    assertThat(lastExpression(
+      """
+      x = (num**2 for num in numbers())
+      x
+      """
+    ).typeV2().unwrappedType()).isEqualTo(PythonType.UNKNOWN);
+  }
+
   @Test
   void return_type_of_call_expression_1() {
     assertThat(lastExpression(
diff --git a/python-frontend/src/test/java/org/sonar/python/types/v2/TypesTestUtils.java b/python-frontend/src/test/java/org/sonar/python/types/v2/TypesTestUtils.java
@@ -36,6 +36,8 @@ public class TypesTestUtils {
   public static final PythonType BOOL_TYPE = BUILTINS.resolveMember("bool").get();
   public static final PythonType STR_TYPE = BUILTINS.resolveMember("str").get();
   public static final PythonType LIST_TYPE = BUILTINS.resolveMember("list").get();
+  public static final PythonType SET_TYPE = BUILTINS.resolveMember("set").get();
+  public static final PythonType DICT_TYPE = BUILTINS.resolveMember("dict").get();
   public static final PythonType NONE_TYPE = BUILTINS.resolveMember("NoneType").get();
 
   public static FileInput parseAndInferTypes(String... code) {
diff --git a/sonar-python-plugin/src/main/java/org/sonar/plugins/python/SymbolVisitor.java b/sonar-python-plugin/src/main/java/org/sonar/plugins/python/SymbolVisitor.java
@@ -26,6 +26,7 @@
 import org.sonar.api.batch.sensor.symbol.NewSymbolTable;
 import org.sonar.plugins.python.api.tree.ClassDef;
 import org.sonar.plugins.python.api.tree.ComprehensionExpression;
+import org.sonar.plugins.python.api.tree.DictCompExpression;
 import org.sonar.plugins.python.api.tree.FileInput;
 import org.sonar.plugins.python.api.tree.FunctionDef;
 import org.sonar.plugins.python.api.tree.LambdaExpression;
@@ -69,7 +70,7 @@ public void visitPyListOrSetCompExpression(ComprehensionExpression tree) {
   }
 
   @Override
-  public void visitDictCompExpression(DictCompExpressionImpl tree) {
+  public void visitDictCompExpression(DictCompExpression tree) {
     tree.localVariables().forEach(this::handleSymbol);
     super.visitDictCompExpression(tree);
   }

Original file line number	Diff line number	Diff line change
`@@ -43,6 +43,7 @@`
`43`	`43`	`import org.sonar.plugins.python.api.tree.ComprehensionExpression;`
`44`	`44`	`import org.sonar.plugins.python.api.tree.ComprehensionIf;`
`45`	`45`	`import org.sonar.plugins.python.api.tree.ConditionalExpression;`
	`46`	`+import org.sonar.plugins.python.api.tree.DictCompExpression;`
`46`	`47`	`import org.sonar.plugins.python.api.tree.Expression;`
`47`	`48`	`import org.sonar.plugins.python.api.tree.ExpressionList;`
`48`	`49`	`import org.sonar.plugins.python.api.tree.FunctionDef;`
`@@ -194,7 +195,7 @@ public void visitComprehensionIf(ComprehensionIf tree) {`
`194`	`195`	`}`
`195`	`196`
`196`	`197`	`@Override`
`197`		`- public void visitDictCompExpression(DictCompExpressionImpl tree) {`
	`198`	`+ public void visitDictCompExpression(DictCompExpression tree) {`
`198`	`199`	`// ignore, not broken down in the cfg`
`199`	`200`	`}`
`200`	`201`
Original file line number	Diff line number	Diff line change
`@@ -431,7 +431,7 @@ public void visitKeyValuePattern(KeyValuePattern keyValuePattern) {`
`431`	`431`	`}`
`432`	`432`
`433`	`433`	`@Override`
`434`		`- public void visitDictCompExpression(DictCompExpressionImpl tree) {`
	`434`	`+ public void visitDictCompExpression(DictCompExpression tree) {`
`435`	`435`	`scan(tree.keyExpression());`
`436`	`436`	`scan(tree.valueExpression());`
`437`	`437`	`scan(tree.comprehensionFor());`