SONARPY-1826 Enable flow sensitive type inference for function types

Author: guillaume-dequenne

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

@@ -218,3 +218,40 @@ def calling_global_func():
 def using_nonlocal_var():
     nonlocal some_nonlocal_var
     some_nonlocal_var()  # Noncompliant
+
+
+def reassigned_function():
+    if cond:
+        def my_callable(): ...
+        my_callable()  # OK
+    else:
+        def my_callable(): ...
+        my_callable = 42
+        my_callable()  # Noncompliant
+
+
+
+def recursive_with_try_finally(x):
+    if x is False:
+        print("recursion!")
+        return
+    recursive_with_try_finally(False) # Noncompliant
+    try:
+        recursive_with_try_finally(False) # Noncompliant
+    finally:
+        recursive_with_try_finally = None
+        recursive_with_try_finally(False) # Noncompliant
+    recursive_with_try_finally(False)  # Noncompliant
+
+
+
+def nested_recursive_try_finally():
+    def my_rec(x):
+        if x is False:
+            print("yeah")
+            return
+        my_rec(False)
+    try:
+        my_rec(True)
+    finally:
+        my_rec = None

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

@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+public class SymbolV2Utils {
+
+  private SymbolV2Utils() {}
+
+  public static boolean isFunctionOrClassDeclaration(UsageV2 usageV2) {
+    return usageV2.kind() == UsageV2.Kind.FUNC_DECLARATION || usageV2.kind() == UsageV2.Kind.CLASS_DECLARATION;
+  }
+}

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

@@ -35,8 +35,8 @@
 import org.sonar.plugins.python.api.tree.Parameter;
 import org.sonar.plugins.python.api.tree.StatementList;
 import org.sonar.plugins.python.api.tree.Tree;
-import org.sonar.python.semantic.v2.types.Assignment;
 import org.sonar.python.semantic.v2.types.FlowSensitiveTypeInference;
+import org.sonar.python.semantic.v2.types.Propagation;
 import org.sonar.python.semantic.v2.types.PropagationVisitor;
 import org.sonar.python.semantic.v2.types.TrivialTypeInferenceVisitor;
 import org.sonar.python.semantic.v2.types.TryStatementVisitor;
@@ -111,9 +111,9 @@ private static void inferTypesAndMemberAccessSymbols(Tree scopeTree,
   ) {
     PropagationVisitor propagationVisitor = new PropagationVisitor();
     scopeTree.accept(propagationVisitor);
-    Set<Name> assignedNames = propagationVisitor.assignmentsByLhs().values().stream()
+    Set<Name> assignedNames = propagationVisitor.propagationsByLhs().values().stream()
       .flatMap(Collection::stream)
-      .map(Assignment::lhsName)
+      .map(Propagation::lhsName)
       .collect(Collectors.toSet());
 
     TryStatementVisitor tryStatementVisitor = new TryStatementVisitor();
@@ -135,6 +135,7 @@ private static void flowSensitiveTypeInference(ControlFlowGraph cfg, Set<SymbolV
     FlowSensitiveTypeInference flowSensitiveTypeInference = new FlowSensitiveTypeInference(
       trackedVars,
       propagationVisitor.assignmentsByAssignmentStatement(),
+      propagationVisitor.definitionsByDefinitionStatement(),
       Map.of());
 
     flowSensitiveTypeInference.compute(cfg);

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

@@ -21,32 +21,23 @@
 
 import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import org.sonar.plugins.python.api.tree.Expression;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.python.semantic.v2.SymbolV2;
-import org.sonar.python.semantic.v2.UsageV2;
-import org.sonar.python.tree.NameImpl;
 import org.sonar.python.types.HasTypeDependencies;
 import org.sonar.python.types.v2.PythonType;
-import org.sonar.python.types.v2.UnionType;
 
-public class Assignment {
+public class Assignment extends Propagation {
 
-  final SymbolV2 lhsSymbol;
-  Name lhsName;
   Expression rhs;
-  Set<SymbolV2> variableDependencies = new HashSet<>();
-  Set<Assignment> dependents = new HashSet<>();
-  Map<SymbolV2, Set<Assignment>> assignmentsByLhs;
+  Map<SymbolV2, Set<Propagation>> propagationsByLhs;
 
-  public Assignment(SymbolV2 lhsSymbol, Name lhsName, Expression rhs, Map<SymbolV2, Set<Assignment>> assignmentsByLhs) {
-    this.lhsSymbol = lhsSymbol;
-    this.lhsName = lhsName;
+  public Assignment(SymbolV2 lhsSymbol, Name lhsName, Expression rhs, Map<SymbolV2, Set<Propagation>> propagationsByLhs) {
+    super(lhsSymbol, lhsName);
     this.rhs = rhs;
-    this.assignmentsByLhs = assignmentsByLhs;
+    this.propagationsByLhs = propagationsByLhs;
   }
 
   void computeDependencies(Expression expression, Set<SymbolV2> trackedVars) {
@@ -58,45 +49,28 @@ void computeDependencies(Expression expression, Set<SymbolV2> trackedVars) {
         SymbolV2 symbol = name.symbolV2();
         if (symbol != null && trackedVars.contains(symbol)) {
           variableDependencies.add(symbol);
-          assignmentsByLhs.get(symbol).forEach(a -> a.dependents.add(this));
+          propagationsByLhs.get(symbol).forEach(a -> a.dependents.add(this));
         }
       } else if (e instanceof HasTypeDependencies hasTypeDependencies) {
         workList.addAll(hasTypeDependencies.typeDependencies());
       }
     }
   }
 
-  boolean areDependenciesReady(Set<SymbolV2> initializedVars) {
-    return initializedVars.containsAll(variableDependencies);
-  }
-
-  /** @return true if the propagation effectively changed the inferred type of lhs */
-  public boolean propagate(Set<SymbolV2> initializedVars) {
-    PythonType rhsType = rhs.typeV2();
-    if (initializedVars.add(lhsSymbol)) {
-      lhsSymbol.usages().stream().map(UsageV2::tree).filter(NameImpl.class::isInstance).map(NameImpl.class::cast).forEach(n -> n.typeV2(rhsType));
-      return true;
-    } else {
-      PythonType currentType = lhsName.typeV2();
-      PythonType newType = UnionType.or(rhsType, currentType);
-      lhsSymbol.usages().stream().map(UsageV2::tree).filter(NameImpl.class::isInstance).map(NameImpl.class::cast).forEach(n -> n.typeV2(newType));
-      return !newType.equals(currentType);
-    }
-  }
-
   public Name lhsName() {
     return lhsName;
   }
 
+  @Override
+  public PythonType rhsType() {
+    return rhs.typeV2();
+  }
+
   public SymbolV2 lhsSymbol() {
     return lhsSymbol;
   }
 
   public Expression rhs() {
     return rhs;
   }
-
-  public Set<Assignment> dependents() {
-    return dependents;
-  }
 }

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

@@ -0,0 +1,48 @@
+/*
+ * 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.types;
+
+import org.sonar.plugins.python.api.tree.Name;
+import org.sonar.python.semantic.v2.SymbolV2;
+import org.sonar.python.types.v2.PythonType;
+
+/**
+ * This represents a class or function definition
+ * It can be modelled as an assignment of a class / function type to their symbol
+ */
+public class Definition extends Propagation {
+
+  public Definition(SymbolV2 symbol, Name name) {
+    super(symbol, name);
+  }
+
+  public SymbolV2 lhsSymbol() {
+    return lhsSymbol;
+  }
+
+  public Name lhsName() {
+    return lhsName;
+  }
+
+  @Override
+  public PythonType rhsType() {
+    return lhsName.typeV2();
+  }
+}

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

@@ -42,12 +42,18 @@
 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<String, PythonType> parameterTypesByName;
 
-  public FlowSensitiveTypeInference(Set<SymbolV2> trackedVars,
-                                    Map<Statement, Assignment> assignmentsByAssignmentStatement, Map<String, PythonType> parameterTypesByName) {
+  public FlowSensitiveTypeInference(
+    Set<SymbolV2> trackedVars,
+    Map<Statement, Assignment> assignmentsByAssignmentStatement,
+    Map<Statement, Definition> definitionsByDefinitionStatement,
+    Map<String, PythonType> parameterTypesByName
+  ) {
     this.trackedVars = trackedVars;
     this.assignmentsByAssignmentStatement = assignmentsByAssignmentStatement;
+    this.definitionsByDefinitionStatement = definitionsByDefinitionStatement;
     this.parameterTypesByName = parameterTypesByName;
   }
 
@@ -84,6 +90,8 @@ public void updateProgramState(Tree element, ProgramState programState) {
         // update lhs
         updateTree(assignment.variable(), state);
       }
+    } else if (element instanceof FunctionDef functionDef) {
+      handleDefinition(functionDef, state);
     } else {
       // TODO: isinstance visitor
       updateTree(element, state);
@@ -133,4 +141,14 @@ private void handleAssignment(Statement assignmentStatement, TypeInferenceProgra
         }
       });
   }
+
+  private void handleDefinition(Statement definitionStatement, TypeInferenceProgramState programState) {
+    Optional.ofNullable(definitionsByDefinitionStatement.get(definitionStatement))
+      .ifPresent(definition -> {
+        SymbolV2 symbol = definition.lhsSymbol();
+        if (trackedVars.contains(symbol)) {
+          programState.setTypes(symbol, Set.of(definition.lhsName.typeV2()));
+        }
+      });
+  }
 }

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

@@ -0,0 +1,108 @@
+/*
+ * 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.types;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.stream.Stream;
+import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.Name;
+import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.python.semantic.v2.SymbolV2;
+import org.sonar.python.semantic.v2.SymbolV2Utils;
+import org.sonar.python.semantic.v2.UsageV2;
+import org.sonar.python.tree.NameImpl;
+import org.sonar.python.types.v2.PythonType;
+import org.sonar.python.types.v2.UnionType;
+
+public abstract class Propagation {
+
+  final Set<SymbolV2> variableDependencies;
+  final Set<Propagation> dependents;
+
+  final SymbolV2 lhsSymbol;
+  final Name lhsName;
+
+  protected Propagation(SymbolV2 lhsSymbol, Name lhsName) {
+    this.lhsSymbol = lhsSymbol;
+    this.lhsName = lhsName;
+    this.variableDependencies = new HashSet<>();
+    this.dependents = new HashSet<>();
+  }
+
+  /**
+   * This is used for AST-based type inference in try/catch statements
+   * @return true if the propagation effectively changed the inferred type of assignment LHS
+   */
+  boolean propagate(Set<SymbolV2> initializedVars) {
+    PythonType rhsType = rhsType();
+    if (initializedVars.add(lhsSymbol)) {
+      getSymbolNonDeclarationUsageTrees(lhsSymbol)
+        .filter(NameImpl.class::isInstance)
+        .map(NameImpl.class::cast)
+        .forEach(n -> n.typeV2(rhsType));
+      return true;
+    } else {
+      PythonType currentType = currentType(lhsName);
+      if (currentType == null) {
+        return false;
+      }
+      PythonType newType = UnionType.or(rhsType, currentType);
+      getSymbolNonDeclarationUsageTrees(lhsSymbol)
+        .filter(NameImpl.class::isInstance)
+        .map(NameImpl.class::cast)
+        .forEach(n -> n.typeV2(newType));
+      return !newType.equals(currentType);
+    }
+  }
+
+  public static Stream<Tree> getSymbolNonDeclarationUsageTrees(SymbolV2 symbol) {
+    return symbol.usages()
+      .stream()
+      // Function and class definition names will always have FunctionType and ClassType respectively
+      // so they are filtered out of type propagation
+      .filter(u -> !SymbolV2Utils.isFunctionOrClassDeclaration(u))
+      .map(UsageV2::tree);
+  }
+
+  boolean areDependenciesReady(Set<SymbolV2> initializedVars) {
+    return initializedVars.containsAll(variableDependencies);
+  }
+
+  Set<Propagation> dependents() {
+    return dependents;
+  }
+
+  public abstract Name lhsName();
+
+  public abstract PythonType rhsType();
+
+  static PythonType currentType(Name lhsName) {
+    return lhsName.symbolV2().usages()
+      .stream()
+      .filter(u -> !SymbolV2Utils.isFunctionOrClassDeclaration(u))
+      .map(UsageV2::tree)
+      .filter(Expression.class::isInstance)
+      .map(Expression.class::cast)
+      .findFirst()
+      .map(Expression::typeV2)
+      .orElse(null);
+  }
+}