Merge pull request #661 from koxudaxi/use_multiResolveCalleeFunction_…

…instead_of_getResolvedPsiElements

Use multiResolveCalleeFunction instead of getResolvedPsiElements

CHANGELOG.md

@@ -6,6 +6,7 @@
 - Optimize resolving pydantic class [[#658](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/658)]
 - Improve dynamic model field detection [[#659](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/659)]
 - Improve test coverage [[#660](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/660)]
+- Use multiResolveCalleeFunction instead of getResolvedPsiElements [[#661](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/661)]
 
 ## 0.3.17 - 2022-12-16
 - Support Union operator [[#602](https://github.com/koxudaxi/pydantic-pycharm-plugin/pull/602)]

src/com/koxudaxi/pydantic/Pydantic.kt

@@ -223,10 +223,9 @@ internal fun isPydanticRegex(stringLiteralExpression: StringLiteralExpression):
     val pyKeywordArgument = stringLiteralExpression.parent as? PyKeywordArgument ?: return false
     if (pyKeywordArgument.keyword != "regex") return false
     val pyCallExpression = pyKeywordArgument.parent.parent as? PyCallExpression ?: return false
-    val referenceExpression = pyCallExpression.callee as? PyReferenceExpression ?: return false
-    val context = TypeEvalContext.userInitiated(referenceExpression.project, referenceExpression.containingFile)
-    return getResolvedPsiElements(referenceExpression, context)
-        .filterIsInstance<PyFunction>().any { pyFunction -> pyFunction.isPydanticField || pyFunction.isConStr || pyFunction.isCustomModelField }
+    val context = TypeEvalContext.userInitiated(pyCallExpression.project, pyCallExpression.containingFile)
+    return pyCallExpression.multiResolveCalleeFunction(PyResolveContext.defaultContext(context)).filterIsInstance<PyFunction>()
+        .any { pyFunction -> pyFunction.isPydanticField || pyFunction.isConStr || pyFunction.isCustomModelField }
 }
 
 internal fun getClassVariables(pyClass: PyClass, context: TypeEvalContext): Sequence<PyTargetExpression> {
@@ -586,11 +585,10 @@ internal fun getFieldFromPyExpression(
     context: TypeEvalContext,
     pydanticVersion: KotlinVersion?,
 ): PyCallExpression? {
-    val callee = (psiElement as? PyCallExpression)
-        ?.let { it.callee as? PyReferenceExpression }
-        ?: return null
+    if (psiElement !is PyCallExpression) return null
     val versionZero = pydanticVersion?.major == 0
-    if (!getResolvedPsiElements(callee, context).any {
+    if (
+        !psiElement.multiResolveCalleeFunction(PyResolveContext.defaultContext(context)).any {
             when {
                 versionZero -> isPydanticSchemaByPsiElement(it, context)
                 else -> it.isPydanticField || it.isCustomModelField

src/com/koxudaxi/pydantic/PydanticTypeProvider.kt

@@ -9,6 +9,7 @@ import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.isBitwiseOrUnionAvailable
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.impl.*
+import com.jetbrains.python.psi.resolve.PyResolveContext
 import com.jetbrains.python.psi.types.*
 import com.koxudaxi.pydantic.PydanticConfigService.Companion.getInstance
 import one.util.streamex.StreamEx
@@ -262,8 +263,7 @@ class PydanticTypeProvider : PyTypeProviderBase() {
     ): PydanticDynamicModelClassType? {
         val arguments = pyCallExpression.arguments.toList()
         if (arguments.isEmpty()) return null
-        val referenceExpression = (pyCallExpression.callee as? PyReferenceExpression) ?: return null
-        val pyFunction = getResolvedPsiElements(referenceExpression, context)
+        val pyFunction = pyCallExpression.multiResolveCalleeFunction(PyResolveContext.defaultContext(context))
             .asSequence()
             .filterIsInstance<PyFunction>()
             .map { it.takeIf { pyFunction -> pyFunction.isPydanticCreateModel } }.firstOrNull()