1+ package space.whitememory.pythoninlayparams
2+
3+ import com.intellij.codeInsight.hints.FactoryInlayHintsCollector
4+ import com.intellij.codeInsight.hints.InlayHintsSink
5+ import com.intellij.codeInsight.hints.InlayInfo
6+ import com.intellij.openapi.editor.Editor
7+ import com.intellij.psi.PsiElement
8+ import com.intellij.psi.util.PsiTreeUtil
9+ import com.jetbrains.python.psi.*
10+ import com.jetbrains.python.psi.types.TypeEvalContext
11+
12+ @Suppress(" UnstableApiUsage"
13+ class PythonInlayHintsCollector (
14+ editor : Editor , private val settings : PythonInlayHintsProvider .Settings
15+ ) : FactoryInlayHintsCollector(editor) {
16+
17+ private val forbiddenBuiltinFiles = setOf (" builtins.pyi" " typing.pyi"
18+
19+ override fun collect (element : PsiElement , editor : Editor , sink : InlayHintsSink ): Boolean {
20+ // This method gets every element in the editor,
21+ // so we have to verify it's a Python call expression
22+ if (element !is PyCallExpression || element is PyDecorator ) {
23+ return true
24+ }
25+
26+ // Don't show hints if there's no arguments
27+ // Or the only argument is unpacking (*list, **dict)
28+ if (element.arguments.isEmpty() || (element.arguments.size == 1 && element.arguments[0 ] is PyStarArgument )) {
29+ return true
30+ }
31+
32+ // Try to resolve the object that made this call
33+ var resolved = element.callee?.reference?.resolve() ? : return true
34+ if (isForbiddenBuiltinElement(resolved)) {
35+ return true
36+ }
37+
38+ var classAttributes = listOf<PyTargetExpression >()
39+ if (resolved is PyTargetExpression && settings.lambdaHints) {
40+ // TODO: Handle cases other than lambda expressions
41+ // Use the target to find the lambda expression object, and assign it to get its parameters up ahead
42+ resolved = PsiTreeUtil .getNextSiblingOfType(resolved, PyLambdaExpression ::class .java) ? : return true
43+ } else if (resolved is PyClass && settings.classHints) {
44+ // This call is made by a class (initialization), so we want to find the parameters it takes.
45+ // In order to do so, we first have to check for an init method, and if not found,
46+ // We will use the class attributes instead. (Handle dataclasses, attrs, etc.)
47+ val evalContext = TypeEvalContext .codeAnalysis(element.project, element.containingFile)
48+ val entryMethod = resolved.findInitOrNew(true , evalContext)
49+
50+ resolved = if (entryMethod != null && entryMethod.containingClass == resolved) {
51+ entryMethod
52+ } else {
53+ // Use the class attributes if there's no init with params in the parent classes
54+ // TODO: Make sure wrong attributes are not used
55+ classAttributes = resolved.classAttributes
56+ entryMethod ? : resolved
57+ }
58+ } else if (! settings.functionHints) {
59+ return true
60+ }
61+
62+ val resolvedParameters = getElementFilteredParameters(resolved)
63+ val finalParameters = if (resolvedParameters.isEmpty() && classAttributes.isNotEmpty()) {
64+ // If there's no parameters in the object,
65+ // we use the class attributes instead,
66+ // in case this is a class
67+ classAttributes
68+ } else if (resolvedParameters.isEmpty()) {
69+ return true
70+ } else {
71+ resolvedParameters
72+ }
73+
74+ if (finalParameters.size == 1 ) {
75+ // Don't need a hint if there's only one parameter,
76+ // Make an exception for *args
77+ finalParameters[0 ].let {
78+ if (it !is PyNamedParameter || ! it.isPositionalContainer) return true
79+ }
80+ }
81+
82+ getInlayInfos(finalParameters, element.arguments).forEach {
83+ val hintText = factory.smallText(" ${it.text} :"
84+ val presentation = factory.roundWithBackground(hintText)
85+
86+ sink.addInlineElement(it.offset, false , presentation, false )
87+ }
88+
89+ return true
90+ }
91+
92+ /* *
93+ * Gets the list of [InlayInfo] that represents parameters worth showing along with their offset.
94+ */
95+ private fun getInlayInfos (parameters : List <PyElement >, arguments : Array <PyExpression >): MutableList <InlayInfo > {
96+ val inlayInfos = mutableListOf<InlayInfo >()
97+
98+ parameters.zip(arguments).forEach { (param, arg) ->
99+ val paramName = param.name ? : return @forEach
100+ if (arg is PyStarArgument || arg is PyKeywordArgument ) {
101+ // It's a keyword argument or unpacking,
102+ // we don't need to show hits after this
103+ return inlayInfos
104+ }
105+
106+ if (param is PyNamedParameter ) {
107+ if (param.isPositionalContainer) {
108+ // This is an *args parameter that takes more than one argument
109+ // So we show it and stop the further processing of this call expression
110+ inlayInfos.add(InlayInfo (" ...$paramName "
111+ return inlayInfos
112+ } else if (param.isKeywordContainer) {
113+ // We don't want to show `kwargs` as a hint by accident
114+ return inlayInfos
115+ }
116+ }
117+
118+ if (isHintNameValid(paramName, arg)) {
119+ inlayInfos.add(InlayInfo (paramName, arg.textOffset))
120+ }
121+ }
122+
123+ return inlayInfos
124+ }
125+
126+ /* *
127+ * Checks if the given parameter name is valid for the given argument.
128+ * This is used to skip parameters that start with __, or are the same as the argument.
129+ */
130+ private fun isHintNameValid (name : String , argument : PyExpression ): Boolean {
131+ // TODO: More filters
132+ return name != argument.name?.lowercase() && ! name.startsWith(" __" && name.length > 1
133+ }
134+
135+ /* *
136+ * Get the parameters of the element, but filter out the ones that are not needed.
137+ * For example, if the element is a class method, we don't want to show the __self__ parameter.
138+ */
139+ private fun getElementFilteredParameters (element : PsiElement ): List <PyParameter > {
140+ element.children.forEach {
141+ if (it is PyParameterList ) {
142+ return it.parameters.filter { param -> ! param.isSelf }
143+ }
144+ }
145+ return emptyList()
146+ }
147+
148+ /* *
149+ * Checks if the element is part of the standard library that isn't relevant for these hints.
150+ */
151+ private fun isForbiddenBuiltinElement (element : PsiElement ): Boolean {
152+ // TODO: Implement using PyType.isBuiltin (?),
153+ // although we still want some builtins like datetime.datetime
154+ return element.containingFile.name in forbiddenBuiltinFiles
155+ }
156+ }
0 commit comments