From ff9f7eb1f9d4cb0aa83447016ba415e65554101a Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Mon, 10 May 2021 10:50:48 +0900
Subject: [PATCH 1/2] Improve resolving reference

---
 src/com/koxudaxi/pydantic/Pydantic.kt             | 12 ++----------
 src/com/koxudaxi/pydantic/PydanticTypeProvider.kt |  8 ++++----
 2 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/src/com/koxudaxi/pydantic/Pydantic.kt b/src/com/koxudaxi/pydantic/Pydantic.kt
index ca9dc608..6a8fd46f 100644
--- a/src/com/koxudaxi/pydantic/Pydantic.kt
+++ b/src/com/koxudaxi/pydantic/Pydantic.kt
@@ -5,7 +5,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.Sdk
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
-import com.intellij.psi.ResolveResult
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.QualifiedName
 import com.jetbrains.extensions.ModuleBasedContextAnchor
@@ -245,16 +244,9 @@ private fun getAliasedFieldName(
 }
 
 
-fun getResolveElements(referenceExpression: PyReferenceExpression, context: TypeEvalContext): Array<ResolveResult> {
-    return PyResolveContext.defaultContext().withTypeEvalContext(context).let {
-        referenceExpression.getReference(it).multiResolve(false)
-    }
-
-}
-
-
 fun getResolvedPsiElements(referenceExpression: PyReferenceExpression, context: TypeEvalContext): List<PsiElement> {
-    return getResolveElements(referenceExpression, context).let { PyUtil.filterTopPriorityResults(it) }
+    return PyUtil.multiResolveTopPriority(referenceExpression,
+        PyResolveContext.defaultContext().withTypeEvalContext(context))
 }
 
 val PyType.pyClassTypes: List<PyClassType>
diff --git a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
index ba244da0..70a8c739 100644
--- a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
+++ b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt
@@ -201,7 +201,7 @@ class PydanticTypeProvider : PyTypeProviderBase() {
 
         return getResolvedPsiElements(referenceExpression, context)
             .asSequence()
-            .map {
+            .mapNotNull {
                 when {
                     it is PyClass -> getPydanticTypeForClass(it, context, true, pyCallExpression)
                     it is PyParameter && it.isSelf ->
@@ -672,9 +672,9 @@ class PydanticTypeProvider : PyTypeProviderBase() {
         val callee = (assignedValue as? PyCallExpressionImpl)?.callee ?: return assignedValue
         val referenceExpression = callee.reference?.element as? PyReferenceExpression ?: return ellipsis
 
-        val resolveResults = getResolveElements(referenceExpression, context)
+        val resolveResults = getResolvedPsiElements(referenceExpression, context)
         if (isDataclass) {
-            PyUtil.filterTopPriorityResults(resolveResults)
+            resolveResults
                 .any {
                     it.isDataclassField
                 }
@@ -687,7 +687,7 @@ class PydanticTypeProvider : PyTypeProviderBase() {
         } else {
 
             val versionZero = pydanticVersion?.major == 0
-            PyUtil.filterTopPriorityResults(resolveResults)
+            resolveResults
                 .any {
                     when {
                         versionZero -> isPydanticSchemaByPsiElement(it, context)

From 0caf3a1fe34d385bc1c593338eb41341f4ebdd96 Mon Sep 17 00:00:00 2001
From: Koudai Aono <koxudaxi@gmail.com>
Date: Mon, 10 May 2021 10:52:42 +0900
Subject: [PATCH 2/2] Update history

---
 resources/META-INF/plugin.xml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/resources/META-INF/plugin.xml b/resources/META-INF/plugin.xml
index f05131d9..3d40a42d 100644
--- a/resources/META-INF/plugin.xml
+++ b/resources/META-INF/plugin.xml
@@ -7,6 +7,7 @@
     <h2>version 0.3.1</h2>
     <p>Features</p>
     <ul>
+        <li>Improve resolving reference [#293]</li>
         <li>Improve coding style [#292]</li>
         <li>Support GenericModel [#289]</li>
         <li>Support frozen on config [#288]</li>