SONARPY-1985 Infer a function return type based on its type hint for locally defined functions (#1900)

Author: guillaume-dequenne

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

@@ -49,3 +49,9 @@ def foo():
     def bar():
         ...
 
+def local_function() -> int:
+    ...
+
+def calling_local_function():
+    x = local_function()
+    x() # Noncompliant

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

@@ -35,6 +35,7 @@
 import org.sonar.python.types.v2.FunctionType;
 import org.sonar.python.types.v2.ParameterV2;
 import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.TypeOrigin;
 
 import static org.sonar.python.tree.TreeUtils.locationInFile;
 
@@ -49,6 +50,7 @@ public class FunctionTypeBuilder implements TypeBuilder<FunctionType> {
   private boolean isInstanceMethod;
   private PythonType owner;
   private PythonType returnType = PythonType.UNKNOWN;
+  private TypeOrigin typeOrigin = TypeOrigin.STUB;
   private LocationInFile definitionLocation;
 
   private static final String CLASS_METHOD_DECORATOR = "classmethod";
@@ -110,6 +112,11 @@ public FunctionTypeBuilder withReturnType(PythonType returnType) {
     return this;
   }
 
+  public FunctionTypeBuilder withTypeOrigin(TypeOrigin typeOrigin) {
+    this.typeOrigin = typeOrigin;
+    return this;
+  }
+
   @Override
   public FunctionTypeBuilder withDefinitionLocation(@Nullable LocationInFile definitionLocation) {
     this.definitionLocation = definitionLocation;
@@ -118,7 +125,7 @@ public FunctionTypeBuilder withDefinitionLocation(@Nullable LocationInFile defin
 
   public FunctionType build() {
     return new FunctionType(
-      name, attributes, parameters, returnType, isAsynchronous, hasDecorators, isInstanceMethod, hasVariadicParameter, owner, definitionLocation
+      name, attributes, parameters, returnType, typeOrigin, isAsynchronous, hasDecorators, isInstanceMethod, hasVariadicParameter, owner, definitionLocation
     );
   }
 

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

@@ -47,6 +47,7 @@
 import org.sonar.python.types.v2.ModuleType;
 import org.sonar.python.types.v2.ParameterV2;
 import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.TypeOrigin;
 import org.sonar.python.types.v2.UnionType;
 
 public class SymbolsModuleTypeProvider {
@@ -122,12 +123,14 @@ private PythonType convertToFunctionType(FunctionSymbol symbol, Map<Symbol, Pyth
     InferredType inferredType = ((FunctionSymbolImpl) symbol).declaredReturnType();
     ClassSymbol classSymbol = inferredType.runtimeTypeSymbol();
     var returnType = resolveReturnType(createdTypesBySymbol, classSymbol);
+    TypeOrigin typeOrigin = symbol.isStub() ? TypeOrigin.STUB : TypeOrigin.LOCAL;
 
     FunctionTypeBuilder functionTypeBuilder =
       new FunctionTypeBuilder(symbol.name())
         .withAttributes(List.of())
         .withParameters(parameters)
         .withReturnType(returnType)
+        .withTypeOrigin(typeOrigin)
         .withAsynchronous(symbol.isAsynchronous())
         .withHasDecorators(symbol.hasDecorators())
         .withInstanceMethod(symbol.isInstanceMethod())

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

@@ -79,6 +79,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.TypeOrigin;
 import org.sonar.python.types.v2.TypeSource;
 import org.sonar.python.types.v2.UnionType;
 
@@ -257,13 +258,13 @@ public void visitFunctionDef(FunctionDef functionDef) {
     scan(functionDef.decorators());
     scan(functionDef.typeParams());
     scan(functionDef.parameters());
+    scan(functionDef.returnTypeAnnotation());
     FunctionType functionType = buildFunctionType(functionDef);
     ((NameImpl) functionDef.name()).typeV2(functionType);
     inTypeScope(functionType, () -> {
       // TODO: check scope accuracy
       scan(functionDef.typeParams());
       scan(functionDef.parameters());
-      scan(functionDef.returnTypeAnnotation());
       scan(functionDef.body());
     });
   }
@@ -279,6 +280,12 @@ private FunctionType buildFunctionType(FunctionDef functionDef) {
     if (owner != null) {
       functionTypeBuilder.withOwner(owner);
     }
+    TypeAnnotation typeAnnotation = functionDef.returnTypeAnnotation();
+    if (typeAnnotation != null) {
+      PythonType returnType = typeAnnotation.expression().typeV2();
+      functionTypeBuilder.withReturnType(returnType);
+      functionTypeBuilder.withTypeOrigin(TypeOrigin.LOCAL);
+    }
     FunctionType functionType = functionTypeBuilder.build();
     if (owner != null) {
       if (functionDef.name().symbolV2().hasSingleBindingUsage()) {

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

@@ -51,6 +51,7 @@
 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.TypeOrigin;
 import org.sonar.python.types.v2.TypeSource;
 import org.sonar.python.types.v2.UnionType;
 
@@ -189,11 +190,29 @@ public List<Expression> typeDependencies() {
 
   @Override
   public PythonType typeV2() {
-    TypeSource typeSource = computeTypeSource();
-    PythonType pythonType = returnTypeOfCall(callee().typeV2());
+    PythonType calleeType = callee().typeV2();
+    TypeSource typeSource = computeTypeSource(calleeType);
+    PythonType pythonType = returnTypeOfCall(calleeType);
     return pythonType != PythonType.UNKNOWN ? new ObjectType(pythonType, typeSource) : PythonType.UNKNOWN;
   }
 
+  private TypeSource computeTypeSource(PythonType calleeType) {
+    if (isCalleeLocallyDefinedFunction(calleeType)) {
+      return TypeSource.TYPE_HINT;
+    }
+    return calleeTypeSource();
+  }
+
+  boolean isCalleeLocallyDefinedFunction(PythonType pythonType) {
+    if (pythonType instanceof FunctionType functionType) {
+      return functionType.typeOrigin() == TypeOrigin.LOCAL;
+    }
+    if (pythonType instanceof UnionType unionType) {
+      return unionType.candidates().stream().anyMatch(this::isCalleeLocallyDefinedFunction);
+    }
+    return false;
+  }
+
   static PythonType returnTypeOfCall(PythonType calleeType) {
     if (calleeType instanceof ClassType classType) {
       return classType;
@@ -219,7 +238,7 @@ static PythonType returnTypeOfCall(PythonType calleeType) {
     return PythonType.UNKNOWN;
   }
 
-  TypeSource computeTypeSource() {
+  TypeSource calleeTypeSource() {
     if (callee() instanceof QualifiedExpression qualifiedExpression) {
       return qualifiedExpression.qualifier().typeV2().typeSource();
     }

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

@@ -35,6 +35,7 @@ public final class FunctionType implements PythonType {
   private final List<PythonType> attributes;
   private final List<ParameterV2> parameters;
   private PythonType returnType;
+  private final TypeOrigin typeOrigin;
   private final boolean isAsynchronous;
   private final boolean hasDecorators;
   private final boolean isInstanceMethod;
@@ -50,6 +51,7 @@ public FunctionType(
     List<PythonType> attributes,
     List<ParameterV2> parameters,
     PythonType returnType,
+    TypeOrigin typeOrigin,
     boolean isAsynchronous,
     boolean hasDecorators,
     boolean isInstanceMethod,
@@ -61,6 +63,7 @@ public FunctionType(
     this.attributes = attributes;
     this.parameters = parameters;
     this.returnType = returnType;
+    this.typeOrigin = typeOrigin;
     this.isAsynchronous = isAsynchronous;
     this.hasDecorators = hasDecorators;
     this.isInstanceMethod = isInstanceMethod;
@@ -128,4 +131,8 @@ public void resolveLazyReturnType(PythonType pythonType) {
     }
     this.returnType = pythonType;
   }
+
+  public TypeOrigin typeOrigin() {
+    return typeOrigin;
+  }
 }

python-frontend/src/main/java/org/sonar/python/types/v2/TypeOrigin.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 TypeOrigin {
+  LOCAL,
+  STUB
+}

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

@@ -60,13 +60,16 @@
 import org.sonar.python.types.v2.ObjectType;
 import org.sonar.python.types.v2.ParameterV2;
 import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.TypeOrigin;
 import org.sonar.python.types.v2.TypeSource;
 import org.sonar.python.types.v2.UnionType;
 import org.sonar.python.types.v2.UnknownType;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
 import static org.sonar.python.PythonTestUtils.parse;
+import static org.sonar.python.PythonTestUtils.parseWithoutSymbols;
+import static org.sonar.python.PythonTestUtils.pythonFile;
 import static org.sonar.python.types.v2.TypesTestUtils.DICT_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.FROZENSET_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.INT_TYPE;
@@ -423,7 +426,7 @@ void typeSourceIsExactByDefault() {
 
     CallExpression callExpressionSpy = Mockito.spy(callExpression);
     Expression calleeSpy = Mockito.spy(callExpression.callee());
-    FunctionType functionType = new FunctionType("foo", List.of(), List.of(), INT_TYPE, false, false, false, false, null, null);
+    FunctionType functionType = new FunctionType("foo", List.of(), List.of(), INT_TYPE, TypeOrigin.STUB, false, false, false, false, null, null);
     Mockito.when(calleeSpy.typeV2()).thenReturn(functionType);
     Mockito.when(callExpressionSpy.callee()).thenReturn(calleeSpy);
 
@@ -1855,6 +1858,56 @@ void return_type_of_call_expression_2() {
     ).typeV2().unwrappedType()).isEqualTo(NONE_TYPE);
   }
 
+  @Test
+  void return_type_of_call_of_locally_defined_function() {
+    var type = lastExpression("""
+      def foo() -> int: ...
+      x = foo()
+      x
+      """).typeV2();
+    assertThat(type.unwrappedType()).isEqualTo(INT_TYPE);
+    assertThat(type.typeSource()).isEqualTo(TypeSource.TYPE_HINT);
+  }
+
+  @Test
+  void type_origin_of_stub_function() {
+    FileInput fileInput = inferTypes("""
+      len
+      x = len([1,2])
+      x
+      """);
+
+    FunctionType lenType = (FunctionType) ((ExpressionStatement) fileInput.statements().statements().get(0)).expressions().get(0).typeV2();
+    assertThat(lenType.typeOrigin()).isEqualTo(TypeOrigin.STUB);
+    PythonType xType = ((ExpressionStatement) fileInput.statements().statements().get(2)).expressions().get(0).typeV2();
+    assertThat(xType.unwrappedType()).isEqualTo(INT_TYPE);
+    assertThat(xType.typeSource()).isEqualTo(TypeSource.EXACT);
+
+  }
+
+  @Test
+  void type_origin_of_project_function() {
+    FileInput tree = parseWithoutSymbols(
+      "def foo() -> int: ..."
+    );
+    ProjectLevelSymbolTable projectLevelSymbolTable = new ProjectLevelSymbolTable();
+    projectLevelSymbolTable.addModule(tree, "", pythonFile("mod.py"));
+    ProjectLevelTypeTable projectLevelTypeTable = new ProjectLevelTypeTable(projectLevelSymbolTable);
+
+    var lines = """
+      from mod import foo
+      foo
+      x = foo()
+      x
+      """;
+    FileInput fileInput = inferTypes(lines, projectLevelTypeTable);
+    FunctionType fooType = (FunctionType) ((ExpressionStatement) fileInput.statements().statements().get(1)).expressions().get(0).typeV2();
+    assertThat(fooType.typeOrigin()).isEqualTo(TypeOrigin.LOCAL);
+    PythonType xType = ((ExpressionStatement) fileInput.statements().statements().get(3)).expressions().get(0).typeV2();
+    // Declared return types of local functions are currently not stored in the project level symbol table
+    assertThat(xType).isEqualTo(PythonType.UNKNOWN);
+  }
+
   @Test
   void return_type_of_call_expression_union_type() {
     FileInput fileInput = inferTypes(

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

@@ -33,6 +33,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.sonar.python.PythonTestUtils.parseWithoutSymbols;
+import static org.sonar.python.types.v2.TypesTestUtils.INT_TYPE;
 import static org.sonar.python.types.v2.TypesTestUtils.PROJECT_LEVEL_TYPE_TABLE;
 
 class FunctionTypeTest {
@@ -148,6 +149,14 @@ void declaredTypes() {
     assertThat(functionType.parameters()).extracting(ParameterV2::declaredType).containsExactly(PythonType.UNKNOWN);
   }
 
+  @Test
+  void declared_return_type() {
+    FunctionType functionType = functionType("def fn() -> int: ...");
+    assertThat(functionType.returnType()).isEqualTo(INT_TYPE);
+    functionType = functionType("def fn() -> unknown: ...");
+    assertThat(functionType.returnType()).isEqualTo(PythonType.UNKNOWN);
+  }
+
   @Test
   void decorators() {
     // TODO: SONARPY-1772 Handle decorators