PY-58374 Fix dict keys completion when editing an existing key

GitOrigin-RevId: 41434300f6ea0f776a115a869fb9578a84015471

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

@@ -31,7 +31,10 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
   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
+    var possibleSequenceExpr = if (originalElement is PyStringLiteralExpression) originalElement.parent else originalElement
+    if (possibleSequenceExpr is PyKeyValueExpression) {
+      possibleSequenceExpr = possibleSequenceExpr.parent
+    }
     if (possibleSequenceExpr is PyDictLiteralExpression || possibleSequenceExpr is PySetLiteralExpression) { // it's set literal when user is typing the first key
       addCompletionToCallExpression(originalElement, possibleSequenceExpr as PySequenceExpression, result)
       addCompletionToAssignment(originalElement, possibleSequenceExpr, result)

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

@@ -0,0 +1,10 @@
+from typing import TypedDict
+
+
+class Point(TypedDict):
+    coordinateX: int
+    coordinateY: int
+    coordinateZ: int
+
+
+p: Point = {'coo<caret>': 42, 'coordinateY': 42}

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

@@ -67,12 +67,23 @@ class PyDictLiteralCompletionTest : PyTestCase() {
     assertCompletionContains("'x'")
   }
 
+  // PY-58374
+  fun testEditingKeyWithExistingValue() {
+    val variants = getVariants()
+    assertContainsElements(variants, "coordinateX", "coordinateZ")
+    assertDoesntContain(variants, "coordinateY")
+  }
+
   private fun assertCompletionContains(vararg expected: String) {
+    val variants = getVariants()
+    assertContainsElements(variants, *expected)
+  }
+
+  private fun getVariants(): List<String> {
     myFixture.copyDirectoryToProject(getTestName(false), "")
     myFixture.configureByFile("main.py")
     myFixture.completeBasic()
-    val variants = myFixture.lookupElementStrings ?: emptyList()
-    assertContainsElements(variants, *expected)
+    return myFixture.lookupElementStrings ?: emptyList()
   }