SONARPY-1997 Propagate function return types to the result of call expressions (#1857)

Author: guillaume-dequenne

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

@@ -33,7 +33,9 @@
 import org.sonar.plugins.python.api.symbols.ClassSymbol;
 import org.sonar.plugins.python.api.symbols.FunctionSymbol;
 import org.sonar.plugins.python.api.symbols.Symbol;
+import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.semantic.ClassSymbolImpl;
+import org.sonar.python.semantic.FunctionSymbolImpl;
 import org.sonar.python.semantic.ProjectLevelSymbolTable;
 import org.sonar.python.semantic.SymbolImpl;
 import org.sonar.python.types.v2.ClassType;
@@ -88,8 +90,9 @@ private static ModuleType createEmptyModule(String moduleName, ModuleType parent
 
   private ModuleType createModuleFromSymbols(@Nullable String name, @Nullable ModuleType parent, Collection<Symbol> symbols) {
     var members = new HashMap<String, PythonType>();
+    Map<Symbol, PythonType> createdTypesBySymbol = new HashMap<>();
     symbols.forEach(symbol -> {
-      var type = convertToType(symbol, new HashMap<>());
+      var type = convertToType(symbol, createdTypesBySymbol);
       members.put(symbol.name(), type);
     });
     var module = new ModuleType(name, parent);
@@ -102,7 +105,7 @@ private ModuleType createModuleFromSymbols(@Nullable String name, @Nullable Modu
     return module;
   }
 
-  private static PythonType convertToFunctionType(FunctionSymbol symbol, Map<Symbol, PythonType> createdTypesBySymbol) {
+  private PythonType convertToFunctionType(FunctionSymbol symbol, Map<Symbol, PythonType> createdTypesBySymbol) {
     if (createdTypesBySymbol.containsKey(symbol)) {
       return createdTypesBySymbol.get(symbol);
     }
@@ -112,11 +115,18 @@ private static PythonType convertToFunctionType(FunctionSymbol symbol, Map<Symbo
       .map(SymbolsModuleTypeProvider::convertParameter)
       .toList();
 
+    InferredType inferredType = ((FunctionSymbolImpl) symbol).declaredReturnType();
+    ClassSymbol classSymbol = inferredType.runtimeTypeSymbol();
+    var returnType = PythonType.UNKNOWN;
+    if (classSymbol != null) {
+      returnType = convertToType(classSymbol, createdTypesBySymbol);
+    }
+
     FunctionTypeBuilder functionTypeBuilder =
       new FunctionTypeBuilder(symbol.name())
         .withAttributes(List.of())
         .withParameters(parameters)
-        .withReturnType(PythonType.UNKNOWN)
+        .withReturnType(returnType)
         .withAsynchronous(symbol.isAsynchronous())
         .withHasDecorators(symbol.hasDecorators())
         .withInstanceMethod(symbol.isInstanceMethod())
@@ -187,5 +197,4 @@ private PythonType convertToType(Symbol symbol, Map<Symbol, PythonType> createdT
       case OTHER -> PythonType.UNKNOWN;
     };
   }
-
 }

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

@@ -91,7 +91,7 @@ public class TypeShed {
   public TypeShed(ProjectLevelSymbolTable projectLevelSymbolTable) {
     // workaround to initialize supported python versions used in ClassSymbolImpl
     // TODO: remove once v2 types model will be populated from TypeShed bypassing conversion to symbols
-    org.sonar.python.types.TypeShed.builtinSymbols();
+    builtins = org.sonar.python.types.TypeShed.builtinSymbols();
     typeShedSymbols = new HashMap<>();
     modulesInProgress = new HashSet<>();
     this.projectLevelSymbolTable = projectLevelSymbolTable;

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

@@ -47,8 +47,10 @@
 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.FunctionType;
 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.plugins.python.api.symbols.Symbol.Kind.CLASS;
 import static org.sonar.plugins.python.api.tree.Tree.Kind.SUBSCRIPTION;
@@ -188,6 +190,21 @@ public PythonType typeV2() {
     if (callee().typeV2() instanceof ClassType classType) {
       return new ObjectType(classType);
     }
+    if (callee().typeV2() instanceof FunctionType functionType) {
+      return functionType.returnType();
+    }
+    if (callee().typeV2() instanceof UnionType unionType) {
+      PythonType result = null;
+      for (PythonType candidate : unionType.candidates()) {
+        if (candidate instanceof ClassType classType) {
+          result = UnionType.or(result, classType);
+        }
+        if (candidate instanceof FunctionType functionType) {
+          result = UnionType.or(result, functionType.returnType());
+        }
+      }
+      return result == null ? PythonType.UNKNOWN : new ObjectType(result);
+    }
     return PythonType.UNKNOWN;
   }
 }

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

@@ -26,6 +26,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import org.sonar.api.Beta;
 
 @Beta
@@ -71,7 +72,13 @@ public static PythonType or(Collection<PythonType> candidates) {
   }
 
   @Beta
-  public static PythonType or(PythonType type1, PythonType type2) {
+  public static PythonType or(@Nullable PythonType type1, @Nullable PythonType type2) {
+    if (type1 == null) {
+      return type2;
+    }
+    if (type2 == null) {
+      return type1;
+    }
     if (type1.equals(PythonType.UNKNOWN) || type2.equals(PythonType.UNKNOWN)) {
       return PythonType.UNKNOWN;
     }

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

@@ -1574,6 +1574,65 @@ def reassigned_param(a, param):
     Assertions.assertThat(paramType).isSameAs(PythonType.UNKNOWN);
   }
 
+  @Test
+  void return_type_of_call_expression_1() {
+    assertThat(lastExpression(
+      """
+      x = [1,2,3]
+      a = x.append(42)
+      a
+      """
+    ).typeV2().unwrappedType()).isEqualTo(NONE_TYPE);
+  }
+
+  @Test
+  void return_type_of_call_expression_2() {
+    assertThat(lastExpression(
+      """
+      x = [1,2,3]
+      a = x.sort()
+      a
+      """
+    ).typeV2().unwrappedType()).isEqualTo(NONE_TYPE);
+  }
+
+  @Test
+  void return_type_of_call_expression_union_type() {
+    FileInput fileInput = inferTypes(
+      """
+        class A:
+          def foo(self): ...
+        class B:
+          def bar(self): ...
+        a = A
+        b = B
+        if cond:
+          x = a
+        else:
+          x = b
+        y = x()
+        y
+        """
+    );
+    var classA = TreeUtils.firstChild(fileInput.statements().statements().get(0), ClassDef.class::isInstance)
+      .map(ClassDef.class::cast)
+      .map(ClassDef::name)
+      .map(Expression::typeV2)
+      .map(ClassType.class::cast)
+      .get();
+
+    var classB = TreeUtils.firstChild(fileInput.statements().statements().get(1), ClassDef.class::isInstance)
+      .map(ClassDef.class::cast)
+      .map(ClassDef::name)
+      .map(Expression::typeV2)
+      .map(ClassType.class::cast)
+      .get();
+
+    assertThat(((ExpressionStatement) fileInput.statements().statements().get(6)).expressions().get(0).typeV2()).isInstanceOf(ObjectType.class);
+    UnionType unionType = (UnionType) ((ExpressionStatement) fileInput.statements().statements().get(6)).expressions().get(0).typeV2().unwrappedType();
+    assertThat(unionType.candidates()).containsExactlyInAnyOrder(classA, classB);
+  }
+
   private static FileInput inferTypes(String lines) {
     return inferTypes(lines, PROJECT_LEVEL_TYPE_TABLE);
   }

python-frontend/src/test/java/org/sonar/python/types/v2/UnionTypeTest.java

@@ -77,6 +77,14 @@ void or_with_union_type() {
     assertThat(((UnionType) result).candidates()).containsExactlyInAnyOrder(INT_TYPE, STR_TYPE, BOOL_TYPE);
   }
 
+  @Test
+  void or_with_null() {
+    PythonType type = UnionType.or(INT_TYPE, null);
+    assertThat(type).isEqualTo(INT_TYPE);
+    type = UnionType.or(null, INT_TYPE);
+    assertThat(type).isEqualTo(INT_TYPE);
+  }
+
   @Test
   void or_unionType() {
     FileInput fileInput = parseAndInferTypes("42;\"hello\"");