SONARPY-900 S5886 should not report on async functions returning AsyncGenerator

Author: guillaume-dequenne

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

@@ -44,6 +44,7 @@ public class FunctionReturnTypeCheck extends PythonSubscriptionCheck {
 
   private static final String MESSAGE = "Return a value of type \"%s\" instead of \"%s\" or update function \"%s\" type hint.";
   private static final List<String> ITERABLE_TYPES = Arrays.asList("typing.Generator", "typing.Iterator", "typing.Iterable");
+  private static final List<String> ASYNC_ITERABLE_TYPES = Arrays.asList("typing.AsyncGenerator", "typing.AsyncIterator", "typing.AsyncIterable");
 
   @Override
   public void initialize(Context context) {
@@ -68,14 +69,24 @@ private static void raiseIssues(SubscriptionContext ctx, FunctionDef functionDef
     String functionName = functionDef.name().name();
     String returnTypeName = InferredTypes.typeName(declaredReturnType);
     if (!returnTypeVisitor.yieldStatements.isEmpty()) {
+      boolean isAsyncFunction = functionDef.asyncKeyword() != null;
+      String recommendedSuperType = isAsyncFunction ? "typing.AsyncGenerator" : "typing.Generator";
       // Here we should probably use an equivalent of "canBeOrExtend" (accepting uncertainty) instead of "mustBeOrExtend"
-      if (ITERABLE_TYPES.stream().anyMatch(declaredReturnType::mustBeOrExtend)) {
+      if (ITERABLE_TYPES.stream().anyMatch(declaredReturnType::mustBeOrExtend) || ASYNC_ITERABLE_TYPES.stream().anyMatch(declaredReturnType::mustBeOrExtend)) {
+        if (isMixedUpAnnotation(isAsyncFunction, declaredReturnType)) {
+          returnTypeVisitor.yieldStatements
+            .forEach(y -> {
+              PreciseIssue issue =
+                ctx.addIssue(y, String.format("Annotate function \"%s\" with \"%s\" or one of its supertypes.", functionName, recommendedSuperType));
+              addSecondaries(issue, functionDef);
+            });
+        }
         return;
       }
       returnTypeVisitor.yieldStatements
         .forEach(y -> {
           PreciseIssue issue =
-            ctx.addIssue(y, String.format("Remove this yield statement or annotate function \"%s\" with \"typing.Generator\" or one of its supertypes.", functionName));
+            ctx.addIssue(y, String.format("Remove this yield statement or annotate function \"%s\" with \"%s\" or one of its supertypes.", functionName, recommendedSuperType));
           addSecondaries(issue, functionDef);
         });
     }
@@ -92,6 +103,10 @@ private static void raiseIssues(SubscriptionContext ctx, FunctionDef functionDef
     });
   }
 
+  private static boolean isMixedUpAnnotation(boolean isAsyncFunction, InferredType declaredReturnType) {
+    return isAsyncFunction ? ITERABLE_TYPES.stream().anyMatch(declaredReturnType::mustBeOrExtend) : ASYNC_ITERABLE_TYPES.stream().anyMatch(declaredReturnType::mustBeOrExtend);
+  }
+
   private static void addSecondaries(PreciseIssue issue, FunctionDef functionDef) {
     issue.secondary(functionDef.name(), "Function definition.");
     TypeAnnotation returnTypeAnnotation = functionDef.returnTypeAnnotation();

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

@@ -1,4 +1,4 @@
-from typing import List, SupportsFloat, Set, Dict, NoReturn, Text, Generator, Tuple, Union, AnyStr, Iterator, Iterable, Callable, Optional, TypedDict
+from typing import List, SupportsFloat, Set, Dict, NoReturn, Text, Generator, Tuple, Union, AnyStr, Iterator, Iterable, Callable, Optional, TypedDict, AsyncIterator, AsyncGenerator
 import numpy as np
 
 def builtins():
@@ -313,3 +313,12 @@ def my_dict() -> MyCustomDict:
     users = {1,2,3}
     messages = {1,2,3}
     return dict(user_ids=users, message_ids=messages)
+
+async def no_issue_for_async_iterators() -> AsyncIterator[int]:
+    yield 1
+
+async def no_issue_for_async_generators() -> AsyncGenerator[int]:
+    yield 1
+
+async def async_function_returning_iterator() -> Iterator[int]:
+    yield 1 # Noncompliant {{Annotate function "async_function_returning_iterator" with "typing.AsyncGenerator" or one of its supertypes.}}