Skip to content

Commit fc03417

Browse files
WhiteMemory99WhiteMemory99
committed
Refactor code to generalize the style
1 parent 5f91195 commit fc03417

File tree

9 files changed

+127
-184
lines changed

9 files changed

+127
-184
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,12 +59,13 @@ class PythonInlayParameterHintsProvider : InlayParameterHintsProvider {
5959
if (rootCalleeObject !is PyClass && !functionHints.isEnabled()) return inlayInfos
6060

6161
// Try to resolve the object that made this call, it can be a dataclass/method/function
62+
// codeCompletion is used because simpler contexts are not enough to resolve complex cases
6263
val evalContext = TypeEvalContext.codeCompletion(element.project, element.containingFile)
6364
val resolvedCallee = element.multiResolveCallee(PyResolveContext.defaultContext(evalContext))
6465
if (resolvedCallee.isEmpty() || isForbiddenBuiltinCallable(resolvedCallee[0])) return inlayInfos
6566

6667
// Get the parameters of the call, except `self`, `*` and `/`
67-
val resolvedParameters = resolvedCallee[0].getParameters(evalContext)?.filter { it ->
68+
val resolvedParameters = resolvedCallee[0].getParameters(evalContext)?.filter {
6869
!it.isSelf && it.parameter !is PySingleStarParameter && it.parameter !is PySlashParameter
6970
} ?: return inlayInfos
7071

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

Lines changed: 17 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -13,26 +13,24 @@ import space.whitememory.pythoninlayparams.types.hints.InlayInfoDetails
1313
import space.whitememory.pythoninlayparams.types.hints.PsiInlayInfoDetail
1414

1515
@Suppress("UnstableApiUsage")
16-
abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor, open val settings: Any) :
16+
abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor) :
1717
FactoryInlayHintsCollector(editor) {
18-
abstract fun validateExpression(element: PsiElement): Boolean
1918

2019
abstract val textBeforeTypeHint: String
2120

21+
abstract fun validateExpression(element: PsiElement): Boolean
22+
2223
protected fun getTypeEvalContext(editor: Editor, element: PsiElement): TypeEvalContext {
2324
return TypeEvalContext.codeCompletion(editor.project!!, element.containingFile)
2425
}
2526

2627
protected fun renderTypeHint(element: PyElement, typeEvalContext: TypeEvalContext, sink: InlayHintsSink) {
2728
val typeAnnotation = HintResolver.getExpressionAnnotationType(element, typeEvalContext)
28-
val hintName = HintGenerator.generateTypeHintText(element, typeAnnotation, typeEvalContext)
29-
30-
if (hintName.isEmpty()) {
31-
return
32-
}
29+
val hintName =
30+
HintGenerator.generateTypeHintText(element, typeAnnotation, typeEvalContext).takeIf { it.isNotEmpty() }
31+
?: return
3332

3433
val resolvedHintName = resolveInlayPresentation(hintName)
35-
3634
displayTypeHint(element, sink, resolvedHintName)
3735
}
3836

@@ -48,23 +46,27 @@ abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor, open val se
4846

4947
private fun getInlayPresentationForInlayInfoDetails(infoDetail: InlayInfoDetails): InlayPresentation {
5048
if (infoDetail.rootInlayInfo == null) {
51-
return resolveInlayPresentation(infoDetail.details, separator = infoDetail.separator, limit = infoDetail.limit)
49+
return resolveInlayPresentation(
50+
infoDetail.details,
51+
separator = infoDetail.separator,
52+
limit = infoDetail.limit
53+
)
5254
}
5355

5456
val textPresentation = factory.smallText(infoDetail.rootInlayInfo.text)
55-
val navigationElementProvider: (() -> PsiElement?)? = when(infoDetail.rootInlayInfo) {
56-
is PsiInlayInfoDetail -> {{ infoDetail.rootInlayInfo.element }}
57+
val navigationElementProvider = when (infoDetail.rootInlayInfo) {
58+
is PsiInlayInfoDetail -> {
59+
{ infoDetail.rootInlayInfo.element }
60+
}
61+
5762
else -> null
5863
}
5964

6065
val basePresentation = navigationElementProvider?.let {
6166
factory.psiSingleReference(textPresentation, it)
6267
} ?: textPresentation
6368

64-
if (infoDetail.details.isEmpty()) {
65-
return basePresentation
66-
}
67-
69+
if (infoDetail.details.isEmpty()) return basePresentation
6870
val childDetails = infoDetail.details.map { getInlayPresentationForInlayInfoDetails(it) }
6971

7072
return factory.seq(
@@ -114,12 +116,4 @@ abstract class AbstractPythonInlayTypeHintsCollector(editor: Editor, open val se
114116
}
115117

116118
abstract fun displayTypeHint(element: PyElement, sink: InlayHintsSink, hintName: InlayPresentation)
117-
118-
override fun collect(element: PsiElement, editor: Editor, sink: InlayHintsSink): Boolean {
119-
if (!element.isValid || element.project.isDefault) {
120-
return false
121-
}
122-
123-
return true
124-
}
125119
}

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

Lines changed: 13 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,35 +15,25 @@ import space.whitememory.pythoninlayparams.types.AbstractPythonInlayTypeHintsCol
1515
import space.whitememory.pythoninlayparams.types.hints.HintResolver
1616

1717
@Suppress("UnstableApiUsage")
18-
class PythonFunctionInlayTypeHintsCollector(editor: Editor, settings: Any) :
19-
AbstractPythonInlayTypeHintsCollector(editor, settings) {
18+
class PythonFunctionInlayTypeHintsCollector(editor: Editor) :
19+
AbstractPythonInlayTypeHintsCollector(editor) {
2020

2121
override val textBeforeTypeHint = "->"
2222

2323
override fun validateExpression(element: PsiElement): Boolean {
2424
// Not a function or has only def keyword
25-
if (element !is PyFunction || element.nameNode == null) {
26-
return false
27-
}
25+
if (element !is PyFunction || element.nameNode == null) return false
2826

2927
val colonToken = TokenSet.create(PyTokenTypes.COLON)
3028
return element.node.getChildren(colonToken).isNotEmpty()
3129
}
3230

3331
override fun collect(element: PsiElement, editor: Editor, sink: InlayHintsSink): Boolean {
34-
if (!super.collect(element, editor, sink)) {
35-
return false
36-
}
37-
38-
if (!validateExpression(element)) {
39-
return true
40-
}
32+
if (!element.isValid || element.project.isDefault) return false
33+
if (!validateExpression(element)) return true
4134

4235
val typeEvalContext = getTypeEvalContext(editor, element)
43-
44-
if (!HintResolver.shouldShowTypeHint(element as PyFunction, typeEvalContext)) {
45-
return true
46-
}
36+
if (!HintResolver.shouldShowTypeHint(element as PyFunction, typeEvalContext)) return true
4737

4838
try {
4939
renderTypeHint(element, typeEvalContext, sink)
@@ -55,15 +45,12 @@ class PythonFunctionInlayTypeHintsCollector(editor: Editor, settings: Any) :
5545
}
5646

5747
override fun displayTypeHint(element: PyElement, sink: InlayHintsSink, hintName: InlayPresentation) {
58-
val statementList = PsiTreeUtil.getChildOfType(element, PyParameterList::class.java)
59-
60-
statementList?.let {
61-
sink.addInlineElement(
62-
it.endOffset,
63-
false,
64-
factory.roundWithBackground(factory.seq(factory.smallText("$textBeforeTypeHint "), hintName)),
65-
false
66-
)
67-
}
48+
val statementList = PsiTreeUtil.getChildOfType(element, PyParameterList::class.java) ?: return
49+
sink.addInlineElement(
50+
statementList.endOffset,
51+
false,
52+
factory.roundWithBackground(factory.seq(factory.smallText("$textBeforeTypeHint "), hintName)),
53+
false
54+
)
6855
}
6956
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ class PythonFunctionInlayTypeHintsProvider : InlayHintsProvider<NoSettings> {
2323
editor: Editor,
2424
settings: NoSettings,
2525
sink: InlayHintsSink
26-
): InlayHintsCollector = PythonFunctionInlayTypeHintsCollector(editor, settings)
26+
): InlayHintsCollector = PythonFunctionInlayTypeHintsCollector(editor)
2727

2828
override fun createConfigurable(settings: NoSettings): ImmediateConfigurable {
2929
return object : ImmediateConfigurable {

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

Lines changed: 63 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -3,21 +3,21 @@
33
package space.whitememory.pythoninlayparams.types.hints
44

55
import com.intellij.psi.PsiElement
6-
import com.intellij.util.containers.toArray
76
import com.jetbrains.python.PyNames
87
import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
9-
import com.jetbrains.python.psi.PyClass
108
import com.jetbrains.python.psi.PyElement
119
import com.jetbrains.python.psi.PyFunction
1210
import com.jetbrains.python.psi.PyLambdaExpression
1311
import com.jetbrains.python.psi.types.*
1412

1513
enum class HintGenerator {
16-
UNION_TYPE() {
17-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>? {
18-
if (type !is PyUnionType) {
19-
return null
20-
}
14+
UNION_TYPE {
15+
override fun handleType(
16+
element: PyElement,
17+
type: PyType?,
18+
typeEvalContext: TypeEvalContext
19+
): List<InlayInfoDetails>? {
20+
if (type !is PyUnionType) return null
2121

2222
val generatedValues = type.members
2323
.filterNotNull()
@@ -48,11 +48,13 @@ enum class HintGenerator {
4848
}
4949
},
5050

51-
ASYNC_TYPE() {
52-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>? {
53-
if (type == null || element !is PyFunction) {
54-
return null
55-
}
51+
ASYNC_TYPE {
52+
override fun handleType(
53+
element: PyElement,
54+
type: PyType?,
55+
typeEvalContext: TypeEvalContext
56+
): List<InlayInfoDetails>? {
57+
if (type == null || element !is PyFunction) return null
5658

5759
if (type is PyCollectionType && type.classQName == PyTypingTypeProvider.COROUTINE && element.isAsync) {
5860
return generateTypeHintText(
@@ -64,8 +66,12 @@ enum class HintGenerator {
6466
}
6567
},
6668

67-
COLLECTION_TYPE() {
68-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>? {
69+
COLLECTION_TYPE {
70+
override fun handleType(
71+
element: PyElement,
72+
type: PyType?,
73+
typeEvalContext: TypeEvalContext
74+
): List<InlayInfoDetails>? {
6975
if (
7076
type is PyCollectionType
7177
&& type.name != null
@@ -87,7 +93,8 @@ enum class HintGenerator {
8793
)
8894
}
8995

90-
val elements = type.elementTypes.mapNotNull { generateTypeHintText(element, it, typeEvalContext) }.flatten()
96+
val elements =
97+
type.elementTypes.mapNotNull { generateTypeHintText(element, it, typeEvalContext) }.flatten()
9198

9299
if (elements.isEmpty()) {
93100
return listOf(
@@ -110,11 +117,13 @@ enum class HintGenerator {
110117
}
111118
},
112119

113-
TUPLE_TYPE() {
114-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>? {
115-
if (type !is PyTupleType) {
116-
return null
117-
}
120+
TUPLE_TYPE {
121+
override fun handleType(
122+
element: PyElement,
123+
type: PyType?,
124+
typeEvalContext: TypeEvalContext
125+
): List<InlayInfoDetails>? {
126+
if (type !is PyTupleType) return null
118127

119128
val baseInlayDetail = resolvePsiReference(type, PyNames.TUPLE)
120129

@@ -151,14 +160,15 @@ enum class HintGenerator {
151160
}
152161
},
153162

154-
CLASS_TYPE() {
155-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>? {
156-
if (type !is PyClassType) {
157-
return null
158-
}
163+
CLASS_TYPE {
164+
override fun handleType(
165+
element: PyElement,
166+
type: PyType?,
167+
typeEvalContext: TypeEvalContext
168+
): List<InlayInfoDetails>? {
169+
if (type !is PyClassType) return null
159170

160171
val baseInlayDetail = resolvePsiReference(type, type.declarationElement?.name!!)
161-
162172
val classInlayDetails = listOf(InlayInfoDetails(baseInlayDetail))
163173

164174
if (type.isDefinition) {
@@ -170,19 +180,19 @@ enum class HintGenerator {
170180
return listOf(inlayDetail)
171181
}
172182

173-
if (!type.isDefinition) {
174-
return classInlayDetails
175-
}
183+
if (!type.isDefinition) return classInlayDetails
176184

177185
return null
178186
}
179187
},
180188

181-
FUNCTION_TYPE() {
182-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>? {
183-
if (type !is PyFunctionType) {
184-
return null
185-
}
189+
FUNCTION_TYPE {
190+
override fun handleType(
191+
element: PyElement,
192+
type: PyType?,
193+
typeEvalContext: TypeEvalContext
194+
): List<InlayInfoDetails>? {
195+
if (type !is PyFunctionType) return null
186196

187197
val parametersText = when (type.callable) {
188198
is PyLambdaExpression -> type.callable.parameterList.getPresentableText(false, typeEvalContext)
@@ -213,15 +223,23 @@ enum class HintGenerator {
213223
}
214224
},
215225

216-
ANY_TYPE() {
217-
override fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails> {
226+
ANY_TYPE {
227+
override fun handleType(
228+
element: PyElement,
229+
type: PyType?,
230+
typeEvalContext: TypeEvalContext
231+
): List<InlayInfoDetails> {
218232
val baseInlayDetail = resolvePsiReference(type, type?.name ?: PyNames.UNKNOWN_TYPE)
219233

220234
return listOf(InlayInfoDetails(baseInlayDetail))
221235
}
222236
};
223237

224-
abstract fun handleType(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails>?
238+
abstract fun handleType(
239+
element: PyElement,
240+
type: PyType?,
241+
typeEvalContext: TypeEvalContext
242+
): List<InlayInfoDetails>?
225243

226244
companion object {
227245

@@ -230,13 +248,17 @@ enum class HintGenerator {
230248
return PsiInlayInfoDetail(name, it)
231249
} ?: TextInlayInfoDetail(name)
232250
}
233-
fun generateTypeHintText(element: PyElement, type: PyType?, typeEvalContext: TypeEvalContext): List<InlayInfoDetails> =
251+
252+
fun generateTypeHintText(
253+
element: PyElement,
254+
type: PyType?,
255+
typeEvalContext: TypeEvalContext
256+
): List<InlayInfoDetails> =
234257
values().firstNotNullOf { it.handleType(element, type, typeEvalContext) }
235258
}
236259
}
237260

238261

239-
@Suppress("UnstableApiUsage")
240262
data class InlayInfoDetails(
241263
val rootInlayInfo: InlayInfoDetail?,
242264
val details: List<InlayInfoDetails> = listOf(),
@@ -246,5 +268,5 @@ data class InlayInfoDetails(
246268

247269
sealed class InlayInfoDetail(val text: String)
248270

249-
class TextInlayInfoDetail(text: String): InlayInfoDetail(text)
250-
class PsiInlayInfoDetail(text: String, val element: PsiElement): InlayInfoDetail(text)
271+
class TextInlayInfoDetail(text: String) : InlayInfoDetail(text)
272+
class PsiInlayInfoDetail(text: String, val element: PsiElement) : InlayInfoDetail(text)

0 commit comments

Comments
 (0)