PY-52502 Duplicate completion variants when local variable is used as map key

GitOrigin-RevId: 463225922b2f0201b5d1ed2749f4175a77fd543f

python/python-common-tests/com/jetbrains/python/PythonCommonCompletionTest.java

@@ -308,7 +308,11 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase {
   }
 
   public void testDictKeys() {  // PY-2245
-    doTest();
+    myFixture.configureByFile("dictKeys.py");
+    myFixture.completeBasic();
+    assertSameElements(myFixture.getLookupElementStrings(), "'xyz'");
+    myFixture.type('\n');
+    myFixture.checkResultByFile("dictKeys.after.py");
   }
 
   public void testDictKeys2() { //PY-4181
@@ -319,6 +323,11 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase {
     doTest();
   }
 
+  // PY-52502
+  public void testExpressionDictKey() {
+    doTest();
+  }
+
   // PY-42738
   public void testDictLiteralValueAccessWithDoubleQuotes() {
     final String text = """
@@ -340,7 +349,8 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase {
     final String text = """
     d={ "key1": 222, 'key2': 333, 22: True, False: 22 }
     d[<caret>]""";
-    assertContainsElements(doTestByText(text), "\"key1\"", "'key2'", "22", "False");
+    List<String> suggested = doTestByText(text);
+    assertContainsElements(suggested, "\"key1\"", "'key2'");
   }
 
   // PY-42738
@@ -400,7 +410,7 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase {
       d["xxx"]=25
       d['yyy']=26
       d[<caret>]""";
-    assertContainsElements(doTestByText(text), "30", "False", "\"xxx\"", "'yyy'");
+    assertContainsElements(doTestByText(text), "\"xxx\"", "'yyy'");
   }
 
   public void testNoParensForDecorator() {  // PY-2210

python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictKeyNamesCompletionContributor.java

@@ -122,15 +122,11 @@ public final class PyDictKeyNamesCompletionContributor extends CompletionContrib
   private static void addAdditionalKeys(final PsiFile file, final PsiElement operand, final DictKeyCompletionResultSet result) {
     Collection<PySubscriptionExpression> subscriptionExpressions = PsiTreeUtil.findChildrenOfType(file, PySubscriptionExpression.class);
     for (PySubscriptionExpression expr : subscriptionExpressions) {
-      if (expr.getOperand().getText().equals(operand.getText())) {
-        if (expr.getParent() instanceof PyAssignmentStatement assignmentStatement) {
-          if (expr.equals(assignmentStatement.getLeftHandSideExpression())) {
-            PyExpression key = expr.getIndexExpression();
-            if (key != null) {
-              result.addKey(key);
-            }
-          }
-        }
+      if (expr.getOperand().getText().equals(operand.getText()) &&
+          expr.getParent() instanceof PyAssignmentStatement assignmentStatement &&
+          expr.equals(assignmentStatement.getLeftHandSideExpression()) &&
+          expr.getIndexExpression() instanceof PyStringLiteralExpression key) {
+        result.addKey(key);
       }
     }
   }
@@ -141,7 +137,9 @@ public final class PyDictKeyNamesCompletionContributor extends CompletionContrib
   private static void addDictLiteralKeys(final PyDictLiteralExpression dict, final DictKeyCompletionResultSet result) {
     PyKeyValueExpression[] keyValues = dict.getElements();
     for (PyKeyValueExpression expression : keyValues) {
-      result.addKey(expression.getKey());
+      if (expression.getKey() instanceof PyStringLiteralExpression key) {
+        result.addKey(key);
+      }
     }
   }
 
@@ -154,11 +152,9 @@ public final class PyDictKeyNamesCompletionContributor extends CompletionContrib
       myIsInsideString = isInsideString;
     }
 
-    void addKey(@NotNull PyExpression keyExpression) {
+    void addKey(@NotNull PyStringLiteralExpression keyExpression) {
       if (myIsInsideString) {
-        if (keyExpression instanceof PyStringLiteralExpression stringLiteralExpression) {
-          addElement(stringLiteralExpression.getStringValue());
-        }
+        addElement(keyExpression.getStringValue());
       }
       else {
         addElement(keyExpression.getText());

python/testData/completion/dictKeys.after.py

@@ -1,2 +1,2 @@
-tasks = {key: "123"}
-tasks[key]
+tasks = {'xyz': "123"}
+tasks['xyz']

python/testData/completion/dictKeys.py

@@ -1,2 +1,2 @@
-tasks = {key: "123"}
-tasks[ke<caret>]
+tasks = {'xyz': "123"}
+tasks[xy<caret>]

python/testData/completion/expressionDictKey.after.py

@@ -0,0 +1,3 @@
+yyy='yyy'
+x = {yyy: -1}
+x[yyy]

python/testData/completion/expressionDictKey.py

@@ -0,0 +1,3 @@
+yyy='yyy'
+x = {yyy: -1}
+x[yy<caret>]