diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt index e1a42b1ceafb..5b80ceba4265 100644 --- a/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt @@ -8,6 +8,8 @@ import com.intellij.patterns.PlatformPatterns import com.intellij.util.ProcessingContext import com.jetbrains.extensions.python.afterDefInMethod import com.jetbrains.python.PyNames +import com.jetbrains.python.codeInsight.stdlib.DUNDER_POST_INIT +import com.jetbrains.python.codeInsight.stdlib.DATACLASSES_INITVAR_TYPE import com.jetbrains.python.codeInsight.stdlib.parseDataclassParameters import com.jetbrains.python.psi.PySubscriptionExpression import com.jetbrains.python.psi.PyTargetExpression @@ -38,7 +40,7 @@ class PyDataclassPostInitCompletionContributor : CompletionContributor() { if (name != null && annotationValue != null) { val type = typeEvalContext.getType(element) - if (type is PyClassType && type.classQName == "dataclasses.InitVar") { + if (type is PyClassType && type.classQName == DATACLASSES_INITVAR_TYPE) { val typeHint = annotationValue.indexExpression.let { if (it == null) "" else ": ${it.text}" } postInitParameters.add(name + typeHint) } @@ -48,7 +50,7 @@ class PyDataclassPostInitCompletionContributor : CompletionContributor() { true } - addMethodToResult(result, cls, typeEvalContext, "__post_init__", postInitParameters.joinToString(prefix = "(", postfix = ")")) + addMethodToResult(result, cls, typeEvalContext, DUNDER_POST_INIT, postInitParameters.joinToString(prefix = "(", postfix = ")")) } } } diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt index 3222fc4b6e03..740a581814f2 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassUtil.kt @@ -9,6 +9,10 @@ import com.jetbrains.python.psi.resolve.PyResolveContext import com.jetbrains.python.psi.types.TypeEvalContext +const val DATACLASSES_INITVAR_TYPE = "dataclasses.InitVar" +const val DUNDER_POST_INIT = "__post_init__" + + fun parseDataclassParameters(cls: PyClass, context: TypeEvalContext): DataclassParameters? { val decorators = cls.decoratorList ?: return null diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt index dff9501620f5..22ca205023ff 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyDataclassesTypeProvider.kt @@ -17,7 +17,7 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() { } override fun getParameterType(param: PyNamedParameter, func: PyFunction, context: TypeEvalContext): Ref? { - if (!param.isPositionalContainer && !param.isKeywordContainer && param.annotationValue == null && func.name == "__post_init__") { + if (!param.isPositionalContainer && !param.isKeywordContainer && param.annotationValue == null && func.name == DUNDER_POST_INIT) { val cls = func.containingClass if (cls != null && parseDataclassParameters(cls, context)?.init == true) { @@ -78,7 +78,7 @@ class PyDataclassesTypeProvider : PyTypeProviderBase() { private fun getTypeForParameter(element: PyTargetExpression, context: TypeEvalContext): PyType? { val type = context.getType(element) - if (type is PyCollectionType && type is PyClassType && type.classQName == "dataclasses.InitVar") { + if (type is PyCollectionType && type is PyClassType && type.classQName == DATACLASSES_INITVAR_TYPE) { return type.elementTypes.firstOrNull() } return type