Fix bugs around async functions (#22)

* Fix invalid display of coroutine which return another coroutine

* Fix invalid display of AsyncGenerator

Author: divaltor

src/main/kotlin/space/whitememory/pythoninlayparams/types/hints/HintGenerator.kt

@@ -48,24 +48,6 @@ enum class HintGenerator {
         }
     },
 
-    ASYNC_TYPE {
-        override fun handleType(
-            element: PyElement,
-            type: PyType?,
-            typeEvalContext: TypeEvalContext
-        ): List<InlayInfoDetails>? {
-            if (type == null || element !is PyFunction) return null
-
-            if (type is PyCollectionType && type.classQName == PyTypingTypeProvider.COROUTINE && element.isAsync) {
-                return generateTypeHintText(
-                    element, PyTypingTypeProvider.coroutineOrGeneratorElementType(type)?.get(), typeEvalContext
-                )
-            }
-
-            return null
-        }
-    },
-
     COLLECTION_TYPE {
         override fun handleType(
             element: PyElement,

src/main/kotlin/space/whitememory/pythoninlayparams/types/hints/HintResolver.kt

@@ -63,6 +63,8 @@ enum class HintResolver {
 
             // Handle case `var = async_func()` without `await` keyword
             if (assignedValue is PyCallExpression) return true
+            // Handle case 'var = await async_func()` which return `Coroutine` inside
+            if (assignedValue is PyPrefixExpression && assignedValue.operator == PyTokenTypes.AWAIT_KEYWORD) return true
 
             if (typeAnnotation is PyClassType && isElementInsideTypingModule(typeAnnotation.pyClass)) return false
 
@@ -398,7 +400,13 @@ enum class HintResolver {
         }
 
         fun getExpressionAnnotationType(element: PyElement, typeEvalContext: TypeEvalContext): PyType? {
-            if (element is PyFunction) return typeEvalContext.getReturnType(element)
+            if (element is PyFunction) {
+                if (element.isAsync && !element.isGenerator) {
+                    return element.getReturnStatementType(typeEvalContext)
+                }
+
+                return typeEvalContext.getReturnType(element)
+            }
             if (element is PyTargetExpression) return typeEvalContext.getType(element)
 
             return null
@@ -428,14 +436,6 @@ enum class HintResolver {
                 }
             }
 
-            if (element is PyFunction && element.isAsync) {
-                val functionType = PyTypingTypeProvider.coroutineOrGeneratorElementType(typeAnnotation)?.get()
-
-                if (functionType is PyNoneType || PyTypeChecker.isUnknown(functionType, false, typeEvalContext)) {
-                    return false
-                }
-            }
-
             if (PyTypeChecker.isUnknown(typeAnnotation, false, typeEvalContext)) return false
 
             return true