SONARPY-1836 Infer item type for basic loop over list iteration (#1805)

Author: maksim-grebeniuk-sonarsource

python-checks/src/test/resources/checks/nonCallableCalled.py

@@ -310,3 +310,42 @@ def nested_function_in_try_catch():
         ...
     def bar():
         foo()
+
+
+
+def f1():
+    ...
+
+def f2():
+    ...
+
+def callable_from_loop_try_except():
+    l = [f1, f2]
+    try:
+        for i in l:
+            i()
+    except:
+        ...
+
+def non_callable_from_loop_try_except():
+    l = ["f1", "f2"]
+    try:
+        for i in l:
+            i() # Noncompliant
+    except:
+        ...
+
+
+def callable_from_loop_append_noncallable():
+    l = [f1, f2]
+    l.append("1")
+    for i in l:
+        i() # FN
+
+
+def callable_from_loop_append_noncallable():
+    l = ["1"]
+    possible_modiffication(l)
+    for i in l:
+        # FP
+        i() # Noncompliant

python-frontend/src/main/java/org/sonar/python/semantic/v2/types/FlowSensitiveTypeInference.java

@@ -27,6 +27,7 @@
 import org.sonar.plugins.python.api.tree.AssignmentStatement;
 import org.sonar.plugins.python.api.tree.CompoundAssignmentStatement;
 import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.ForStatement;
 import org.sonar.plugins.python.api.tree.FunctionDef;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.Parameter;
@@ -88,6 +89,8 @@ public void updateProgramState(Tree element, ProgramState programState) {
       handleDefinition(functionDef, state);
     } else if (element instanceof Parameter parameter) {
       handleParameter(parameter, state);
+    } else if (isForLoopAssignment(element)) {
+      handleLoopAssignment(element, state);
     } else {
       // Here we should run "isinstance" visitor when we handle declared types, to avoid FPs when type guard checks are made
       updateTree(element, state);
@@ -110,6 +113,27 @@ private static void updateTree(Tree tree, TypeInferenceProgramState state) {
     tree.accept(new ProgramStateTypeInferenceVisitor(state));
   }
 
+
+  private static boolean isForLoopAssignment(Tree tree) {
+    return tree instanceof Name && tree.parent() instanceof ForStatement forStatement && forStatement.expressions().contains(tree);
+  }
+
+  private void handleLoopAssignment(Tree element, TypeInferenceProgramState state) {
+    Optional.of(element)
+      .map(Tree::parent)
+      .filter(ForStatement.class::isInstance)
+      .map(ForStatement.class::cast)
+      .ifPresent(forStatement -> {
+        forStatement.testExpressions().forEach(t -> updateTree(t, state));
+        Optional.ofNullable(assignmentsByAssignmentStatement.get(forStatement))
+          .filter(assignment -> trackedVars.contains(assignment.lhsSymbol()))
+          .ifPresent(assignment -> Optional.of(assignment)
+            .map(Assignment::rhsType)
+            .ifPresent(collectionItemType -> state.setTypes(assignment.lhsSymbol(), Set.of(collectionItemType)))
+          );
+      });
+  }
+
   private void handleAssignment(Statement assignmentStatement, TypeInferenceProgramState programState) {
     Optional.ofNullable(assignmentsByAssignmentStatement.get(assignmentStatement))
       .ifPresent(assignment -> {

python-frontend/src/main/java/org/sonar/python/semantic/v2/types/LoopAssignment.java

@@ -0,0 +1,50 @@
+/*
+ * 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.semantic.v2.types;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Stream;
+import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.Name;
+import org.sonar.python.semantic.v2.SymbolV2;
+import org.sonar.python.types.v2.ObjectType;
+import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.TriBool;
+
+public class LoopAssignment extends Assignment {
+  public LoopAssignment(SymbolV2 lhsSymbol, Name lhsName, Expression rhs, Map<SymbolV2, Set<Propagation>> propagationsByLhs) {
+    super(lhsSymbol, lhsName, rhs, propagationsByLhs);
+  }
+
+  @Override
+  public PythonType rhsType() {
+    return Optional.of(super.rhsType())
+      .filter(ObjectType.class::isInstance)
+      .map(ObjectType.class::cast)
+      .filter(t -> t.hasMember("__iter__") == TriBool.TRUE)
+      .map(ObjectType::attributes)
+      .map(Collection::stream)
+      .flatMap(Stream::findFirst)
+      .orElse(PythonType.UNKNOWN);
+  }
+}

python-frontend/src/main/java/org/sonar/python/semantic/v2/types/PropagationVisitor.java

@@ -31,11 +31,13 @@
 import org.sonar.plugins.python.api.tree.BaseTreeVisitor;
 import org.sonar.plugins.python.api.tree.CompoundAssignmentStatement;
 import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.ForStatement;
 import org.sonar.plugins.python.api.tree.FunctionDef;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.Statement;
 import org.sonar.python.semantic.v2.SymbolV2;
+import org.sonar.python.tree.NameImpl;
 
 public class PropagationVisitor extends BaseTreeVisitor {
   private final Map<SymbolV2, Set<Propagation>> propagationsByLhs;
@@ -111,10 +113,35 @@ public void visitAnnotatedAssignment(AnnotatedAssignment annotatedAssignment){
     }
   }
 
+  @Override
+  public void visitForStatement(ForStatement forStatement) {
+    scan(forStatement.testExpressions());
+    if (forStatement.testExpressions().size() == 1 && forStatement.expressions().size() == 1) {
+      forStatement
+        .testExpressions()
+        .stream()
+        .findFirst()
+        .ifPresent(rhsExpression -> forStatement.expressions().stream()
+            .findFirst()
+            .filter(NameImpl.class::isInstance)
+            .map(NameImpl.class::cast)
+            .ifPresent(i -> {
+              var symbol = i.symbolV2();
+              var assignment = new LoopAssignment(symbol, i, rhsExpression, propagationsByLhs);
+              assignmentsByAssignmentStatement.put(forStatement, assignment);
+              propagationsByLhs.computeIfAbsent(symbol, s -> new HashSet<>()).add(assignment);
+            })
+        );
+
+    }
+    scan(forStatement.body());
+    scan(forStatement.elseClause());
+  }
+
   private void processAssignment(Statement assignmentStatement, Expression lhsExpression, Expression rhsExpression){
     if (lhsExpression instanceof Name lhs && lhs.symbolV2() != null) {
       var symbol = lhs.symbolV2();
-      Assignment assignment = new Assignment(symbol, lhs, rhsExpression, propagationsByLhs);
+      var assignment = new Assignment(symbol, lhs, rhsExpression, propagationsByLhs);
       assignmentsByAssignmentStatement.put(assignmentStatement, assignment);
       propagationsByLhs.computeIfAbsent(symbol, s -> new HashSet<>()).add(assignment);
     }

python-frontend/src/main/java/org/sonar/python/semantic/v2/types/TrivialTypeInferenceVisitor.java

@@ -71,6 +71,7 @@
 import org.sonar.python.types.v2.ModuleType;
 import org.sonar.python.types.v2.ObjectType;
 import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.UnionType;
 
 import static org.sonar.python.semantic.SymbolUtils.pathOf;
 import static org.sonar.python.tree.TreeUtils.locationInFile;
@@ -153,11 +154,22 @@ public void visitNone(NoneExpression noneExpression) {
 
   @Override
   public void visitListLiteral(ListLiteral listLiteral) {
-    ModuleType builtins = this.projectLevelTypeTable.getModule();
+    var builtins = this.projectLevelTypeTable.getModule();
     scan(listLiteral.elements());
-    List<PythonType> pythonTypes = listLiteral.elements().expressions().stream().map(Expression::typeV2).distinct().toList();
+
+    var candidateTypes = listLiteral.elements()
+      .expressions()
+      .stream()
+      .map(Expression::typeV2)
+      .distinct()
+      .toList();
+
+    var elementsType = UnionType.or(candidateTypes);
+
+    var attributes = new ArrayList<PythonType>();
+    attributes.add(elementsType);
     PythonType listType = builtins.resolveMember("list").orElse(PythonType.UNKNOWN);
-    ((ListLiteralImpl) listLiteral).typeV2(new ObjectType(listType, pythonTypes, new ArrayList<>()));
+    ((ListLiteralImpl) listLiteral).typeV2(new ObjectType(listType, attributes, new ArrayList<>()));
   }
 
   @Override

python-frontend/src/main/java/org/sonar/python/types/v2/UnionType.java

@@ -20,6 +20,7 @@
 package org.sonar.python.types.v2;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
@@ -60,6 +61,12 @@ public boolean isCompatibleWith(PythonType another) {
       .anyMatch(candidate -> candidate.isCompatibleWith(another));
   }
 
+  public static PythonType or(Collection<PythonType> candidates) {
+    return candidates
+      .stream()
+      .reduce(new UnionType(new HashSet<>()), UnionType::or);
+  }
+
   public static PythonType or(PythonType type1, PythonType type2) {
     if (type1.equals(PythonType.UNKNOWN) || type2.equals(PythonType.UNKNOWN)) {
       return PythonType.UNKNOWN;
@@ -70,6 +77,9 @@ public static PythonType or(PythonType type1, PythonType type2) {
     Set<PythonType> types = new HashSet<>();
     addTypes(type1, types);
     addTypes(type2, types);
+    if (types.size() == 1) {
+      return types.iterator().next();
+    }
     return new UnionType(types);
   }