diff --git a/python/src/META-INF/python-core-common.xml b/python/src/META-INF/python-core-common.xml index fcbcd2f6d7dc..3801dc04a336 100644 --- a/python/src/META-INF/python-core-common.xml +++ b/python/src/META-INF/python-core-common.xml @@ -85,6 +85,8 @@ implementationClass="com.jetbrains.python.codeInsight.completion.PyMetaClassCompletionContributor"/> + diff --git a/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt b/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt new file mode 100644 index 000000000000..e1a42b1ceafb --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/completion/PyDataclassPostInitCompletionContributor.kt @@ -0,0 +1,55 @@ +/* + * Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. + */ +package com.jetbrains.python.codeInsight.completion + +import com.intellij.codeInsight.completion.* +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.parseDataclassParameters +import com.jetbrains.python.psi.PySubscriptionExpression +import com.jetbrains.python.psi.PyTargetExpression +import com.jetbrains.python.psi.types.PyClassType + +class PyDataclassPostInitCompletionContributor : CompletionContributor() { + + override fun handleAutoCompletionPossibility(context: AutoCompletionContext) = autoInsertSingleItem(context) + + init { + extend(CompletionType.BASIC, PlatformPatterns.psiElement().afterDefInMethod(), MyCompletionProvider) + } + + private object MyCompletionProvider : CompletionProvider() { + + override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext?, result: CompletionResultSet) { + val cls = parameters.getPyClass() ?: return + val typeEvalContext = parameters.getTypeEvalContext() + + if (parseDataclassParameters(cls, typeEvalContext)?.init == true) { + val postInitParameters = mutableListOf(PyNames.CANONICAL_SELF) + + cls.processClassLevelDeclarations { element, _ -> + if (element is PyTargetExpression && element.annotationValue != null) { + val name = element.name + val annotationValue = element.annotation?.value as? PySubscriptionExpression + + if (name != null && annotationValue != null) { + val type = typeEvalContext.getType(element) + + if (type is PyClassType && type.classQName == "dataclasses.InitVar") { + val typeHint = annotationValue.indexExpression.let { if (it == null) "" else ": ${it.text}" } + postInitParameters.add(name + typeHint) + } + } + } + + true + } + + addMethodToResult(result, cls, typeEvalContext, "__post_init__", postInitParameters.joinToString(prefix = "(", postfix = ")")) + } + } + } +} \ No newline at end of file diff --git a/python/testData/completion/dataclassPostInit/a.after.py b/python/testData/completion/dataclassPostInit/a.after.py new file mode 100644 index 000000000000..4c5d2378067d --- /dev/null +++ b/python/testData/completion/dataclassPostInit/a.after.py @@ -0,0 +1,7 @@ +import dataclasses + +@dataclasses.dataclass +class A1: + x: int = 0 + + def __post_init__(self): \ No newline at end of file diff --git a/python/testData/completion/dataclassPostInit/a.py b/python/testData/completion/dataclassPostInit/a.py new file mode 100644 index 000000000000..a4c8f8176de8 --- /dev/null +++ b/python/testData/completion/dataclassPostInit/a.py @@ -0,0 +1,7 @@ +import dataclasses + +@dataclasses.dataclass +class A1: + x: int = 0 + + def __post \ No newline at end of file diff --git a/python/testData/completion/dataclassPostInit/dataclasses.py b/python/testData/completion/dataclassPostInit/dataclasses.py new file mode 100644 index 000000000000..3a2b30d36707 --- /dev/null +++ b/python/testData/completion/dataclassPostInit/dataclasses.py @@ -0,0 +1,3 @@ +def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False, + hash=None, frozen=False): + pass \ No newline at end of file diff --git a/python/testData/completion/dataclassPostInitNoInit/a.after.py b/python/testData/completion/dataclassPostInitNoInit/a.after.py new file mode 100644 index 000000000000..0d50f3dba181 --- /dev/null +++ b/python/testData/completion/dataclassPostInitNoInit/a.after.py @@ -0,0 +1,7 @@ +import dataclasses + +@dataclasses.dataclass(init=False) +class A1: + x: int = 0 + + def __post \ No newline at end of file diff --git a/python/testData/completion/dataclassPostInitNoInit/a.py b/python/testData/completion/dataclassPostInitNoInit/a.py new file mode 100644 index 000000000000..b5e2b0efe55f --- /dev/null +++ b/python/testData/completion/dataclassPostInitNoInit/a.py @@ -0,0 +1,7 @@ +import dataclasses + +@dataclasses.dataclass(init=False) +class A1: + x: int = 0 + + def __post \ No newline at end of file diff --git a/python/testData/completion/dataclassPostInitNoInit/dataclasses.py b/python/testData/completion/dataclassPostInitNoInit/dataclasses.py new file mode 100644 index 000000000000..3a2b30d36707 --- /dev/null +++ b/python/testData/completion/dataclassPostInitNoInit/dataclasses.py @@ -0,0 +1,3 @@ +def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False, + hash=None, frozen=False): + pass \ No newline at end of file diff --git a/python/testData/completion/dataclassWithInitVarPostInit/a.after.py b/python/testData/completion/dataclassWithInitVarPostInit/a.after.py new file mode 100644 index 000000000000..82dda47f86ad --- /dev/null +++ b/python/testData/completion/dataclassWithInitVarPostInit/a.after.py @@ -0,0 +1,8 @@ +import dataclasses + +@dataclasses.dataclass +class A1: + x: int = 0 + y: dataclasses.InitVar[int] = 1 + + def __post_init__(self, y: int): \ No newline at end of file diff --git a/python/testData/completion/dataclassWithInitVarPostInit/a.py b/python/testData/completion/dataclassWithInitVarPostInit/a.py new file mode 100644 index 000000000000..42fad562550f --- /dev/null +++ b/python/testData/completion/dataclassWithInitVarPostInit/a.py @@ -0,0 +1,8 @@ +import dataclasses + +@dataclasses.dataclass +class A1: + x: int = 0 + y: dataclasses.InitVar[int] = 1 + + def __post \ No newline at end of file diff --git a/python/testData/completion/dataclassWithInitVarPostInit/dataclasses.py b/python/testData/completion/dataclassWithInitVarPostInit/dataclasses.py new file mode 100644 index 000000000000..e2c6ed38f1c6 --- /dev/null +++ b/python/testData/completion/dataclassWithInitVarPostInit/dataclasses.py @@ -0,0 +1,11 @@ +class _InitVarMeta(type): + def __getitem__(self, params): + return self + +class InitVar(metaclass=_InitVarMeta): + pass + + +def dataclass(_cls=None, *, init=True, repr=True, eq=True, order=False, + hash=None, frozen=False): + pass \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java index 97e06bc29419..975e55b1bf55 100644 --- a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java +++ b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java @@ -254,6 +254,21 @@ public class Py3CompletionTest extends PyTestCase { assertDoesntContain(suggested, "meta_method"); } + // PY-27398 + public void testDataclassPostInit() { + runWithLanguageLevel(LanguageLevel.PYTHON37, this::doMultiFileTest); + } + + // PY-27398 + public void testDataclassWithInitVarPostInit() { + runWithLanguageLevel(LanguageLevel.PYTHON37, this::doMultiFileTest); + } + + // PY-27398 + public void testDataclassPostInitNoInit() { + runWithLanguageLevel(LanguageLevel.PYTHON37, this::doMultiFileTest); + } + @Override protected String getTestDataPath() { return super.getTestDataPath() + "/completion";