SONARPY-2000 Propagate types of imported names (#1863)

Author: guillaume-dequenne

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

@@ -29,6 +29,7 @@
 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.ImportName;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.Statement;
@@ -41,13 +42,13 @@
 public class FlowSensitiveTypeInference extends ForwardAnalysis {
   private final Set<SymbolV2> trackedVars;
   private final Map<Statement, Assignment> assignmentsByAssignmentStatement;
-  private final Map<Statement, Definition> definitionsByDefinitionStatement;
+  private final Map<Statement, Set<Definition>> definitionsByDefinitionStatement;
   private final Map<String, PythonType> parameterTypesByName;
 
   public FlowSensitiveTypeInference(
     Set<SymbolV2> trackedVars,
     Map<Statement, Assignment> assignmentsByAssignmentStatement,
-    Map<Statement, Definition> definitionsByDefinitionStatement,
+    Map<Statement, Set<Definition>> definitionsByDefinitionStatement,
     Map<String, PythonType> parameterTypesByName
   ) {
     this.trackedVars = trackedVars;
@@ -86,7 +87,9 @@ public void updateProgramState(Tree element, ProgramState programState) {
         updateTree(assignment.variable(), state);
       }
     } else if (element instanceof FunctionDef functionDef) {
-      handleDefinition(functionDef, state);
+      handleDefinitions(functionDef, state);
+    } else if (element instanceof ImportName importName) {
+      handleDefinitions(importName, state);
     } else if (element instanceof Parameter parameter) {
       handleParameter(parameter, state);
     } else if (isForLoopAssignment(element)) {
@@ -150,13 +153,13 @@ private void handleAssignment(Statement assignmentStatement, TypeInferenceProgra
       });
   }
 
-  private void handleDefinition(Statement definitionStatement, TypeInferenceProgramState programState) {
+  private void handleDefinitions(Statement definitionStatement, TypeInferenceProgramState programState) {
     Optional.ofNullable(definitionsByDefinitionStatement.get(definitionStatement))
-      .ifPresent(definition -> {
-        SymbolV2 symbol = definition.lhsSymbol();
+      .ifPresent(definitions -> definitions.forEach(d -> {
+        SymbolV2 symbol = d.lhsSymbol();
         if (trackedVars.contains(symbol)) {
-          programState.setTypes(symbol, Set.of(definition.lhsName.typeV2()));
+          programState.setTypes(symbol, Set.of(d.lhsName.typeV2()));
         }
-      });
+      }));
   }
 }

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

@@ -33,6 +33,7 @@
 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.ImportName;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.Statement;
@@ -42,7 +43,7 @@
 public class PropagationVisitor extends BaseTreeVisitor {
   private final Map<SymbolV2, Set<Propagation>> propagationsByLhs;
   private final Map<Statement, Assignment> assignmentsByAssignmentStatement;
-  private final Map<Statement, Definition> definitionsByDefinitionStatement;
+  private final Map<Statement, Set<Definition>> definitionsByDefinitionStatement;
 
   public PropagationVisitor() {
     propagationsByLhs = new HashMap<>();
@@ -55,7 +56,7 @@ public Map<Statement, Assignment> assignmentsByAssignmentStatement() {
     return assignmentsByAssignmentStatement;
   }
 
-  public Map<Statement, Definition> definitionsByDefinitionStatement() {
+  public Map<Statement, Set<Definition>> definitionsByDefinitionStatement() {
     return definitionsByDefinitionStatement;
   }
 
@@ -69,7 +70,7 @@ public void visitFunctionDef(FunctionDef functionDef) {
     Name name = functionDef.name();
     var symbol = name.symbolV2();
     Definition definition = new Definition(symbol, name);
-    definitionsByDefinitionStatement.put(functionDef, definition);
+    definitionsByDefinitionStatement.computeIfAbsent(functionDef, k -> new HashSet<>()).add(definition);
     propagationsByLhs.computeIfAbsent(symbol, s -> new HashSet<>()).add(definition);
   }
 
@@ -113,6 +114,32 @@ public void visitAnnotatedAssignment(AnnotatedAssignment annotatedAssignment){
     }
   }
 
+  @Override
+  public void visitImportName(ImportName importName) {
+    super.visitImportName(importName);
+    importName.modules()
+      .forEach(aliasedName -> {
+        var names = aliasedName.dottedName().names();
+        Name alias = aliasedName.alias();
+        if (alias != null) {
+          SymbolV2 aliasSymbol = alias.symbolV2();
+          Definition definition = new Definition(aliasSymbol, alias);
+          definitionsByDefinitionStatement.computeIfAbsent(importName, k -> new HashSet<>()).add(definition);
+          propagationsByLhs.computeIfAbsent(aliasSymbol, s -> new HashSet<>()).add(definition);
+        } else {
+          for (int i = names.size() - 1; i >= 0; i--) {
+            Name name = names.get(i);
+            SymbolV2 symbolV2 = name.symbolV2();
+            if (symbolV2 != null) {
+              Definition definition = new Definition(symbolV2, name);
+              definitionsByDefinitionStatement.computeIfAbsent(importName, k -> new HashSet<>()).add(definition);
+              propagationsByLhs.computeIfAbsent(symbolV2, s -> new HashSet<>()).add(definition);
+            }
+          }
+        }
+      });
+  }
+
   @Override
   public void visitForStatement(ForStatement forStatement) {
     scan(forStatement.testExpressions());

python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java

@@ -1633,6 +1633,46 @@ def bar(self): ...
     assertThat(unionType.candidates()).containsExactlyInAnyOrder(classA, classB);
   }
 
+  @Test
+  @Disabled("Duplicate PythonType for NONE_TYPE")
+  void imported_symbol_call_return_type() {
+    assertThat(lastExpression(
+      """
+      import fcntl
+      ret = fcntl.flock(..., ...)
+      ret
+      """
+    ).typeV2()).isEqualTo(NONE_TYPE);
+  }
+
+  @Test
+  void basic_imported_symbol() {
+    assertThat(lastExpression(
+      """
+      import fcntl
+      fcntl
+      """
+    ).typeV2()).isInstanceOf(ModuleType.class);
+  }
+
+  @Test
+  void basic_imported_symbols() {
+    FileInput fileInput = inferTypes(
+      """
+      import fcntl, math
+      fcntl
+      math
+      """
+    );
+    PythonType fnctlModule = ((ExpressionStatement) fileInput.statements().statements().get(1)).expressions().get(0).typeV2();
+    assertThat(fnctlModule).isInstanceOf(ModuleType.class);
+    assertThat(fnctlModule.name()).isEqualTo("fcntl");
+    PythonType mathModule = ((ExpressionStatement) fileInput.statements().statements().get(2)).expressions().get(0).typeV2();
+    assertThat(mathModule).isInstanceOf(ModuleType.class);
+    assertThat(mathModule.name()).isEqualTo("math");
+    assertThat(((UnionType) mathModule.resolveMember("acos").get()).candidates()).allMatch(FunctionType.class::isInstance);
+  }
+
   private static FileInput inferTypes(String lines) {
     return inferTypes(lines, PROJECT_LEVEL_TYPE_TABLE);
   }