Skip to content

Add support for async functions #15

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
WhiteMemory99 merged 1 commit into from
Sep 24, 2022
Merged

Add support for async functions #15

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ...kotlin/space/whitememory/pythoninlayparams/types/AbstractPythonInlayTypeHintsCollector.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor, open val se

protected fun renderTypeHint(element: PyElement, typeEvalContext: TypeEvalContext, sink: InlayHintsSink) {
val typeAnnotation = HintResolver.getExpressionAnnotationType(element, typeEvalContext)
val hintName = HintGenerator.generateTypeHintText(typeAnnotation, typeEvalContext)
val hintName = HintGenerator.generateTypeHintText(element, typeAnnotation, typeEvalContext)

displayTypeHint(element, sink, hintName)
}
Expand Down
63 changes: 41 additions & 22 deletions src/main/kotlin/space/whitememory/pythoninlayparams/types/hints/HintGenerator.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,22 @@
package space.whitememory.pythoninlayparams.types.hints

import com.jetbrains.python.PyNames
import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
import com.jetbrains.python.psi.PyElement
import com.jetbrains.python.psi.PyFunction
import com.jetbrains.python.psi.PyLambdaExpression
import com.jetbrains.python.psi.types.*

enum class HintGenerator {
UNION_TYPE {
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
UNION_TYPE() {
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
if (type !is PyUnionType) {
return null
}

val generatedValues = type.members
.filterNotNull()
.map { generateTypeHintText(it, typeEvalContext) }
.map { generateTypeHintText(element, it, typeEvalContext) }
.distinct()

if (PyNames.NONE in generatedValues) {
Expand All @@ -24,8 +27,24 @@ enum class HintGenerator {
}
},

COLLECTION_TYPE {
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
ASYNC_TYPE() {
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
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, type: PyType?, typeEvalContext: TypeEvalContext): String? {
if (
type is PyCollectionType
&& type.name != null
Expand All @@ -41,7 +60,7 @@ enum class HintGenerator {
return collectionName
}

val elements = type.elementTypes.mapNotNull { generateTypeHintText(it, typeEvalContext) }
val elements = type.elementTypes.mapNotNull { generateTypeHintText(element, it, typeEvalContext) }

if (elements.isEmpty()) {
return collectionName
Expand All @@ -54,8 +73,8 @@ enum class HintGenerator {
}
},

TUPLE_TYPE {
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
TUPLE_TYPE() {
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
if (type !is PyTupleType) {
return null
}
Expand All @@ -65,20 +84,20 @@ enum class HintGenerator {
}

if (type.elementCount > 2) {
val firstElement = generateTypeHintText(type.elementTypes[0], typeEvalContext)
val secondElement = generateTypeHintText(type.elementTypes[1], typeEvalContext)

val firstElement = generateTypeHintText(element, type.elementTypes[0], typeEvalContext)
val secondElement = generateTypeHintText(element, type.elementTypes[1], typeEvalContext)
return "${PyNames.TUPLE}[$firstElement, $secondElement, ...]"
}

return type.elementTypes
.mapNotNull { generateTypeHintText(it, typeEvalContext) }
.mapNotNull { generateTypeHintText(element, it, typeEvalContext) }
.joinToString(separator = ", ", prefix = "${PyNames.TUPLE}[", postfix = "]")
}
},

CLASS_TYPE {
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
CLASS_TYPE() {
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
if (type is PyClassType && type.isDefinition) {
return "${PyNames.TYPE.replaceFirstChar { it.titlecaseChar() }}[${type.declarationElement?.name}]"
}
Expand All @@ -87,8 +106,8 @@ enum class HintGenerator {
}
},

FUNCTION_TYPE {
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
FUNCTION_TYPE() {
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String? {
if (type !is PyFunctionType) {
return null
}
Expand All @@ -104,20 +123,20 @@ enum class HintGenerator {

val callableReturnType = typeEvalContext.getReturnType(type.callable)

return "$parametersText -> (${generateTypeHintText(callableReturnType, typeEvalContext)})"
return "$parametersText -> (${generateTypeHintText(element, callableReturnType, typeEvalContext)})"
}
},

ANY_TYPE {
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String {
ANY_TYPE() {
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String {
return type?.name ?: PyNames.UNKNOWN_TYPE
}
};

abstract fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String?
abstract fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String?

companion object {
fun generateTypeHintText(type: PyType?, typeEvalContext: TypeEvalContext): String =
values().firstNotNullOf { it.handleType(type, typeEvalContext) }
fun generateTypeHintText(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): String =
values().firstNotNullOf { it.handleType(element, type, typeEvalContext) }
}
}
1 change: 0 additions & 1 deletion src/main/kotlin/space/whitememory/pythoninlayparams/types/hints/HintResolver.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,6 @@ enum class HintResolver {

if (
typeAnnotation == null
|| (element is PyFunction && element.isAsync) // TODO: Handle async functions properly
|| (element is PyFunction && typeAnnotation is PyNoneType)
|| ((element is PyFunction || element is PyTargetExpression) && (element as PyTypeCommentOwner).typeCommentAnnotation != null)
|| (element is PyAnnotationOwner && element.annotation != null)
Expand Down