SONARPY-2151 S6982: Fix fn when import torch is used instead of the full import name torch.nn

Author: Seppli11

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

@@ -31,6 +31,7 @@
 import org.sonar.plugins.python.api.tree.Name;
 import org.sonar.plugins.python.api.tree.QualifiedExpression;
 import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.plugins.python.api.types.InferredType;
 import org.sonar.python.tree.TreeUtils;
 
 @Rule(key = "S6982")
@@ -56,8 +57,10 @@ private static boolean isLoadStateDictCall(CallExpression callExpr) {
     // 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 qualifiedExpr.qualifier().type().mustBeOrExtend("torch.nn.modules.module.Module")
-        && LOAD_STATE_DICT_NAME.equals(qualifiedExpr.name().name());
+      InferredType qualifierType = qualifiedExpr.qualifier().type();
+      boolean isModule = qualifierType.mustBeOrExtend("torch.nn.modules.module.Module")
+        || qualifierType.mustBeOrExtend("torch.nn.Module");
+      return isModule && LOAD_STATE_DICT_NAME.equals(qualifiedExpr.name().name());
     }
     return false;
   }

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

@@ -1,45 +1,54 @@
-import torch.nn
-import torchvision.models as models
 
-def noncompliant():
-    model = models.vgg16()
-    model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant {{Set the module in training or evaluation mode.}}
-   #^^^^^^^^^^^^^^^^^^^^^
-    ...
+import torchvision.models as models
 
-class CustomModule(torch.nn.Module):
-    pass
+def torch_nn_imported():
+    import torch.nn
+    def noncompliant():
+        model = models.vgg16()
+        model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant {{Set the module in training or evaluation mode.}}
+       #^^^^^^^^^^^^^^^^^^^^^
+        ...
 
-def noncompliant():
-    model = CustomModule()
-    model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
-    ...
+    class CustomModule(torch.nn.Module):
+        pass
 
-def noncompliant():
-    model = models.vgg16()
-    model.train()
-    model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
+    def noncompliant():
+        model = CustomModule()
+        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()
+    def noncompliant():
+        model = models.vgg16()
+        model.train()
+        model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
 
-    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()
 
-    other_model = model
+        model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
 
-def compliant(model):
-    model.load_state_dict(torch.load('model_weights.pth'))
+        other_model = model
 
+    def compliant(model):
+        model.load_state_dict(torch.load('model_weights.pth'))
 
 
-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():
+        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 torch_imported():
+    import torch
 
+    class CustomModule(torch.nn.Module):
+        pass
 
+    def noncompliant():
+        model = CustomModule()
+        model.load_state_dict(torch.load('model_weights.pth')) # Noncompliant
+        ...