SONARPY-1801 infer type of subscription expression when it is a class…

… with generic type
SonarSource · Apr 26, 2024 · 6bc838b · 6bc838b
1 parent 71dac54
commit 6bc838b
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/python-frontend/src/main/java/org/sonar/python/tree/SubscriptionExpressionImpl.java b/python-frontend/src/main/java/org/sonar/python/tree/SubscriptionExpressionImpl.java
@@ -28,6 +28,8 @@
 import org.sonar.plugins.python.api.tree.Token;
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.plugins.python.api.tree.TreeVisitor;
+import org.sonar.python.types.v2.ClassType;
+import org.sonar.python.types.v2.PythonType;
 
 public class SubscriptionExpressionImpl extends PyTree implements SubscriptionExpression {
 
@@ -77,4 +79,13 @@ public List<Tree> computeChildren() {
   public Kind getKind() {
     return Kind.SUBSCRIPTION;
   }
+
+  @Override
+  public PythonType typeV2() {
+    var objectType = object().typeV2();
+    if (objectType instanceof ClassType classType) {
+      return classType;
+    }
+    return PythonType.UNKNOWN;
+  }
 }
diff --git a/python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java b/python-frontend/src/test/java/org/sonar/python/semantic/v2/TypeInferenceV2Test.java
@@ -45,6 +45,7 @@
 import org.sonar.python.types.v2.ClassType;
 import org.sonar.python.types.v2.FunctionType;
 import org.sonar.python.types.v2.ModuleType;
+import org.sonar.python.types.v2.ObjectType;
 import org.sonar.python.types.v2.PythonType;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -264,6 +265,35 @@ def foo():
     Assertions.assertThat(expressionStatement.expressions().get(0).typeV2()).isEqualTo(PythonType.UNKNOWN);
   }
 
+  @Test
+  void inferForSubscriptionExpressionFromClassType() {
+    FileInput root = inferTypes("""
+      def foo():
+        a = list[int](1, 2, 3)
+      """);
+
+    var functionDef = (FunctionDef) root.statements().statements().get(0);
+    var assignmentStatement = (AssignmentStatement) functionDef.body().statements().get(0);
+    var type = (ObjectType) assignmentStatement.assignedValue().typeV2();
+    assertThat(type.type()).isInstanceOf(ClassType.class)
+      .extracting(PythonType::name)
+      .isEqualTo("list");
+  }
+
+  @Test
+  void inferTypeForSubscriptionExpressionFromObjectType() {
+    FileInput root = inferTypes("""
+      def foo():
+        my_list = ["hello"]
+        a = my_list[0]
+      """);
+
+    var functionDef = (FunctionDef) root.statements().statements().get(0);
+    var assignmentStatement = (AssignmentStatement) functionDef.body().statements().get(1);
+    var type = assignmentStatement.assignedValue().typeV2();
+    assertThat(type).isEqualTo(PythonType.UNKNOWN);
+  }
+
   private FileInput inferTypes(String lines) {
     return inferTypes(lines, new HashMap<>());
   }