Add completion for __post_init__ method (PY-27398)

python/src/META-INF/python-core-common.xml

@@ -85,6 +85,8 @@
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyMetaClassCompletionContributor"/>
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyStructuralTypeAttributesCompletionContributor"/>
+    <completion.contributor language="Python"
+                            implementationClass="com.jetbrains.python.codeInsight.completion.PyDataclassPostInitCompletionContributor"/>
 
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyStringFormatCompletionContributor" />

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<CompletionParameters>() {
+
+    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 = ")"))
+      }
+    }
+  }
+}

python/testData/completion/dataclassPostInit/a.after.py

@@ -0,0 +1,7 @@
+import dataclasses
+
+@dataclasses.dataclass
+class A1:
+    x: int = 0
+
+    def __post_init__(self):

python/testData/completion/dataclassPostInit/a.py

@@ -0,0 +1,7 @@
+import dataclasses
+
+@dataclasses.dataclass
+class A1:
+    x: int = 0
+
+    def __post<caret>

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

python/testData/completion/dataclassPostInitNoInit/a.after.py

@@ -0,0 +1,7 @@
+import dataclasses
+
+@dataclasses.dataclass(init=False)
+class A1:
+    x: int = 0
+
+    def __post

python/testData/completion/dataclassPostInitNoInit/a.py

@@ -0,0 +1,7 @@
+import dataclasses
+
+@dataclasses.dataclass(init=False)
+class A1:
+    x: int = 0
+
+    def __post<caret>

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

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):

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<caret>

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

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";