use optimized typeEvalContext

fix to handle models which have __init__ or __new__ methods
koxudaxi · Sep 19, 2019 · 7daaaa3 · 7daaaa3
1 parent f168c61
commit 7daaaa3
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 20 deletions.
diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
@@ -1,9 +1,15 @@
 <idea-plugin url="https://github.com/koxudaxi/pydantic-pycharm-plugin">
     <id>com.koxudaxi.pydantic</id>
     <name>Pydantic</name>
-    <version>0.0.19</version>
+    <version>0.0.20</version>
     <vendor email="koaxudai@gmail.com">Koudai Aono @koxudaxi</vendor>
     <change-notes><![CDATA[
+    <h2>version 0.0.20</h2>
+    <p>Features, BugFixes</p>
+    <ul>
+        <li>Support all features by parameters [#67] </li>
+        <li>Fix to handle models which have __init__ or __new__ methods [#67] </li>
+    </ul>
     <h2>version 0.0.19</h2>
     <p>BugFixes</p>
     <ul>

diff --git a/src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt b/src/com/koxudaxi/pydantic/PydanticCompletionContributor.kt
@@ -55,12 +55,6 @@ class PydanticCompletionContributor : CompletionContributor() {
             return "${typeHint}$defaultValue ${pyClass.name}"
         }
 
-        protected fun getPydanticPyClassTypesFromPyNamedParameter(pyNamedParameter: PyNamedParameter, typeEvalContext: TypeEvalContext): PyClassType? {
-            return pyNamedParameter.getArgumentType(typeEvalContext)?.let {
-                getPyClassTypeByPyTypes(it)
-            }?.firstOrNull { it1 -> isPydanticModel(it1.pyClass) }
-        }
-
 
         private fun addFieldElement(pyClass: PyClass, results: LinkedHashMap<String, LookupElement>,
                                     typeEvalContext: TypeEvalContext,

diff --git a/src/com/koxudaxi/pydantic/PydanticFieldRenameFactory.kt b/src/com/koxudaxi/pydantic/PydanticFieldRenameFactory.kt
@@ -60,7 +60,7 @@ class PydanticFieldRenameFactory : AutomaticRenamerFactory {
                     }
                 is PyKeywordArgument ->
                     element.name?.let { name ->
-                        getPyClassByPyKeywordArgument(element, TypeEvalContext.codeAnalysis(element.project, element.containingFile))
+                        getPyClassByPyKeywordArgument(element, TypeEvalContext.userInitiated(element.project, element.containingFile))
                                 ?.let { pyClass ->
                                     addAllElement(pyClass, name, added)
                                 }

diff --git a/src/com/koxudaxi/pydantic/PydanticFieldSearchExecutor.kt b/src/com/koxudaxi/pydantic/PydanticFieldSearchExecutor.kt
@@ -18,7 +18,7 @@ class PydanticFieldSearchExecutor : QueryExecutorBase<PsiReference, ReferencesSe
             is PyKeywordArgument -> run<RuntimeException> {
                 element.name
                         ?.let { elementName ->
-                            getPyClassByPyKeywordArgument(element, TypeEvalContext.deepCodeInsight(element.project))
+                            getPyClassByPyKeywordArgument(element, TypeEvalContext.userInitiated(element.project, element.containingFile))
                                     ?.takeIf { pyClass -> isPydanticModel(pyClass) }
                                     ?.let { pyClass -> searchDirectReferenceField(pyClass, elementName, consumer) }
                         }
@@ -54,7 +54,9 @@ class PydanticFieldSearchExecutor : QueryExecutorBase<PsiReference, ReferencesSe
 
             PsiTreeUtil.getParentOfType(psiReference.element, PyNamedParameter::class.java)
                     ?.let { param ->
-                        param.getArgumentType(TypeEvalContext.deepCodeInsight(psiReference.element.project))
+                        param.getArgumentType(TypeEvalContext.userInitiated(
+                                psiReference.element.project,
+                                psiReference.element.containingFile))
                                 ?.let { pyType ->
                                     getPyClassTypeByPyTypes(pyType)
                                             .firstOrNull { pyClassType -> isPydanticModel(pyClassType.pyClass) }
@@ -64,9 +66,7 @@ class PydanticFieldSearchExecutor : QueryExecutorBase<PsiReference, ReferencesSe
                                                 }
                                             }
                                 }
-
                     }
-
         }
 
     }

diff --git a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
@@ -99,20 +99,14 @@ class PydanticTypeProvider : PyTypeProviderBase() {
         if (!isPydanticModel(pyClass, context)) return null
         val clsType = (context.getType(pyClass) as? PyClassLikeType) ?: return null
         val ellipsis = PyElementGenerator.getInstance(pyClass.project).createEllipsis()
-        val resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context)
 
         val collected = linkedMapOf<String, PyCallableParameter>()
 
         for (currentType in StreamEx.of(clsType).append(pyClass.getAncestorTypes(context))) {
-            if (currentType == null ||
-                    !currentType.resolveMember(PyNames.INIT, null, AccessDirection.READ, resolveContext, false).isNullOrEmpty() ||
-                    !currentType.resolveMember(PyNames.NEW, null, AccessDirection.READ, resolveContext, false).isNullOrEmpty() ||
-                    currentType !is PyClassType) {
-                continue
-            }
+            if ( currentType !is PyClassType) continue
 
             val current = currentType.pyClass
-            if (!isPydanticModel(current, context)) return null
+            if (!isPydanticModel(current, context)) continue
 
             getClassVariables(current, context)
                     .mapNotNull { fieldToParameter(it, ellipsis, context, current) }

diff --git a/testData/completion/classFields.py b/testData/completion/classFields.py
@@ -0,0 +1,13 @@
+from builtins import *
+from pydantic import BaseModel
+
+
+class A(BaseModel):
+    abc: str
+    cde = str('abc')
+    efg: str = str('abc')
+
+class B(A):
+    hij: str
+
+B.<caret>
diff --git a/testSrc/com/koxudaxi/pydantic/PydanticCompletionTest.kt b/testSrc/com/koxudaxi/pydantic/PydanticCompletionTest.kt
@@ -306,6 +306,14 @@ open class PydanticCompletionTest : PydanticTestCase() {
         )
     }
 
+    fun testClassFields() {
+        doFieldTest(
+                listOf(
+                        Pair("___slots__", "BaseModel")
+                )
+        )
+    }
+
     fun testField() {
         doFieldTest(
                 listOf(