Skip to content

Commit ec3fb73

Browse files
divaltordivaltor
authored
Add support for inlay type hint annotations (#6)
* Add support of type hints * Fix invalid handling of LiteralExpression * Fix invalid handling of LiteralExpression and add cases for `set` and `dict comprehension` * Add ignore for _ variables * Fix typings for call chains from class * Add missed bracket * Add tuple support and refactor code * Up limit for Union type * Add class attribute ignore * Add enum type handling * Delete panel * Refactor code, add settings * Minor lambda refactor * Minor refactoring * Add support for ignoring of typing.pyi module and some generic types * Remove explicit type annotation * Handle case when all values inside dict are literal expressions * Clean up code * Add additional check for dict literal expression * Fix invalid dict check * Add handling of #type comments * Remove duplicates from union type * Add resolve for globals\locals keywords * Fix presentation of callable argument if it typing was null * Hide hint for empty literal dict * Add missed peelArgument call * Ignore dict hint if it is empty or just empty "dict" value * Fix description of settings * Add try catch blocks for hints render
1 parent e3941af commit ec3fb73

9 files changed

Lines changed: 864 additions & 2 deletions

File tree

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package space.whitememory.pythoninlayparams
2+
3+
import com.intellij.codeInsight.hints.FactoryInlayHintsCollector
4+
import com.intellij.codeInsight.hints.InlayHintsSink
5+
import com.intellij.openapi.editor.Editor
6+
import com.intellij.psi.PsiElement
7+
import com.jetbrains.python.psi.PyElement
8+
import com.jetbrains.python.psi.types.TypeEvalContext
9+
import space.whitememory.pythoninlayparams.hints.HintGenerator
10+
import space.whitememory.pythoninlayparams.hints.HintResolver
11+
12+
@Suppress("UnstableApiUsage")
13+
abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor, open val settings: Any): FactoryInlayHintsCollector(editor) {
14+
abstract fun validateExpression(element: PsiElement): Boolean
15+
16+
abstract val textBeforeTypeHint: String
17+
18+
protected fun getTypeEvalContext(editor: Editor, element: PsiElement): TypeEvalContext {
19+
return TypeEvalContext.codeCompletion(editor.project!!, element.containingFile)
20+
}
21+
22+
protected fun renderTypeHint(element: PyElement, typeEvalContext: TypeEvalContext, sink: InlayHintsSink) {
23+
val typeAnnotation = HintResolver.getExpressionAnnotationType(element, typeEvalContext)
24+
val hintName = HintGenerator.generateTypeHintText(typeAnnotation, typeEvalContext)
25+
26+
displayTypeHint(element, sink, hintName)
27+
}
28+
29+
abstract fun displayTypeHint(element: PyElement, sink: InlayHintsSink, hintName: String)
30+
31+
override fun collect(element: PsiElement, editor: Editor, sink: InlayHintsSink): Boolean {
32+
if (!element.isValid || element.project.isDefault) {
33+
return false
34+
}
35+
36+
return true
37+
}
38+
}

‎src/main/kotlin/space/whitememory/pythoninlayparams/functions/PythonFunctionInlayTypeHintsCollector.kt‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package space.whitememory.pythoninlayparams.functions
2+
3+
import com.intellij.codeInsight.hints.InlayHintsSink
4+
import com.intellij.openapi.editor.Editor
5+
import com.intellij.psi.PsiElement
6+
import com.intellij.psi.util.PsiTreeUtil
7+
import com.intellij.refactoring.suggested.endOffset
8+
import com.jetbrains.python.psi.PyElement
9+
import com.jetbrains.python.psi.PyFunction
10+
import com.jetbrains.python.psi.PyParameterList
11+
import com.jetbrains.python.psi.PyStatementList
12+
import space.whitememory.pythoninlayparams.AbstractPythonInlayTypeHintsCollector
13+
import space.whitememory.pythoninlayparams.hints.HintGenerator
14+
import space.whitememory.pythoninlayparams.hints.HintResolver
15+
16+
@Suppress("UnstableApiUsage")
17+
class PythonFunctionInlayTypeHintsCollector(editor: Editor, settings: Any) :
18+
AbstractPythonInlayTypeHintsCollector(editor, settings) {
19+
20+
override val textBeforeTypeHint = "->"
21+
22+
override fun validateExpression(element: PsiElement): Boolean {
23+
return element is PyFunction
24+
}
25+
26+
override fun collect(element: PsiElement, editor: Editor, sink: InlayHintsSink): Boolean {
27+
if (!super.collect(element, editor, sink)) {
28+
return false
29+
}
30+
31+
if (!validateExpression(element)) {
32+
return true
33+
}
34+
35+
val typeEvalContext = getTypeEvalContext(editor, element)
36+
37+
if (!HintResolver.shouldShowTypeHint(element as PyFunction, typeEvalContext)) {
38+
return true
39+
}
40+
41+
try {
42+
renderTypeHint(element, typeEvalContext, sink)
43+
} catch (e: Exception) {
44+
return true
45+
}
46+
47+
return true
48+
}
49+
50+
override fun displayTypeHint(element: PyElement, sink: InlayHintsSink, hintName: String) {
51+
val statementList = PsiTreeUtil.getChildOfType(element, PyParameterList::class.java)
52+
53+
statementList?.let {
54+
sink.addInlineElement(
55+
it.endOffset,
56+
false,
57+
factory.roundWithBackground(factory.smallText("$textBeforeTypeHint $hintName")),
58+
false
59+
)
60+
}
61+
}
62+
}

‎src/main/kotlin/space/whitememory/pythoninlayparams/functions/PythonFunctionInlayTypeHintsProvider.kt‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package space.whitememory.pythoninlayparams.functions
2+
3+
import com.intellij.codeInsight.hints.*
4+
import com.intellij.openapi.editor.Editor
5+
import com.intellij.psi.PsiFile
6+
import com.intellij.ui.dsl.builder.panel
7+
import javax.swing.JComponent
8+
9+
@Suppress("UnstableApiUsage")
10+
class PythonFunctionInlayTypeHintsProvider: InlayHintsProvider<NoSettings> {
11+
12+
override val key: SettingsKey<NoSettings> = SettingsKey("python.inlay.function.types")
13+
override val name = "Function type hints"
14+
override val description = "Show the return type of functions in the editor"
15+
override val previewText = null
16+
17+
override val group = InlayGroup.TYPES_GROUP
18+
19+
override fun createSettings(): NoSettings = NoSettings()
20+
21+
override fun getCollectorFor(
22+
file: PsiFile,
23+
editor: Editor,
24+
settings: NoSettings,
25+
sink: InlayHintsSink
26+
): InlayHintsCollector = PythonFunctionInlayTypeHintsCollector(editor, settings)
27+
28+
override fun createConfigurable(settings: NoSettings): ImmediateConfigurable {
29+
return object : ImmediateConfigurable {
30+
override fun createComponent(listener: ChangeListener): JComponent = panel { }
31+
}
32+
}
33+
}

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

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
package space.whitememory.pythoninlayparams.hints
2+
3+
import com.jetbrains.python.PyNames
4+
import com.jetbrains.python.psi.PyLambdaExpression
5+
import com.jetbrains.python.psi.types.*
6+
7+
enum class HintGenerator {
8+
UNION_TYPE() {
9+
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
10+
if (type !is PyUnionType) {
11+
return null
12+
}
13+
14+
val generatedValues = type.members
15+
.filterNotNull()
16+
.map { generateTypeHintText(it, typeEvalContext) }
17+
.distinct()
18+
19+
if (PyNames.NONE in generatedValues) {
20+
return generatedValues.joinToString(separator = " | ", limit = 3)
21+
}
22+
23+
return generatedValues.joinToString(separator = " | ", limit = 2)
24+
}
25+
},
26+
27+
COLLECTION_TYPE() {
28+
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
29+
if (
30+
type is PyCollectionType
31+
&& type.name != null
32+
&& type !is PyTupleType
33+
&& type.elementTypes.isNotEmpty()
34+
) {
35+
val collectionName = when (type) {
36+
is PyTypedDictType -> "dict"
37+
else -> type.name
38+
}
39+
40+
if (type.elementTypes.all { it == null }) {
41+
return collectionName
42+
}
43+
44+
val elements = type.elementTypes.mapNotNull { generateTypeHintText(it, typeEvalContext) }
45+
46+
if (elements.isEmpty()) {
47+
return collectionName
48+
}
49+
50+
return elements.joinToString(separator = ", ", limit = 3, prefix = "$collectionName[", postfix = "]")
51+
}
52+
53+
return null
54+
}
55+
},
56+
57+
TUPLE_TYPE() {
58+
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
59+
if (type !is PyTupleType) {
60+
return null
61+
}
62+
63+
if (type.elementCount == 0 || type.elementTypes.filterNotNull().isEmpty()) {
64+
return PyNames.TUPLE
65+
}
66+
67+
if (type.elementCount > 2) {
68+
val firstElement = generateTypeHintText(type.elementTypes[0], typeEvalContext)
69+
val secondElement = generateTypeHintText(type.elementTypes[1], typeEvalContext)
70+
71+
return "${PyNames.TUPLE}[$firstElement, $secondElement, ...]"
72+
}
73+
74+
return type.elementTypes
75+
.mapNotNull { generateTypeHintText(it, typeEvalContext) }
76+
.joinToString(separator = ", ", prefix = "${PyNames.TUPLE}[", postfix = "]")
77+
}
78+
},
79+
80+
CLASS_TYPE() {
81+
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
82+
if (type is PyClassType && type.isDefinition) {
83+
return "${PyNames.TYPE.replaceFirstChar { it.titlecaseChar() }}[${type.declarationElement?.name}]"
84+
}
85+
86+
return null
87+
}
88+
},
89+
90+
FUNCTION_TYPE() {
91+
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String? {
92+
if (type !is PyFunctionType) {
93+
return null
94+
}
95+
96+
val parametersText = when (type.callable) {
97+
is PyLambdaExpression -> type.callable.parameterList.getPresentableText(false, typeEvalContext)
98+
else -> "(${type.callable.parameterList.parameters
99+
.filter { !it.isSelf }
100+
.joinToString(separator = ", ") { "${HintUtil.getParameterAnnotationValue(it)}" }})"
101+
}
102+
103+
val callableReturnType = typeEvalContext.getReturnType(type.callable)
104+
105+
return "$parametersText -> (${generateTypeHintText(callableReturnType, typeEvalContext)})"
106+
}
107+
},
108+
109+
ANY_TYPE() {
110+
override fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String {
111+
return type?.name ?: PyNames.UNKNOWN_TYPE
112+
}
113+
};
114+
115+
abstract fun handleType(type: PyType?, typeEvalContext: TypeEvalContext): String?
116+
117+
companion object {
118+
fun generateTypeHintText(type: PyType?, typeEvalContext: TypeEvalContext): String =
119+
values().firstNotNullOf { it.handleType(type, typeEvalContext) }
120+
}
121+
}

0 commit comments

Comments
 (0)