Skip to content

Commit af23359

Browse files
divaltordivaltor
authored
Add support for async functions (#15)
1 parent 37dba8e commit af23359

3 files changed

Lines changed: 42 additions & 24 deletions

File tree

‎src/main/kotlin/space/whitememory/pythoninlayparams/types/AbstractPythonInlayTypeHintsCollector.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor, open val se
2222

2323
protected fun renderTypeHint(element: PyElement, typeEvalContext: TypeEvalContext, sink: InlayHintsSink) {
2424
val typeAnnotation = HintResolver.getExpressionAnnotationType(element, typeEvalContext)
25-
val hintName = HintGenerator.generateTypeHintText(typeAnnotation, typeEvalContext)
25+
val hintName = HintGenerator.generateTypeHintText(element, typeAnnotation, typeEvalContext)
2626

2727
displayTypeHint(element, sink, hintName)
2828
}

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

Lines changed: 41 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,22 @@
11
package space.whitememory.pythoninlayparams.types.hints
22

33
import com.jetbrains.python.PyNames
4+
import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
5+
import com.jetbrains.python.psi.PyElement
6+
import com.jetbrains.python.psi.PyFunction
47
import com.jetbrains.python.psi.PyLambdaExpression
58
import com.jetbrains.python.psi.types.*
69

710
enum class HintGenerator {
8-
UNION_TYPE {
9-
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
11+
UNION_TYPE() {
12+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
1013
if (type !is PyUnionType) {
1114
return null
1215
}
1316

1417
val generatedValues = type.members
1518
.filterNotNull()
16-
.map { generateTypeHintText(it, typeEvalContext) }
19+
.map { generateTypeHintText(element, it, typeEvalContext) }
1720
.distinct()
1821

1922
if (PyNames.NONE in generatedValues) {
@@ -24,8 +27,24 @@ enum class HintGenerator {
2427
}
2528
},
2629

27-
COLLECTION_TYPE {
28-
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
30+
ASYNC_TYPE() {
31+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
32+
if (type == null || element !is PyFunction) {
33+
return null
34+
}
35+
36+
if (type is PyCollectionType && type.classQName == PyTypingTypeProvider.COROUTINE && element.isAsync) {
37+
return generateTypeHintText(
38+
element, PyTypingTypeProvider.coroutineOrGeneratorElementType(type)?.get(), typeEvalContext
39+
)
40+
}
41+
42+
return null
43+
}
44+
},
45+
46+
COLLECTION_TYPE() {
47+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
2948
if (
3049
type is PyCollectionType
3150
&& type.name != null
@@ -41,7 +60,7 @@ enum class HintGenerator {
4160
return collectionName
4261
}
4362

44-
val elements = type.elementTypes.mapNotNull { generateTypeHintText(it, typeEvalContext) }
63+
val elements = type.elementTypes.mapNotNull { generateTypeHintText(element, it, typeEvalContext) }
4564

4665
if (elements.isEmpty()) {
4766
return collectionName
@@ -54,8 +73,8 @@ enum class HintGenerator {
5473
}
5574
},
5675

57-
TUPLE_TYPE {
58-
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
76+
TUPLE_TYPE() {
77+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
5978
if (type !is PyTupleType) {
6079
return null
6180
}
@@ -65,20 +84,20 @@ enum class HintGenerator {
6584
}
6685

6786
if (type.elementCount > 2) {
68-
val firstElement = generateTypeHintText(type.elementTypes[0], typeEvalContext)
69-
val secondElement = generateTypeHintText(type.elementTypes[1], typeEvalContext)
70-
87+
val firstElement = generateTypeHintText(element, type.elementTypes[0], typeEvalContext)
88+
val secondElement = generateTypeHintText(element, type.elementTypes[1], typeEvalContext)
89+
7190
return "${PyNames.TUPLE}[$firstElement, $secondElement, ...]"
7291
}
7392

7493
return type.elementTypes
75-
.mapNotNull { generateTypeHintText(it, typeEvalContext) }
94+
.mapNotNull { generateTypeHintText(element, it, typeEvalContext) }
7695
.joinToString(separator = ", ", prefix = "${PyNames.TUPLE}[", postfix = "]")
7796
}
7897
},
7998

80-
CLASS_TYPE {
81-
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
99+
CLASS_TYPE() {
100+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
82101
if (type is PyClassType && type.isDefinition) {
83102
return "${PyNames.TYPE.replaceFirstChar { it.titlecaseChar() }}[${type.declarationElement?.name}]"
84103
}
@@ -87,8 +106,8 @@ enum class HintGenerator {
87106
}
88107
},
89108

90-
FUNCTION_TYPE {
91-
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
109+
FUNCTION_TYPE() {
110+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
92111
if (type !is PyFunctionType) {
93112
return null
94113
}
@@ -104,20 +123,20 @@ enum class HintGenerator {
104123

105124
val callableReturnType = typeEvalContext.getReturnType(type.callable)
106125

107-
return "$parametersText -> (${generateTypeHintText(callableReturnType, typeEvalContext)})"
126+
return "$parametersText -> (${generateTypeHintText(element, callableReturnType, typeEvalContext)})"
108127
}
109128
},
110129

111-
ANY_TYPE {
112-
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String {
130+
ANY_TYPE() {
131+
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String {
113132
return type?.name ?: PyNames.UNKNOWN_TYPE
114133
}
115134
};
116135

117-
abstract fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String?
136+
abstract fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String?
118137

119138
companion object {
120-
fun generateTypeHintText(type: PyType?, typeEvalContext: TypeEvalContext): String =
121-
values().firstNotNullOf { it.handleType(type, typeEvalContext) }
139+
fun generateTypeHintText(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String =
140+
values().firstNotNullOf { it.handleType(element, type, typeEvalContext) }
122141
}
123142
}

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -458,7 +458,6 @@ enum class HintResolver {
458458

459459
if (
460460
typeAnnotation == null
461-
|| (element is PyFunction && element.isAsync) // TODO: Handle async functions properly
462461
|| (element is PyFunction && typeAnnotation is PyNoneType)
463462
|| ((element is PyFunction || element is PyTargetExpression) && (element as PyTypeCommentOwner).typeCommentAnnotation != null)
464463
|| (element is PyAnnotationOwner && element.annotation != null)

0 commit comments

Comments
 (0)