diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e5ad75a..729e2ae0 100644 --- a/CHANGELOG.md +++ b/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)] diff --git a/src/com/koxudaxi/pydantic/Pydantic.kt b/src/com/koxudaxi/pydantic/Pydantic.kt index 9ef8e112..81e0d347 100644 --- a/src/com/koxudaxi/pydantic/Pydantic.kt +++ b/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().any { pyFunction -> pyFunction.isPydanticField || pyFunction.isConStr || pyFunction.isCustomModelField } + val context = TypeEvalContext.userInitiated(pyCallExpression.project, pyCallExpression.containingFile) + return pyCallExpression.multiResolveCalleeFunction(PyResolveContext.defaultContext(context)).filterIsInstance() + .any { pyFunction -> pyFunction.isPydanticField || pyFunction.isConStr || pyFunction.isCustomModelField } } internal fun getClassVariables(pyClass: PyClass, context: TypeEvalContext): Sequence { @@ -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 diff --git a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt b/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt index 4532ff36..f46547e7 100644 --- a/src/com/koxudaxi/pydantic/PydanticTypeProvider.kt +++ b/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() .map { it.takeIf { pyFunction -> pyFunction.isPydanticCreateModel } }.firstOrNull()