SONARPY-2150: Fix FP on S6982 when method is used on optimizers (#1984)

* SONARPY-2150 check type of module

* SONARPY-2150 improve coverage

Author: Seppli11

python-checks/src/main/java/org/sonar/python/checks/TorchModuleModeShouldBeSetAfterLoadingCheck.java

@@ -27,23 +27,18 @@
 import org.sonar.plugins.python.api.PythonSubscriptionCheck;
 import org.sonar.plugins.python.api.symbols.Symbol;
 import org.sonar.plugins.python.api.symbols.Usage;
-import org.sonar.plugins.python.api.tree.Argument;
 import org.sonar.plugins.python.api.tree.CallExpression;
-import org.sonar.plugins.python.api.tree.Expression;
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.QualifiedExpression;
-import org.sonar.plugins.python.api.tree.RegularArgument;
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.python.cfg.fixpoint.ReachingDefinitionsAnalysis;
 import org.sonar.python.tree.TreeUtils;
 
 @Rule(key = "S6982")
 public class TorchModuleModeShouldBeSetAfterLoadingCheck extends PythonSubscriptionCheck {
   private static final Set<String> STATE_SETTING_FUNCTION_FQNS = Set.of("eval", "train");
-  private static final String TORCH_LOAD_FQN = "torch.load";
   private static final String LOAD_STATE_DICT_NAME = "load_state_dict";
   private static final String MESSAGE = "Set the module in training or evaluation mode.";
-  private static final int IS_TORCH_LOAD_CALL_MAX_RECURSIVE_COUNTER = 10;
 
   private ReachingDefinitionsAnalysis reachingDefinitionsAnalysis;
 
@@ -54,7 +49,7 @@ public void initialize(Context context) {
 
     context.registerSyntaxNodeConsumer(Tree.Kind.CALL_EXPR, ctx -> {
       CallExpression callExpr = (CallExpression) ctx.syntaxNode();
-      List<Usage> receiverUsages = getForwardUsages(callExpr);
+      List<Usage> receiverUsages = getForwardUsagesOfReceiver(callExpr);
       if (isLoadStateDictCall(callExpr) && !hasEvalOrTrainUsage(receiverUsages) && !isModelPassedOn(receiverUsages)) {
         ctx.addIssue(callExpr.callee(), MESSAGE);
       }
@@ -65,33 +60,14 @@ private boolean isLoadStateDictCall(CallExpression callExpr) {
     // To properly check if the correct load_state_dict is called, typeshed type information would be required.
     // Since this is currently not possible, we check if the parameter to load_state_dict is torch.load(...),
     // with the assumption that if torch.load is passed to this load_state_dict, it is probably the correct method
-    if(callExpr.callee() instanceof QualifiedExpression qualifiedExpr) {
-      return LOAD_STATE_DICT_NAME.equals(qualifiedExpr.name().name()) && containsTorchLoadCall(callExpr.arguments());
+    if (callExpr.callee() instanceof QualifiedExpression qualifiedExpr) {
+      return qualifiedExpr.qualifier().type().mustBeOrExtend("torch.nn.modules.module.Module")
+        && LOAD_STATE_DICT_NAME.equals(qualifiedExpr.name().name());
     }
     return false;
   }
 
-  private boolean containsTorchLoadCall(List<Argument> args) {
-    return args.stream()
-      .flatMap(TreeUtils.toStreamInstanceOfMapper(RegularArgument.class))
-      .anyMatch(arg -> isTorchLoadCall(arg.expression(), 0));
-  }
-
-  private boolean isTorchLoadCall(Expression expr, int recursiveCounter) {
-    if (recursiveCounter > IS_TORCH_LOAD_CALL_MAX_RECURSIVE_COUNTER) {
-      return false;
-    } else if (expr instanceof CallExpression callExpr) {
-      Symbol calleeSymbol = callExpr.calleeSymbol();
-      return calleeSymbol != null && TORCH_LOAD_FQN.equals(calleeSymbol.fullyQualifiedName());
-    } else if (expr instanceof Name name) {
-      return reachingDefinitionsAnalysis.valuesAtLocation(name).stream()
-        .anyMatch(definitionExpr -> isTorchLoadCall(definitionExpr, recursiveCounter + 1));
-    } else {
-      return false;
-    }
-  }
-
-  private static List<Usage> getForwardUsages(CallExpression callExpr) {
+  private static List<Usage> getForwardUsagesOfReceiver(CallExpression callExpr) {
     List<Usage> usages = getFunctionCallReceiverName(callExpr)
       .flatMap(name -> Optional.ofNullable(name.symbol()))
       .map(Symbol::usages)

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

@@ -1,4 +1,4 @@
-import torch
+import torch.nn
 import torchvision.models as models
 
 def noncompliant():
@@ -7,47 +7,39 @@ def noncompliant():
    #^^^^^^^^^^^^^^^^^^^^^
     ...
 
-def noncompliant(model):
-    model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
+class CustomModule(torch.nn.Module):
+    pass
 
 def noncompliant():
+    model = CustomModule()
     model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
+    ...
 
 def noncompliant():
-    get_model().load_state_dict(torch.load('model_weights.pth')) # Noncompliant
-
-def noncompliant(model):
-    weights = torch.load('model_weights.pth')
-    weights2 = weights
-    model.load_state_dict(weights2) # Noncompliant
+    model = models.vgg16()
+    model.train()
+    model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
 
 def noncompliant():
     model = models.vgg16()
+    model.load_state_dict(torch.load('model_weights.pth'))
     model.train()
+
     model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
+
     other_model = model
 
 def compliant(model):
-    weights = weights
-    model.load_state_dict(weights)
+    model.load_state_dict(torch.load('model_weights.pth'))
 
-def compliant():
-    model1 = models.vgg16()
-    model1.load_state_dict(torch.load('model_weights.pth'))
-    model1.eval()
 
-def compliant():
-    model2 = models.vgg16()
-    model2.load_state_dict(torch.load('model_weights.pth'))
-    other_model = model2
-    model2.train()
 
 def compliant():
     model3 = models.vgg16()
     model3.load_state_dict(torch.load('model_weights.pth')) # Ok if model is passed as argument to a function do not raise at all train or eval could be called in such functions
     foo(model3)
 
-def compliant():
-    # Ok since no torch.load() result is passed as an argument
-    model.load_state_dict(1 + 1)
-    model.load_state_dict((lambda x: x)())
+
+
+
+