SONARPY-2023 Translate type source for binary expressions (#1896)

Author: maksim-grebeniuk-sonarsource

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

@@ -34,6 +34,11 @@
 import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.types.HasTypeDependencies;
 import org.sonar.python.types.InferredTypes;
+import org.sonar.python.types.v2.ClassType;
+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.types.InferredTypes.DECL_INT;
 import static org.sonar.python.types.InferredTypes.DECL_STR;
@@ -137,6 +142,22 @@ public InferredType type() {
     return InferredTypes.anyType();
   }
 
+  @Override
+  public PythonType typeV2() {
+    if (is(Kind.AND, Kind.OR)) {
+      return UnionType.or(leftOperand.typeV2(), rightOperand.typeV2());
+    }
+    if (is(Kind.PLUS)
+        && leftOperand.typeV2() instanceof ObjectType leftObjectType
+        && leftObjectType.unwrappedType() instanceof ClassType leftClassType
+        && rightOperand.typeV2() instanceof ObjectType rightObjectType
+        && rightObjectType.unwrappedType() instanceof ClassType rightClassType
+        && leftClassType == rightClassType) {
+      return new ObjectType(leftClassType, TypeSource.min(leftObjectType.typeSource(), rightObjectType.typeSource()));
+    }
+    return PythonType.UNKNOWN;
+  }
+
   private static boolean shouldReturnDeclaredStr(InferredType leftType, InferredType rightType) {
     return (leftType.equals(DECL_STR) && rightType.equals(DECL_STR)) ||
       (leftType.equals(STR) && rightType.equals(DECL_STR)) ||

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

@@ -19,7 +19,26 @@
  */
 package org.sonar.python.types.v2;
 
+import java.util.Comparator;
+import java.util.stream.Stream;
+
 public enum TypeSource {
-  EXACT,
-  TYPE_HINT
+  TYPE_HINT(0),
+  EXACT(1);
+
+  private final int score;
+
+  TypeSource(int score) {
+    this.score = score;
+  }
+
+  public int score() {
+    return score;
+  }
+
+  public static TypeSource min(TypeSource... typeSources) {
+    return Stream.of(typeSources)
+      .min(Comparator.comparing(TypeSource::score))
+      .orElse(EXACT);
+  }
 }

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

@@ -21,6 +21,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
@@ -64,6 +65,13 @@ public boolean isCompatibleWith(PythonType another) {
       .anyMatch(candidate -> candidate.isCompatibleWith(another));
   }
 
+  @Override
+  public TypeSource typeSource() {
+    return candidates.stream().map(PythonType::typeSource)
+      .min(Comparator.comparing(TypeSource::score))
+      .orElse(TypeSource.EXACT);
+  }
+
   @Beta
   public static PythonType or(Collection<PythonType> candidates) {
     if (candidates.isEmpty()) {

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

@@ -2089,6 +2089,116 @@ def foo(x: list):
     Assertions.assertThat(xType.unwrappedType()).isSameAs(LIST_TYPE);
   }
 
+  @Test
+  void assignmentOrTest() {
+    var yType = (UnionType) lastExpression("""
+      def foo(x: int):
+        y = x or "str"
+        y
+      """).typeV2();
+
+    Assertions.assertThat(yType.candidates())
+      .allMatch(ObjectType.class::isInstance)
+      .extracting(PythonType::unwrappedType)
+      .containsOnly(INT_TYPE, STR_TYPE);
+    Assertions.assertThat(yType.typeSource()).isSameAs(TypeSource.TYPE_HINT);
+  }
+
+  @Test
+  void assignmentPlusTest() {
+    var fileInput = inferTypes("""
+      class A: ...
+      def foo(x: int, y: str):
+        a = x + 1
+        b = x + y
+        c = A + x
+        d = x + A
+        e = 2 + 3
+        a
+        b
+        c
+        d
+        e
+      """);
+
+    var statements = TreeUtils.firstChild(fileInput, FunctionDef.class::isInstance)
+      .map(FunctionDef.class::cast)
+      .map(FunctionDef::body)
+      .map(StatementList::statements)
+      .orElseGet(List::of);
+
+    var aType = ((ExpressionStatement) statements.get(statements.size() - 5)).expressions().get(0).typeV2();
+    var bType = ((ExpressionStatement) statements.get(statements.size() - 4)).expressions().get(0).typeV2();
+    var cType = ((ExpressionStatement) statements.get(statements.size() - 3)).expressions().get(0).typeV2();
+    var dType = ((ExpressionStatement) statements.get(statements.size() - 2)).expressions().get(0).typeV2();
+    var eType = ((ExpressionStatement) statements.get(statements.size() - 1)).expressions().get(0).typeV2();
+
+    Assertions.assertThat(aType.unwrappedType()).isSameAs(INT_TYPE);
+    Assertions.assertThat(aType.typeSource()).isSameAs(TypeSource.TYPE_HINT);
+
+    Assertions.assertThat(bType).isSameAs(PythonType.UNKNOWN);
+    Assertions.assertThat(cType).isSameAs(PythonType.UNKNOWN);
+    Assertions.assertThat(dType).isSameAs(PythonType.UNKNOWN);
+
+    Assertions.assertThat(eType.unwrappedType()).isSameAs(INT_TYPE);
+    Assertions.assertThat(eType.typeSource()).isSameAs(TypeSource.EXACT);
+  }
+
+  @Test
+  void assignmentPlusTest2() {
+    var fileInput = inferTypes("""
+      class A: ...
+      def foo(x):
+        if x:
+          t = int
+        else:
+          t = str
+        a = t()
+        b = 1 + a
+        c = a + 1
+        a
+        b
+        c
+      """);
+
+    var statements = TreeUtils.firstChild(fileInput, FunctionDef.class::isInstance)
+      .map(FunctionDef.class::cast)
+      .map(FunctionDef::body)
+      .map(StatementList::statements)
+      .orElseGet(List::of);
+
+    var aType = ((ExpressionStatement) statements.get(statements.size() - 3)).expressions().get(0).typeV2();
+    var bType = ((ExpressionStatement) statements.get(statements.size() - 2)).expressions().get(0).typeV2();
+    var cType = ((ExpressionStatement) statements.get(statements.size() - 1)).expressions().get(0).typeV2();
+
+    Assertions.assertThat(aType).isInstanceOf(ObjectType.class);
+    Assertions.assertThat(aType.unwrappedType()).isInstanceOf(UnionType.class);
+    var candidates = ((UnionType) aType.unwrappedType()).candidates();
+    Assertions.assertThat(candidates).containsOnly(INT_TYPE, STR_TYPE);
+
+    Assertions.assertThat(bType).isSameAs(PythonType.UNKNOWN);
+    Assertions.assertThat(cType).isSameAs(PythonType.UNKNOWN);
+  }
+
+  @Test
+  void assignmentMinusTest() {
+    var fileInput = inferTypes("""
+      def foo():
+        f = 2 - 3
+        f
+      """);
+
+    var statements = TreeUtils.firstChild(fileInput, FunctionDef.class::isInstance)
+      .map(FunctionDef.class::cast)
+      .map(FunctionDef::body)
+      .map(StatementList::statements)
+      .orElseGet(List::of);
+
+    var fType = ((ExpressionStatement) statements.get(statements.size() - 1)).expressions().get(0).typeV2();
+    Assertions.assertThat(fType).isSameAs(PythonType.UNKNOWN);
+  }
+
+
   private static FileInput inferTypes(String lines) {
     return inferTypes(lines, PROJECT_LEVEL_TYPE_TABLE);
   }