From 5390bf619dd75b26fb6133d5618f4de766e980a9 Mon Sep 17 00:00:00 2001
From: Ghislain Piot <ghislain.piot@sonarsource.com>
Date: Thu, 5 Dec 2024 14:09:07 +0100
Subject: [PATCH] SONARPY-2370 Fix FP on S5644 for generic types defined
 through type parameter syntax (#2209)

---
 .../checks/ItemOperationsTypeCheck.java       | 61 ++++++++++++++-----
 .../itemOperations_getitem.py                 |  5 ++
 2 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/python-checks/src/main/java/org/sonar/python/checks/ItemOperationsTypeCheck.java b/python-checks/src/main/java/org/sonar/python/checks/ItemOperationsTypeCheck.java
index 92816c78df..bbb82dc69f 100644
--- a/python-checks/src/main/java/org/sonar/python/checks/ItemOperationsTypeCheck.java
+++ b/python-checks/src/main/java/org/sonar/python/checks/ItemOperationsTypeCheck.java
@@ -16,18 +16,24 @@
  */
 package org.sonar.python.checks;
 
+import java.util.Collection;
 import java.util.Map;
+import java.util.Optional;
 import javax.annotation.Nullable;
 import org.sonar.check.Rule;
 import org.sonar.plugins.python.api.LocationInFile;
 import org.sonar.plugins.python.api.symbols.ClassSymbol;
 import org.sonar.plugins.python.api.symbols.FunctionSymbol;
 import org.sonar.plugins.python.api.symbols.Symbol;
+import org.sonar.plugins.python.api.tree.ArgList;
 import org.sonar.plugins.python.api.tree.CallExpression;
+import org.sonar.plugins.python.api.tree.ClassDef;
 import org.sonar.plugins.python.api.tree.Expression;
-import org.sonar.plugins.python.api.tree.HasSymbol;
+import org.sonar.plugins.python.api.tree.RegularArgument;
+import org.sonar.plugins.python.api.tree.SubscriptionExpression;
 import org.sonar.plugins.python.api.tree.Tree;
 import org.sonar.plugins.python.api.types.InferredType;
+import org.sonar.python.tree.TreeUtils;
 import org.sonar.python.types.InferredTypes;
 
 import static org.sonar.plugins.python.api.symbols.Symbol.Kind.CLASS;
@@ -39,30 +45,27 @@ public class ItemOperationsTypeCheck extends ItemOperationsType {
 
   @Override
   public boolean isValidSubscription(Expression subscriptionObject, String requiredMethod, @Nullable String classRequiredMethod,
-                                     Map<LocationInFile, String> secondaries) {
+    Map<LocationInFile, String> secondaries) {
 
     if (subscriptionObject.is(Tree.Kind.GENERATOR_EXPR)) {
       return false;
     }
-    if (subscriptionObject.is(Tree.Kind.CALL_EXPR)) {
-      Symbol subscriptionCalleeSymbol = ((CallExpression) subscriptionObject).calleeSymbol();
-      if (subscriptionCalleeSymbol != null && subscriptionCalleeSymbol.is(FUNCTION) && ((FunctionSymbol) subscriptionCalleeSymbol).isAsynchronous()) {
-        FunctionSymbol functionSymbol = (FunctionSymbol) subscriptionCalleeSymbol;
-        secondaries.put(functionSymbol.definitionLocation(), String.format(SECONDARY_MESSAGE, functionSymbol.name()));
-        return false;
-      }
+    if (isInvalidSubscriptionCallExpr(subscriptionObject, secondaries)) {
+      return false;
     }
-    if (subscriptionObject instanceof HasSymbol hasSymbol) {
-      Symbol symbol = hasSymbol.symbol();
-      if (symbol == null || isTypingOrCollectionsSymbol(symbol)) {
+
+    var symbolOptional = TreeUtils.getSymbolFromTree(subscriptionObject);
+
+    if (symbolOptional.isPresent()) {
+      var symbol = symbolOptional.get();
+      if (isTypingOrCollectionsSymbol(symbol)) {
         return true;
       }
       if (symbol.is(FUNCTION, CLASS)) {
-        secondaries.put(symbol.is(FUNCTION) ?
-          ((FunctionSymbol) symbol).definitionLocation() : ((ClassSymbol) symbol).definitionLocation(), String.format(SECONDARY_MESSAGE, symbol.name()));
-        return canHaveMethod(symbol, requiredMethod, classRequiredMethod);
+        return isValidSubscriptionSymbol(symbol, subscriptionObject, secondaries, requiredMethod, classRequiredMethod);
       }
     }
+
     InferredType type = subscriptionObject.type();
     String typeName = InferredTypes.typeName(type);
     String secondaryMessage = typeName != null ? String.format(SECONDARY_MESSAGE, typeName) : DEFAULT_SECONDARY_MESSAGE;
@@ -70,6 +73,34 @@ public boolean isValidSubscription(Expression subscriptionObject, String require
     return type.canHaveMember(requiredMethod);
   }
 
+  private static boolean isValidSubscriptionSymbol(Symbol symbol, Expression subscriptionObject, Map<LocationInFile, String> secondaries, String requiredMethod,
+    @Nullable String classRequiredMethod) {
+    LocationInFile locationInFile = symbol.is(FUNCTION) ? ((FunctionSymbol) symbol).definitionLocation() : ((ClassSymbol) symbol).definitionLocation();
+    secondaries.put(locationInFile, SECONDARY_MESSAGE.formatted(symbol.name()));
+    return isSubscriptionInClassArg(subscriptionObject) || canHaveMethod(symbol, requiredMethod, classRequiredMethod);
+  }
+
+  private static boolean isInvalidSubscriptionCallExpr(Expression expression, Map<LocationInFile, String> secondaries) {
+    if (expression instanceof CallExpression callExpression
+      && callExpression.calleeSymbol() instanceof FunctionSymbol functionSymbol
+      && functionSymbol.isAsynchronous()) {
+      secondaries.put(functionSymbol.definitionLocation(), SECONDARY_MESSAGE.formatted(functionSymbol.name()));
+      return true;
+    }
+    return false;
+  }
+
+  private static boolean isSubscriptionInClassArg(Expression subscriptionObject) {
+    return Optional.ofNullable(((ClassDef) TreeUtils.firstAncestorOfKind(subscriptionObject, Tree.Kind.CLASSDEF))).map(ClassDef::args).map(ArgList::arguments)
+      .stream()
+      .flatMap(Collection::stream)
+      .flatMap(TreeUtils.toStreamInstanceOfMapper(RegularArgument.class))
+      .map(RegularArgument::expression)
+      .flatMap(TreeUtils.toStreamInstanceOfMapper(SubscriptionExpression.class))
+      .map(SubscriptionExpression::object)
+      .anyMatch(subscriptionObject::equals);
+  }
+
   @Override
   public String message(@Nullable String name, String missingMethod) {
     if (name != null) {
diff --git a/python-checks/src/test/resources/checks/itemOperationsTypeCheck/itemOperations_getitem.py b/python-checks/src/test/resources/checks/itemOperationsTypeCheck/itemOperations_getitem.py
index 304e3ca4f7..afa8f04d4e 100644
--- a/python-checks/src/test/resources/checks/itemOperationsTypeCheck/itemOperations_getitem.py
+++ b/python-checks/src/test/resources/checks/itemOperationsTypeCheck/itemOperations_getitem.py
@@ -216,3 +216,8 @@ class ExtendedMock(Mock):
 
     def custom_mock():
         a = ExtendedMock()[42]
+
+
+class MyGenericClass[T]: ...
+
+class MyGenericSubType(MyGenericClass[str]): ...