PY-42637 Fix quotes for TypedDict keys completion in dict literals

- when there are other keys present in the literal, use quotes from there
- when the literal is empty, use double quotes by default
- when there already are quotes, suggest just the key names

GitOrigin-RevId: 5f37cd25c0e47521e78a01ab84053571369b24a8

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

@@ -25,19 +25,22 @@ class PyDictLiteralCompletionContributor : CompletionContributor() {
 }
 
 private class DictLiteralCompletionProvider : CompletionProvider<CompletionParameters?>() {
+
+  private val DEFAULT_QUOTE = "\""
+
   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)
+      addCompletionToCallExpression(originalElement, possibleSequenceExpr as PySequenceExpression, result)
       addCompletionToAssignment(originalElement, possibleSequenceExpr, result)
       addCompletionToReturnStatement(originalElement, possibleSequenceExpr, result)
     }
   }
 
   private fun addCompletionToCallExpression(originalElement: PsiElement,
-                                            possibleSequenceExpr: PyTypedElement,
+                                            possibleSequenceExpr: PySequenceExpression,
                                             result: CompletionResultSet) {
     val callExpression = PsiTreeUtil.getParentOfType(originalElement, PyCallExpression::class.java)
     if (callExpression != null) {
@@ -51,12 +54,25 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
       if (params.size < argumentIndex) return
       val expectedType = params[argumentIndex].getType(typeEvalContext)
       val actualType = typeEvalContext.getType(possibleSequenceExpr)
-      addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+
+      addCompletionForTypedDictKeys(expectedType, actualType, result, getForcedQuote(possibleSequenceExpr, originalElement))
     }
   }
 
+  private fun getForcedQuote(possibleSequenceExpr: PySequenceExpression, originalElement: PsiElement): String {
+    if (originalElement !is PyStringLiteralExpression) {
+      val usedQuotes = possibleSequenceExpr.elements.mapNotNull { if (it is PyKeyValueExpression) it.key else it }
+        .filterIsInstance<PyStringLiteralExpression>()
+        .flatMap { it.stringElements }
+        .map { it.quote }
+        .toSet()
+      return usedQuotes.singleOrNull() ?: DEFAULT_QUOTE
+    }
+    return ""
+  }
+
   private fun addCompletionToAssignment(originalElement: PsiElement,
-                                        possibleSequenceExpr: PyTypedElement,
+                                        possibleSequenceExpr: PySequenceExpression,
                                         result: CompletionResultSet) {
     val assignment = PsiTreeUtil.getParentOfType(originalElement, PyAssignmentStatement::class.java)
     if (assignment != null) {
@@ -67,7 +83,7 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
       if (targetToValue?.first != null && targetToValue.second != null) {
         val expectedType = typeEvalContext.getType(targetToValue.first as PyTypedElement)
         val actualType = typeEvalContext.getType(targetToValue.second as PyTypedElement)
-        addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+        addCompletionForTypedDictKeys(expectedType, actualType, result, getForcedQuote(possibleSequenceExpr, originalElement))
       }
       else { //multiple target expressions and there is a PsiErrorElement
         val targetExpr = assignment.assignedValue
@@ -80,7 +96,7 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
           if (index < assignment.targets.size) {
             val expectedType = typeEvalContext.getType(assignment.targets[index])
             val actualType = typeEvalContext.getType(element)
-            addCompletionForTypedDictKeys(expectedType, actualType, result, originalElement !is PyStringLiteralExpression)
+            addCompletionForTypedDictKeys(expectedType, actualType, result, getForcedQuote(possibleSequenceExpr, originalElement))
           }
         }
       }
@@ -88,7 +104,7 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
   }
 
   private fun addCompletionToReturnStatement(originalElement: PsiElement,
-                                             possibleSequenceExpr: PyTypedElement,
+                                             possibleSequenceExpr: PySequenceExpression,
                                              result: CompletionResultSet) {
     val returnStatement = PsiTreeUtil.getParentOfType(originalElement, PyReturnStatement::class.java)
     if (returnStatement != null) {
@@ -100,7 +116,7 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
         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)
+          addCompletionForTypedDictKeys(expectedType, actualType, result, getForcedQuote(possibleSequenceExpr, originalElement))
         }
       }
     }
@@ -109,7 +125,7 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
   private fun addCompletionForTypedDictKeys(expectedType: PyType?,
                                             actualType: PyType?,
                                             dictCompletion: CompletionResultSet,
-                                            addQuotes: Boolean) {
+                                            quote: String) {
     if (expectedType is PyTypedDictType) {
       val keys =
         when {
@@ -121,7 +137,7 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
       for (key in keys) {
         dictCompletion.addElement(
           LookupElementBuilder
-            .create(if (addQuotes) "'$key'" else key)
+            .create("$quote$key$quote")
             .withTypeText("dict key")
             .withIcon(IconManager.getInstance().getPlatformIcon(com.intellij.ui.PlatformIcons.Parameter))
         )

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

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

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

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

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

@@ -2,13 +2,11 @@
 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'")
+    assertCompletionContains("\"x\"", "\"y\"")
   }
 
   fun testEmptyLiteralsInCallExpressionsWithQuotes() {
@@ -24,7 +22,7 @@ class PyDictLiteralCompletionTest : PyTestCase() {
   }
 
   fun testEmptyLiteralsInAssignments() {
-    assertCompletionContains("'x'", "'y'")
+    assertCompletionContains("\"x\"", "\"y\"")
   }
 
   fun testEmptyLiteralsInAssignmentsWithQuotes() {
@@ -32,7 +30,7 @@ class PyDictLiteralCompletionTest : PyTestCase() {
   }
 
   fun testNotEmptyLiteralsInAssignments() {
-    assertCompletionContains("'y'")
+    assertCompletionContains("\"y\"")
   }
 
   fun testNotEmptyLiteralsInAssignmentsWithQuotes() {
@@ -40,7 +38,7 @@ class PyDictLiteralCompletionTest : PyTestCase() {
   }
 
   fun testEmptyLiteralsInReturnStatements() {
-    assertCompletionContains("'x'", "'y'")
+    assertCompletionContains("\"x\"", "\"y\"")
   }
 
   fun testEmptyLiteralsInReturnStatementsWithQuotes() {
@@ -59,6 +57,16 @@ class PyDictLiteralCompletionTest : PyTestCase() {
     assertCompletionContains("b")
   }
 
+  // PY-42637
+  fun testNotEmptyLiteralsInReturnStatementsWithSeveralTypesOfQuotes() {
+    assertCompletionContains("\"x\"")
+  }
+
+  // PY-42637
+  fun testNotEmptyLiteralsInReturnStatementsWithOnlySingleQuotes() {
+    assertCompletionContains("'x'")
+  }
+
   private fun assertCompletionContains(vararg expected: String) {
     myFixture.copyDirectoryToProject(getTestName(false), "")
     myFixture.configureByFile("main.py")