use optimized typeEvalContext

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

Hide file tree

Showing 4 changed files with 7 additions and 19 deletions.
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) }