Add type and default value to popup of auto-completion

resources/META-INF/plugin.xml

@@ -4,6 +4,12 @@
     <version>0.0.15</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.0.16</h2>
+    <p>Features, BugFixes</p>
+    <ul>
+        <li>Add types and default values to popup of auto-completion [#54] </li>
+        <li>Fix class imported path on auto-completion [#54] </li>
+    </ul>
     <h2>version 0.0.15</h2>
     <p>Features</p>
     <ul>

src/com/koxudaxi/pydantic/Pydantic.kt

@@ -13,32 +13,32 @@ const val SCHEMA_Q_NAME = "pydantic.schema.Schema"
 const val FIELD_Q_NAME = "pydantic.field.Field"
 const val BASE_SETTINGS_Q_NAME = "pydantic.env_settings.BaseSettings"
 
-fun getPyClassByPyCallExpression(pyCallExpression: PyCallExpression): PyClass? {
+internal fun getPyClassByPyCallExpression(pyCallExpression: PyCallExpression): PyClass? {
     return pyCallExpression.callee?.reference?.resolve() as? PyClass
 }
 
-fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument): PyClass? {
+internal fun getPyClassByPyKeywordArgument(pyKeywordArgument: PyKeywordArgument): PyClass? {
     val pyCallExpression = PsiTreeUtil.getParentOfType(pyKeywordArgument, PyCallExpression::class.java) ?: return null
     return getPyClassByPyCallExpression(pyCallExpression)
 }
 
-fun isPydanticModel(pyClass: PyClass, context: TypeEvalContext? = null): Boolean {
-    return isSubClassOfPydanticBaseModel(pyClass, context) || isPydanticDataclass(pyClass)
+internal fun isPydanticModel(pyClass: PyClass, context: TypeEvalContext? = null): Boolean {
+    return isSubClassOfPydanticBaseModel(pyClass, context) || isPydanticDataclass(pyClass) && !isPydanticBaseModel(pyClass)
 }
 
-fun isPydanticBaseModel(pyClass: PyClass): Boolean {
+internal fun isPydanticBaseModel(pyClass: PyClass): Boolean {
     return pyClass.qualifiedName == BASE_MODEL_Q_NAME
 }
 
-fun isSubClassOfPydanticBaseModel(pyClass: PyClass, context: TypeEvalContext?): Boolean {
+internal fun isSubClassOfPydanticBaseModel(pyClass: PyClass, context: TypeEvalContext?): Boolean {
     return pyClass.isSubclass(BASE_MODEL_Q_NAME, context)
 }
 
-fun isBaseSetting(pyClass: PyClass, context: TypeEvalContext): Boolean {
+internal fun isBaseSetting(pyClass: PyClass, context: TypeEvalContext): Boolean {
     return pyClass.isSubclass(BASE_SETTINGS_Q_NAME, context)
 }
 
-fun hasDecorator(pyDecoratable: PyDecoratable, refName: String): Boolean {
+internal fun hasDecorator(pyDecoratable: PyDecoratable, refName: String): Boolean {
     pyDecoratable.decoratorList?.decorators?.mapNotNull { it.callee as? PyReferenceExpression }?.forEach {
         PyResolveUtil.resolveImportedElementQNameLocally(it).forEach { decoratorQualifiedName ->
             if (decoratorQualifiedName == QualifiedName.fromDottedString(refName)) return true
@@ -47,14 +47,14 @@ fun hasDecorator(pyDecoratable: PyDecoratable, refName: String): Boolean {
     return false
 }
 
-fun isPydanticDataclass(pyClass: PyClass): Boolean {
+internal fun isPydanticDataclass(pyClass: PyClass): Boolean {
     return hasDecorator(pyClass, DATA_CLASS_Q_NAME)
 }
 
-fun isPydanticField(pyClass: PyClass, context: TypeEvalContext): Boolean {
+internal fun isPydanticField(pyClass: PyClass, context: TypeEvalContext): Boolean {
     return pyClass.isSubclass(SCHEMA_Q_NAME, context) || pyClass.isSubclass(FIELD_Q_NAME, context)
 }
 
-fun validatorMethod(pyFunction: PyFunction): Boolean {
+internal fun validatorMethod(pyFunction: PyFunction): Boolean {
     return hasDecorator(pyFunction, VALIDATOR_Q_NAME)
-}
+}

src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt

@@ -6,75 +6,134 @@ import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.icons.AllIcons
 import com.intellij.patterns.PlatformPatterns.psiElement
 import com.intellij.util.ProcessingContext
-import com.jetbrains.python.PythonLanguage
+import com.jetbrains.python.PyTokenTypes
 import com.jetbrains.python.codeInsight.completion.getTypeEvalContext
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
+import com.jetbrains.python.documentation.PythonDocumentationProvider
 import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.resolve.PyResolveUtil
 import com.jetbrains.python.psi.types.TypeEvalContext
+import javax.swing.Icon
+
 
 class PydanticCompletionContributor : CompletionContributor() {
     init {
         extend(CompletionType.BASIC,
-                psiElement()
-                        .withLanguage(PythonLanguage.getInstance())
-                        .and(psiElement().withParents(
-                                PyReferenceExpression::class.java,
-                                PyArgumentList::class.java,
-                                PyCallExpression::class.java)
-                        ),
+                psiElement(PyTokenTypes.IDENTIFIER).withParents(
+                        PyReferenceExpression::class.java,
+                        PyArgumentList::class.java,
+                        PyCallExpression::class.java),
                 KeywordArgumentCompletionProvider)
+        extend(CompletionType.BASIC,
+                psiElement(PyTokenTypes.IDENTIFIER)
+                        .afterLeaf(psiElement(PyTokenTypes.DOT))
+                        .withParent(psiElement(PyReferenceExpression::class.java)),
+                FieldCompletionProvider)
     }
 
-    private object KeywordArgumentCompletionProvider : CompletionProvider<CompletionParameters>() {
-
-        override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
-            val pyArgumentList = parameters.position.parent!!.parent !!as PyArgumentList
-            val pyCallExpression = pyArgumentList.parent !!as PyCallExpression
-            val pyClass = getPyClassByPyCallExpression(pyCallExpression) ?: return
-            val typeEvalContext = parameters.getTypeEvalContext()
-
-            if (!isPydanticModel(pyClass, typeEvalContext)) return
-
-            val definedSet = pyArgumentList.children
-                    .mapNotNull { (it as? PyKeywordArgument)?.name }
-                    .map { "${it}=" }
-                    .toHashSet()
+    private abstract class PydanticCompletionProvider : CompletionProvider<CompletionParameters>() {
 
-            val newElements: LinkedHashMap<String, LookupElement> = LinkedHashMap()
+        abstract val icon: Icon
 
-            pyClass.getAncestorClasses(typeEvalContext)
-                    .filter { isPydanticModel(it) && !isPydanticBaseModel(it) }
-                    .forEach { addFieldElement(it, definedSet, newElements, typeEvalContext) }
+        abstract fun getLookupNameFromFieldName(fieldName: String): String
 
-            addFieldElement(pyClass, definedSet, newElements, typeEvalContext)
 
-            result.runRemainingContributors(parameters)
-            { completionResult ->
-                completionResult.lookupElement.lookupString
-                        .takeIf { name -> !newElements.containsKey(name) && !definedSet.contains(name) }
-                        ?.let { result.passResult(completionResult) }
-            }
-            result.addAllElements(newElements.values)
+        private fun getTypeHint(pyClass: PyClass, typeEvalContext: TypeEvalContext, pyTargetExpression: PyTargetExpression): String {
+            val defaultValue = pyTargetExpression.findAssignedValue()?.text?.let { "=$it" } ?: ""
+            val typeHint = PythonDocumentationProvider.getTypeHint(typeEvalContext.getType(pyTargetExpression), typeEvalContext)
+            return "${typeHint}$defaultValue ${pyClass.name}"
         }
 
-        private fun addFieldElement(pyClass: PyClass, excludes: HashSet<String>, results: LinkedHashMap<String, LookupElement>, typeEvalContext: TypeEvalContext) {
+        private fun addFieldElement(pyClass: PyClass, results: LinkedHashMap<String, LookupElement>, typeEvalContext: TypeEvalContext, excludes: HashSet<String>?) {
             pyClass.classAttributes
                     .asReversed()
                     .asSequence()
                     .filterNot { PyTypingTypeProvider.isClassVar(it, typeEvalContext) }
                     .filter { it.name != null }
                     .forEach {
-                        val elementName = "${it.name!!}="
-                        if (!excludes.contains(elementName)) {
-                            val className = pyClass.qualifiedName ?: pyClass.name
+                        val elementName = getLookupNameFromFieldName(it.name!!)
+                        if (excludes == null || !excludes.contains(elementName)) {
                             val element = PrioritizedLookupElement.withGrouping(
                                     LookupElementBuilder
                                             .createWithSmartPointer(elementName, it)
-                                            .withTypeText(className).withIcon(AllIcons.Nodes.Parameter), 1)
+                                            .withTypeText(getTypeHint(pyClass, typeEvalContext, it))
+                                            .withIcon(icon), 1)
                             results[elementName] = PrioritizedLookupElement.withPriority(element, 100.0)
                         }
                     }
         }
+
+        protected fun addAllFieldElement(parameters: CompletionParameters, result: CompletionResultSet,
+                                         pyClass: PyClass, typeEvalContext: TypeEvalContext,
+                                         excludes: HashSet<String>? = null) {
+
+            val newElements: LinkedHashMap<String, LookupElement> = LinkedHashMap()
+
+            pyClass.getAncestorClasses(typeEvalContext)
+                    .filter { isPydanticModel(it) }
+                    .forEach { addFieldElement(it, newElements, typeEvalContext, excludes) }
+
+            addFieldElement(pyClass, newElements, typeEvalContext, excludes)
+
+            result.runRemainingContributors(parameters)
+            { completionResult ->
+                completionResult.lookupElement.lookupString
+                        .takeIf { name -> !newElements.containsKey(name) && (excludes == null || !excludes.contains(name)) }
+                        ?.let { result.passResult(completionResult) }
+            }
+            result.addAllElements(newElements.values)
+        }
+
+    }
+
+    private object KeywordArgumentCompletionProvider : PydanticCompletionProvider() {
+        override fun getLookupNameFromFieldName(fieldName: String): String {
+            return "${fieldName}="
+        }
+
+        override val icon: Icon = AllIcons.Nodes.Parameter
+
+        override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
+            val pyArgumentList = parameters.position.parent!!.parent!! as PyArgumentList
+            val pyCallExpression = pyArgumentList.parent!! as PyCallExpression
+            val pyClass = getPyClassByPyCallExpression(pyCallExpression) ?: return
+            val typeEvalContext = parameters.getTypeEvalContext()
+
+            if (!isPydanticModel(pyClass, typeEvalContext)) return
+
+            val definedSet = pyArgumentList.children
+                    .mapNotNull { (it as? PyKeywordArgument)?.name }
+                    .map { "${it}=" }
+                    .toHashSet()
+
+            addAllFieldElement(parameters, result, pyClass, typeEvalContext, definedSet)
+        }
     }
 
+    private object FieldCompletionProvider : PydanticCompletionProvider() {
+        override fun getLookupNameFromFieldName(fieldName: String): String {
+            return fieldName
+        }
+
+        override val icon: Icon = AllIcons.Nodes.Field
+
+        override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
+            val pyCallExpression = when (val instance = parameters.position.parent.firstChild) {
+                is PyReferenceExpression -> PyResolveUtil.fullResolveLocally(instance) as? PyCallExpression ?: return
+                is PyCallExpression -> instance
+                else -> return
+            }
+            val pyClass = getPyClassByPyCallExpression(pyCallExpression) ?: return
+
+            val typeEvalContext = parameters.getTypeEvalContext()
+
+            if (!isPydanticModel(pyClass, typeEvalContext)) return
+
+
+            addAllFieldElement(parameters, result, pyClass, typeEvalContext)
+        }
+
+    }
+
+
 }

src/com/koxudaxi/pydantic/PydanticFieldRenameFactory.kt

@@ -73,7 +73,7 @@ class PydanticFieldRenameFactory : AutomaticRenamerFactory {
             addClassAttributes(pyClass, elementName)
             addKeywordArguments(pyClass, elementName)
             pyClass.getAncestorClasses(null)
-                    .filter { !isPydanticBaseModel(it) && isPydanticModel(it) && !added.contains(it) }
+                    .filter { isPydanticModel(it) && !added.contains(it) }
                     .forEach { addAllElement(it, elementName, added) }
 
             PyClassInheritorsSearch.search(pyClass, true)

src/com/koxudaxi/pydantic/PydanticFieldSearchExecutor.kt

@@ -36,7 +36,6 @@ private fun searchDirectReferenceField(pyClass: PyClass, elementName: String, co
     if (searchField(pyClass, elementName, consumer)) return true
 
     return pyClass.getAncestorClasses(null)
-            .filterNot { isPydanticBaseModel(it) }
             .firstOrNull { isPydanticModel(it) && searchDirectReferenceField(it, elementName, consumer) } != null
 }
 

testData/completion/instance.py

@@ -0,0 +1,13 @@
+from builtins import *
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    abc: str
+    cde: str = str('abc')
+    efg: str = str('abc')
+
+class B(A):
+    hij: str
+
+A().<caret>

testData/completion/instanceBroken.py

@@ -0,0 +1,13 @@
+from builtins import *
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    abc: str a
+    cde: str = s
+    efg:
+
+class B(A):
+    hij: str
+
+A().<caret>

testData/completion/instanceParent.py

@@ -0,0 +1,12 @@
+from builtins import *
+
+from pydantic import BaseModel
+
+class A(BaseModel):
+    abc: str
+    cde: str
+
+class B(A):
+    efg: str
+
+B().<caret>

testData/completion/instanceParentIsPythonClass.py

@@ -0,0 +1,14 @@
+from builtins import *
+
+from pydantic import BaseModel
+
+
+class A:
+    abc: str
+    cde: str
+    efg: str
+
+class B(BaseModel, A):
+    hij: str
+
+B().<caret>

testData/completion/instancePythonClass.py

@@ -0,0 +1,14 @@
+from builtins import *
+
+from pydantic import BaseModel
+
+
+class A:
+    abc: str
+    cde: str
+    efg: str
+
+class B(A):
+    hij: str
+
+A().<caret>

testData/completion/instancePythonFunction.py

@@ -0,0 +1,9 @@
+from builtins import *
+
+from pydantic import BaseModel
+
+
+def a():
+    pass
+
+a().<caret>

testData/completion/instanceUnResolve.py

@@ -0,0 +1,6 @@
+from builtins import *
+from pydantic import BaseModel
+
+a = b
+
+a.<caret>

testData/completion/keywordArgument.py

@@ -1,3 +1,5 @@
+from builtins import *
+
 from pydantic import BaseModel
 
 

testData/completion/keywordArgumentAssignValue.py

@@ -1,3 +1,5 @@
+from builtins import *
+
 from pydantic import BaseModel
 
 

testData/completion/keywordArgumentInserted.py

@@ -1,3 +1,5 @@
+from builtins import *
+
 from pydantic import BaseModel
 
 

testData/completion/keywordArgumentParent.py

@@ -1,3 +1,5 @@
+from builtins import *
+
 from pydantic import BaseModel
 
 class A(BaseModel):

testData/completion/keywordArgumentPythonClass.py

@@ -1,3 +1,5 @@
+from builtins import *
+
 from pydantic import BaseModel
 
 

testData/completion/keywordArgumentPythonFunction.py

@@ -1,3 +1,5 @@
+from builtins import *
+
 from pydantic import BaseModel