SONARPY-1984 Infer type of function parameters based on its type hints (#1854)

Author: maksim-grebeniuk-sonarsource

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

@@ -28,6 +28,7 @@
 import org.sonar.plugins.python.api.PythonFile;
 import org.sonar.plugins.python.api.cfg.ControlFlowGraph;
 import org.sonar.plugins.python.api.tree.BaseTreeVisitor;
+import org.sonar.plugins.python.api.tree.Expression;
 import org.sonar.plugins.python.api.tree.FileInput;
 import org.sonar.plugins.python.api.tree.FunctionDef;
 import org.sonar.plugins.python.api.tree.Name;
@@ -130,11 +131,13 @@ private static void inferTypesAndMemberAccessSymbols(Tree scopeTree,
   }
 
   private static void flowSensitiveTypeInference(ControlFlowGraph cfg, Set<SymbolV2> trackedVars, PropagationVisitor propagationVisitor) {
-    // TODO: infer parameter type based on type hint or default value assignement
+    // TODO: infer parameter type based on default value assignement
     var parameterTypes = trackedVars
       .stream()
-      .filter(s -> s.usages().stream().map(UsageV2::kind).anyMatch(UsageV2.Kind.PARAMETER::equals))
-      .collect(Collectors.toMap(SymbolV2::name, v -> PythonType.UNKNOWN));
+      .filter(symbol -> symbol.usages()
+        .stream()
+        .anyMatch(usage -> usage.kind() == UsageV2.Kind.PARAMETER))
+      .collect(Collectors.toMap(SymbolV2::name, TypeInferenceV2::getParameterType));
 
     FlowSensitiveTypeInference flowSensitiveTypeInference = new FlowSensitiveTypeInference(
       trackedVars,
@@ -146,6 +149,18 @@ private static void flowSensitiveTypeInference(ControlFlowGraph cfg, Set<SymbolV
     flowSensitiveTypeInference.compute(cfg);
   }
 
+  private static PythonType getParameterType(SymbolV2 symbol) {
+    return symbol.usages()
+      .stream()
+      .filter(usage -> usage.kind() == UsageV2.Kind.PARAMETER)
+      .map(UsageV2::tree)
+      .filter(Expression.class::isInstance)
+      .map(Expression.class::cast)
+      .map(Expression::typeV2)
+      .findFirst()
+      .orElse(PythonType.UNKNOWN);
+  }
+
   private static Set<SymbolV2> getTrackedVars(Set<SymbolV2> localVariables, Set<Name> assignedNames) {
     Set<SymbolV2> trackedVars = new HashSet<>();
     for (SymbolV2 variable : localVariables) {

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

@@ -31,6 +31,7 @@
 import org.sonar.plugins.python.api.tree.ArgList;
 import org.sonar.plugins.python.api.tree.AssignmentStatement;
 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.DictionaryLiteral;
 import org.sonar.plugins.python.api.tree.DottedName;
@@ -44,12 +45,15 @@
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.NoneExpression;
 import org.sonar.plugins.python.api.tree.NumericLiteral;
+import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.QualifiedExpression;
 import org.sonar.plugins.python.api.tree.RegularArgument;
 import org.sonar.plugins.python.api.tree.SetLiteral;
 import org.sonar.plugins.python.api.tree.StringLiteral;
+import org.sonar.plugins.python.api.tree.SubscriptionExpression;
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.plugins.python.api.tree.Tuple;
+import org.sonar.plugins.python.api.tree.TypeAnnotation;
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.semantic.v2.ClassTypeBuilder;
 import org.sonar.python.semantic.v2.FunctionTypeBuilder;
@@ -71,6 +75,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.TypeSource;
 import org.sonar.python.types.v2.UnionType;
 
 import static org.sonar.python.semantic.SymbolUtils.pathOf;
@@ -114,23 +119,23 @@ public void visitTuple(Tuple tuple) {
     }
     ModuleType builtins = this.projectLevelTypeTable.getModule();
     PythonType tupleType = builtins.resolveMember("tuple").orElse(PythonType.UNKNOWN);
-    ((TupleImpl) tuple).typeV2(new ObjectType(tupleType,  attributes, new ArrayList<>()));
+    ((TupleImpl) tuple).typeV2(new ObjectType(tupleType, attributes, new ArrayList<>()));
   }
 
   @Override
   public void visitDictionaryLiteral(DictionaryLiteral dictionaryLiteral) {
     super.visitDictionaryLiteral(dictionaryLiteral);
     ModuleType builtins = this.projectLevelTypeTable.getModule();
     PythonType dictType = builtins.resolveMember("dict").orElse(PythonType.UNKNOWN);
-    ((DictionaryLiteralImpl) dictionaryLiteral).typeV2(new ObjectType(dictType,  new ArrayList<>(), new ArrayList<>()));
+    ((DictionaryLiteralImpl) dictionaryLiteral).typeV2(new ObjectType(dictType, new ArrayList<>(), new ArrayList<>()));
   }
 
   @Override
   public void visitSetLiteral(SetLiteral setLiteral) {
     super.visitSetLiteral(setLiteral);
     ModuleType builtins = this.projectLevelTypeTable.getModule();
     PythonType setType = builtins.resolveMember("set").orElse(PythonType.UNKNOWN);
-    ((SetLiteralImpl) setLiteral).typeV2(new ObjectType(setType,  new ArrayList<>(), new ArrayList<>()));
+    ((SetLiteralImpl) setLiteral).typeV2(new ObjectType(setType, new ArrayList<>(), new ArrayList<>()));
   }
 
   @Override
@@ -329,6 +334,43 @@ public void visitAssignmentStatement(AssignmentStatement assignmentStatement) {
       });
   }
 
+  @Override
+  public void visitParameter(Parameter parameter) {
+    scan(parameter.typeAnnotation());
+    scan(parameter.defaultValue());
+    Optional.ofNullable(parameter.typeAnnotation())
+      .map(TypeAnnotation::expression)
+      .map(TrivialTypeInferenceVisitor::resolveTypeAnnotationExpressionType)
+      .ifPresent(type -> setTypeToName(parameter.name(), type));
+    scan(parameter.name());
+  }
+
+  private static PythonType resolveTypeAnnotationExpressionType(Expression expression) {
+    if (expression instanceof Name name && name.typeV2() != PythonType.UNKNOWN) {
+      return new ObjectType(name.typeV2(), TypeSource.TYPE_HINT);
+    } else if (expression instanceof SubscriptionExpression subscriptionExpression && subscriptionExpression.object().typeV2() != PythonType.UNKNOWN) {
+      var candidateTypes = subscriptionExpression.subscripts()
+        .expressions()
+        .stream()
+        .map(Expression::typeV2)
+        .distinct()
+        .toList();
+
+      var elementsType = UnionType.or(candidateTypes);
+
+      var attributes = new ArrayList<PythonType>();
+      attributes.add(new ObjectType(elementsType, TypeSource.TYPE_HINT));
+      return new ObjectType(subscriptionExpression.object().typeV2(), attributes, new ArrayList<>(), TypeSource.TYPE_HINT);
+    } else if (expression instanceof BinaryExpression binaryExpression) {
+      var left = resolveTypeAnnotationExpressionType(binaryExpression.leftOperand());
+      var right = resolveTypeAnnotationExpressionType(binaryExpression.rightOperand());
+      // TODO: we need to make a decision on should here be a union type of object types or an object type of a union type.
+      //  ATM it is blocked by the generic types resolution redesign
+      return UnionType.or(left, right);
+    }
+    return PythonType.UNKNOWN;
+  }
+
   @Override
   public void visitQualifiedExpression(QualifiedExpression qualifiedExpression) {
     scan(qualifiedExpression.qualifier());

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

@@ -27,10 +27,18 @@
 import org.sonar.plugins.python.api.LocationInFile;
 
 @Beta
-public record ObjectType(PythonType type, List<PythonType> attributes, List<Member> members) implements PythonType {
+public record ObjectType(PythonType type, List<PythonType> attributes, List<Member> members, TypeSource typeSource) implements PythonType {
 
   public ObjectType(PythonType type) {
-    this(type, new ArrayList<>(), new ArrayList<>());
+    this(type, TypeSource.EXACT);
+  }
+
+  public ObjectType(PythonType type, TypeSource typeSource) {
+    this(type, new ArrayList<>(), new ArrayList<>(), typeSource);
+  }
+
+  public ObjectType(PythonType type, List<PythonType> attributes, List<Member> members) {
+    this(type, attributes, members, TypeSource.EXACT);
   }
 
   @Override

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

@@ -74,4 +74,9 @@ default Optional<LocationInFile> definitionLocation() {
   default PythonType unwrappedType() {
     return this;
   }
+
+  @Beta
+  default TypeSource typeSource() {
+    return TypeSource.EXACT;
+  }
 }

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

@@ -42,6 +42,11 @@ public TypeCheckBuilder instancesHaveMember(String memberName) {
     return this;
   }
 
+  public TypeCheckBuilder isTypeHintTypeSource() {
+    predicates.add(new TypeSourceMatcherTypePredicate(TypeSource.TYPE_HINT));
+    return this;
+  }
+
   public TriBool check(PythonType pythonType) {
     TriBool result = TriBool.TRUE;
     for (TypePredicate predicate : predicates) {
@@ -84,4 +89,12 @@ public TriBool test(PythonType pythonType) {
       return TriBool.FALSE;
     }
   }
+
+  record TypeSourceMatcherTypePredicate(TypeSource typeSource) implements TypePredicate {
+
+    @Override
+    public TriBool test(PythonType pythonType) {
+      return pythonType.typeSource() == typeSource ? TriBool.TRUE : TriBool.FALSE;
+    }
+  }
 }

python-frontend/src/main/java/org/sonar/python/types/v2/TypeSource.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.types.v2;
+
+public enum TypeSource {
+  EXACT,
+  TYPE_HINT
+}

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

@@ -0,0 +1,41 @@
+/*
+ * 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;
+
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.sonar.python.types.v2.ObjectType;
+import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.TriBool;
+import org.sonar.python.types.v2.TypeCheckBuilder;
+import org.sonar.python.types.v2.TypeSource;
+
+class TypeCheckerBuilderTest {
+
+  @Test
+  void typeSourceTest() {
+    var builder = new TypeCheckBuilder(null).isTypeHintTypeSource();
+    Assertions.assertThat(builder.check(new ObjectType(PythonType.UNKNOWN, TypeSource.TYPE_HINT)))
+      .isEqualTo(TriBool.TRUE);
+    Assertions.assertThat(builder.check(new ObjectType(PythonType.UNKNOWN, TypeSource.EXACT)))
+      .isEqualTo(TriBool.FALSE);
+  }
+
+}