PY-40007 PY-41546 Add completion for keys of dict literals marked as TypedDicts

- in call expressions
- in assignment statements
- in return statements

GitOrigin-RevId: aea2063f7cd7ffc68903c79fa16a25fdf4ff0f8f

python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml

@@ -102,6 +102,8 @@
                             implementationClass="com.jetbrains.python.codeInsight.completion.PySuperClassAttributesCompletionContributor"/>
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyDictKeyNamesCompletionContributor"/>
+    <completion.contributor language="Python"
+                            implementationClass="com.jetbrains.python.codeInsight.completion.PyDictLiteralCompletionContributor"/>
     <completion.contributor language="Python"
                             implementationClass="com.jetbrains.python.codeInsight.completion.PyParameterCompletionContributor"/>
     <completion.contributor language="Python"

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictLiteralCompletionContributor.kt

@@ -0,0 +1,131 @@
+// Copyright 2000-2022 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.codeInsight.lookup.LookupElementBuilder
+import com.intellij.patterns.PlatformPatterns.psiElement
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.PlatformIcons
+import com.intellij.util.ProcessingContext
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.PyPsiUtils
+import com.jetbrains.python.psi.types.*
+
+/**
+ * Provides completion variants for keys of dict literals marked as TypedDict
+ * @see PyTypedDictType
+ */
+class PyDictLiteralCompletionContributor : CompletionContributor() {
+  init {
+    extend(CompletionType.BASIC, psiElement().inside(PySequenceExpression::class.java), DictLiteralCompletionProvider())
+  }
+}
+
+private class DictLiteralCompletionProvider : CompletionProvider<CompletionParameters?>() {
+  override fun addCompletions(parameters: CompletionParameters, context: ProcessingContext, result: CompletionResultSet) {
+    val originalElement = parameters.originalPosition?.parent ?: return
+
+    val possibleSequenceExpr = if (originalElement is PyStringLiteralExpression) originalElement.parent else originalElement
+    if (possibleSequenceExpr is PyDictLiteralExpression || possibleSequenceExpr is PySetLiteralExpression) { // it's set literal when user is typing the first key
+      addCompletionToCallExpression(originalElement, possibleSequenceExpr as PyTypedElement, result)
+      addCompletionToAssignment(originalElement, possibleSequenceExpr, result)
+      addCompletionToReturnStatement(originalElement, possibleSequenceExpr, result)
+    }
+  }
+
+  private fun addCompletionToCallExpression(originalElement: PsiElement,
+                                            possibleSequenceExpr: PyTypedElement,
+                                            result: CompletionResultSet) {
+    val callExpression = PsiTreeUtil.getParentOfType(originalElement, PyCallExpression::class.java)
+    if (callExpression != null) {
+      val typeEvalContext = TypeEvalContext.codeCompletion(originalElement.project, originalElement.containingFile)
+      val callType = typeEvalContext.getType(callExpression.callee ?: return)
+      if (callType !is PyCallableType) return
+
+      val argumentIndex = PyPsiUtils.findArgumentIndex(callExpression, possibleSequenceExpr)
+      if (argumentIndex < 0) return
+      val params = callType.getParameters(typeEvalContext) ?: return
+      if (params.size < argumentIndex) return
+      val expectedType = params[argumentIndex].getType(typeEvalContext)
+      val actualType = typeEvalContext.getType(possibleSequenceExpr)
+      addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+    }
+  }
+
+  private fun addCompletionToAssignment(originalElement: PsiElement,
+                                        possibleSequenceExpr: PyTypedElement,
+                                        result: CompletionResultSet) {
+    val assignment = PsiTreeUtil.getParentOfType(originalElement, PyAssignmentStatement::class.java)
+    if (assignment != null) {
+      val typeEvalContext = TypeEvalContext.codeCompletion(originalElement.project, originalElement.containingFile)
+      val targetToValue = if (assignment.targets.size == 1) assignment.targets[0] to assignment.assignedValue
+      else assignment.targetsToValuesMapping.firstOrNull { it.second == possibleSequenceExpr }?.let { it.first to it.second }
+
+      if (targetToValue != null) {
+        val expectedType = typeEvalContext.getType(targetToValue.first as PyTypedElement)
+        val actualType = typeEvalContext.getType(targetToValue.second as PyTypedElement)
+        addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+      }
+      else { //multiple target expressions and there is a PsiErrorElement
+        val targetExpr = assignment.assignedValue
+        val element = if (targetExpr is PyTupleExpression) targetExpr.elements.firstOrNull {
+          it == possibleSequenceExpr
+        }
+        else return
+        if (element != null) {
+          val index = targetExpr.elements.indexOf(element)
+          if (index < assignment.targets.size) {
+            val expectedType = typeEvalContext.getType(assignment.targets[index])
+            val actualType = typeEvalContext.getType(element)
+            addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+          }
+        }
+      }
+    }
+  }
+
+  private fun addCompletionToReturnStatement(originalElement: PsiElement,
+                                             possibleSequenceExpr: PyTypedElement,
+                                             result: CompletionResultSet) {
+    val returnStatement = PsiTreeUtil.getParentOfType(originalElement, PyReturnStatement::class.java)
+    if (returnStatement != null) {
+      val typeEvalContext = TypeEvalContext.codeCompletion(originalElement.project, originalElement.containingFile)
+      val owner = ScopeUtil.getScopeOwner(returnStatement)
+      if (owner is PyFunction) {
+        val annotation = owner.annotation
+        val typeCommentAnnotation = owner.typeCommentAnnotation
+        if (annotation != null || typeCommentAnnotation != null) { // to ensure that we have return type specified, not inferred
+          val expectedType = typeEvalContext.getReturnType(owner)
+          val actualType = typeEvalContext.getType(possibleSequenceExpr)
+          addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+        }
+      }
+    }
+  }
+
+  private fun addCompletionForTypedDictKeys(expectedType: PyType?,
+                                            actualType: PyType?,
+                                            dictCompletion: CompletionResultSet,
+                                            addQuotes: Boolean) {
+    if (expectedType is PyTypedDictType) {
+      val keys =
+        when {
+          actualType is PyTypedDictType -> expectedType.fields.keys.filterNot { it in actualType.fields }
+          actualType is PyClassType && PyNames.SET == actualType.name -> expectedType.fields.keys
+          else -> return
+        }
+
+      for (key in keys) {
+        dictCompletion.addElement(
+          LookupElementBuilder
+            .create(if (addQuotes) "'$key'" else key)
+            .withTypeText("dict key")
+            .withIcon(PlatformIcons.PARAMETER_ICON)
+        )
+      }
+    }
+  }
+}

python/testData/completion/dictLiteralCompletion/EmptyLiteralsInAssignments/main.py

@@ -0,0 +1,9 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+p: Point = {<caret>}

python/testData/completion/dictLiteralCompletion/EmptyLiteralsInAssignmentsWithQuotes/main.py

@@ -0,0 +1,9 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+p: Point = {"<caret>"}

python/testData/completion/dictLiteralCompletion/EmptyLiteralsInCallExpressions/main.py

@@ -0,0 +1,13 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def is_even(x: Point) -> bool:
+    pass
+
+
+is_even({<caret>})

python/testData/completion/dictLiteralCompletion/EmptyLiteralsInCallExpressionsWithQuotes/main.py

@@ -0,0 +1,13 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def is_even(x: Point) -> bool:
+    pass
+
+
+is_even({'<caret>'})

python/testData/completion/dictLiteralCompletion/EmptyLiteralsInReturnStatements/main.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def create_point() -> Point:
+    return {<caret>}

python/testData/completion/dictLiteralCompletion/EmptyLiteralsInReturnStatementsWithQuotes/main.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def create_another_point() -> Point:
+    return {'<caret>'}

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInAssignments/main.py

@@ -0,0 +1,9 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+p: Point = {"x": 42, <caret>}

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInAssignmentsWithMultipleTargets/main.py

@@ -0,0 +1,16 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+class NotPoint(TypedDict):
+    a: str
+    b: str
+
+
+p1: Point
+p2: NotPoint
+p1, p2 = {}, {'a': '42', '<caret>'}

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInAssignmentsWithQuotes/main.py

@@ -0,0 +1,9 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+p: Point = {"x": 42, "<caret>"}

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInCallExpressions/main.py

@@ -0,0 +1,13 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def is_even(x: Point) -> bool:
+    pass
+
+
+is_even({'x': 42, <caret>})

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInCallExpressionsWithQuotes/main.py

@@ -0,0 +1,13 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def is_even(x: Point) -> bool:
+    pass
+
+
+is_even({'x': 42, "<caret>"})

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInReturnStatements/main.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def create_point() -> Point:
+    return {'y': 42, <caret>}

python/testData/completion/dictLiteralCompletion/NotEmptyLiteralsInReturnStatementsWithQuotes/main.py

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    x: int
+    y: int
+
+
+def create_another_point() -> Point:
+    return {'y': 42, '<caret>'}

python/testSrc/com/jetbrains/python/PyDictLiteralCompletionTest.kt

@@ -0,0 +1,74 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python
+
+import com.jetbrains.python.fixtures.PyTestCase
+import org.junit.Assert
+import java.util.*
+
+
+class PyDictLiteralCompletionTest : PyTestCase() {
+  fun testEmptyLiteralsInCallExpressions() {
+    assertCompletionContains("'x'", "'y'")
+  }
+
+  fun testEmptyLiteralsInCallExpressionsWithQuotes() {
+    assertCompletionContains("x", "y")
+  }
+
+  fun testNotEmptyLiteralsInCallExpressions() {
+    assertCompletionContains("'y'")
+  }
+
+  fun testNotEmptyLiteralsInCallExpressionsWithQuotes() {
+    assertCompletionContains("y")
+  }
+
+  fun testEmptyLiteralsInAssignments() {
+    assertCompletionContains("'x'", "'y'")
+  }
+
+  fun testEmptyLiteralsInAssignmentsWithQuotes() {
+    assertCompletionContains("x", "y")
+  }
+
+  fun testNotEmptyLiteralsInAssignments() {
+    assertCompletionContains("'y'")
+  }
+
+  fun testNotEmptyLiteralsInAssignmentsWithQuotes() {
+    assertCompletionContains("y")
+  }
+
+  fun testEmptyLiteralsInReturnStatements() {
+    assertCompletionContains("'x'", "'y'")
+  }
+
+  fun testEmptyLiteralsInReturnStatementsWithQuotes() {
+    assertCompletionContains("x", "y")
+  }
+
+  fun testNotEmptyLiteralsInReturnStatements() {
+    assertCompletionContains("'x'")
+  }
+
+  fun testNotEmptyLiteralsInReturnStatementsWithQuotes() {
+    assertCompletionContains("x")
+  }
+
+  fun testNotEmptyLiteralsInAssignmentsWithMultipleTargets() {
+    assertCompletionContains("b")
+  }
+
+  private fun assertCompletionContains(vararg expected: String) {
+    myFixture.copyDirectoryToProject(getTestName(false), "")
+    myFixture.configureByFile("main.py")
+    myFixture.completeBasic()
+    val variants = myFixture.lookupElementStrings ?: emptyList()
+    assertContainsElements(variants, *expected)
+  }
+
+
+  override fun getTestDataPath(): String {
+    return super.getTestDataPath() + "/completion/dictLiteralCompletion/"
+  }
+}