From aefaeaa52a4cacda546d7253b9ceb5dbb953e68b Mon Sep 17 00:00:00 2001
From: Petr <petr.golubev@jetbrains.com>
Date: Tue, 7 May 2024 13:46:19 +0200
Subject: [PATCH] PY-54151 TypedDict completion at callee does not work for
 methods

GitOrigin-RevId: 5a037490d8c24ccdd354e8396540a21294148a7e
---
 .../PyDictLiteralCompletionContributor.kt     | 27 ++++++++----------
 .../PyLiteralTypeCompletionContributor.kt     |  6 ++--
 .../psi/impl/PyCallExpressionHelper.java      | 28 +++++++++++++++++++
 .../EmptyLiteralsInMethodCall/main.py         | 14 ++++++++++
 .../main.py                                   | 14 ++++++++++
 .../python/PyDictLiteralCompletionTest.kt     |  8 ++++++
 6 files changed, 77 insertions(+), 20 deletions(-)
 create mode 100644 python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCall/main.py
 create mode 100644 python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCallWithKeywordArgument/main.py

diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictLiteralCompletionContributor.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictLiteralCompletionContributor.kt
index e9a17d91ca2c..e347ad0a9a79 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictLiteralCompletionContributor.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictLiteralCompletionContributor.kt
@@ -11,8 +11,12 @@ 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.*
+import com.jetbrains.python.psi.impl.PyCallExpressionHelper
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.PyClassType
+import com.jetbrains.python.psi.types.PyType
+import com.jetbrains.python.psi.types.PyTypedDictType
+import com.jetbrains.python.psi.types.TypeEvalContext
 
 /**
  * Provides completion variants for keys of dict literals marked as TypedDict
@@ -45,20 +49,11 @@ private class DictLiteralCompletionProvider : CompletionProvider<CompletionParam
   private fun addCompletionToCallExpression(originalElement: PsiElement,
                                             possibleSequenceExpr: PySequenceExpression,
                                             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, getForcedQuote(possibleSequenceExpr, originalElement))
+    val typeEvalContext = TypeEvalContext.codeCompletion(originalElement.project, originalElement.containingFile)
+    val actualType = typeEvalContext.getType(possibleSequenceExpr)
+    val quote = getForcedQuote(possibleSequenceExpr, originalElement)
+    PyCallExpressionHelper.getMappedParameters(possibleSequenceExpr, PyResolveContext.defaultContext(typeEvalContext)).forEach {
+      addCompletionForTypedDictKeys(it.getType(typeEvalContext), actualType, result, quote)
     }
   }
 
diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyLiteralTypeCompletionContributor.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyLiteralTypeCompletionContributor.kt
index 83fc576fbe73..a58be1346d95 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyLiteralTypeCompletionContributor.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyLiteralTypeCompletionContributor.kt
@@ -39,11 +39,9 @@ private class PyLiteralTypeCompletionProvider : CompletionProvider<CompletionPar
 
     val callSiteExpr = PsiTreeUtil.getParentOfType(position, PyCallSiteExpression::class.java)
     if (callSiteExpr != null) {
-      val parent = position.parent
-      val argumentExpr = if (parent is PyKeywordArgument && parent.valueExpression == position) parent else position
       val types = PyCallExpressionHelper
-        .mapArguments(callSiteExpr, PyResolveContext.defaultContext(typeEvalContext))
-        .mapNotNull { it.mappedParameters[argumentExpr]?.getArgumentType(typeEvalContext) }
+        .getMappedParameters(position, PyResolveContext.defaultContext(typeEvalContext))
+        .mapNotNull { it.getArgumentType(typeEvalContext) }
       addToResult(position, types, result)
       return
     }
diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
index b722c18d4bcf..241ed62008a2 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
@@ -26,6 +26,7 @@ import com.jetbrains.python.psi.types.*;
 import com.jetbrains.python.pyi.PyiUtil;
 import com.jetbrains.python.toolbox.Maybe;
 import one.util.streamex.StreamEx;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -750,6 +751,33 @@ public final class PyCallExpressionHelper {
     return null;
   }
 
+  /**
+   * {@code argument} can be (parenthesized) expression or a value of a {@link PyKeywordArgument}
+   */
+  @ApiStatus.Internal
+  @NotNull
+  public static List<PyCallableParameter> getMappedParameters(@NotNull PyExpression argument,
+                                                              @NotNull PyResolveContext resolveContext) {
+    while (argument.getParent() instanceof PyParenthesizedExpression parenthesizedExpr) {
+      argument = parenthesizedExpr;
+    }
+
+    if (argument.getParent() instanceof PyKeywordArgument keywordArgument && keywordArgument.getValueExpression() == argument) {
+      argument = keywordArgument;
+    }
+
+    PsiElement parent = argument.getParent();
+    if (parent instanceof PyArgumentList) {
+      parent = parent.getParent();
+    }
+    if (!(parent instanceof PyCallSiteExpression callSite)) {
+      return Collections.emptyList();
+    }
+
+    PyExpression finalArgument = argument;
+    return ContainerUtil.mapNotNull(mapArguments(callSite, resolveContext), mapping -> mapping.getMappedParameters().get(finalArgument));
+  }
+
   /**
    * Gets implicit offset from the {@code callableType},
    * should be used with the methods below since they specify correct offset value.
diff --git a/python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCall/main.py b/python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCall/main.py
new file mode 100644
index 000000000000..cbcb4a4cfdde
--- /dev/null
+++ b/python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCall/main.py
@@ -0,0 +1,14 @@
+from typing import TypedDict
+
+
+class Movie (TypedDict):
+    name: str
+    age: int
+
+
+class Foo:
+    def foo(self, movie: Movie) -> None:
+        pass
+
+
+Foo().foo(((({<caret>}))))
diff --git a/python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCallWithKeywordArgument/main.py b/python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCallWithKeywordArgument/main.py
new file mode 100644
index 000000000000..977e524cb104
--- /dev/null
+++ b/python/testData/completion/dictLiteralCompletion/EmptyLiteralsInMethodCallWithKeywordArgument/main.py
@@ -0,0 +1,14 @@
+from typing import TypedDict
+
+
+class Movie (TypedDict):
+    name: str
+    age: int
+
+
+class Foo:
+    def foo(self, movie: Movie) -> None:
+        pass
+
+
+Foo().foo(movie=((({<caret>}))))
diff --git a/python/testSrc/com/jetbrains/python/PyDictLiteralCompletionTest.kt b/python/testSrc/com/jetbrains/python/PyDictLiteralCompletionTest.kt
index 42ffdaf14c37..e863378feb69 100644
--- a/python/testSrc/com/jetbrains/python/PyDictLiteralCompletionTest.kt
+++ b/python/testSrc/com/jetbrains/python/PyDictLiteralCompletionTest.kt
@@ -9,6 +9,14 @@ class PyDictLiteralCompletionTest : PyTestCase() {
     assertCompletionContains("\"x\"", "\"y\"")
   }
 
+  fun testEmptyLiteralsInMethodCall() {
+    assertCompletionContains("\"name\"", "\"age\"")
+  }
+
+  fun testEmptyLiteralsInMethodCallWithKeywordArgument() {
+    assertCompletionContains("\"name\"", "\"age\"")
+  }
+
   fun testEmptyLiteralsInCallExpressionsWithQuotes() {
     assertCompletionContains("x", "y")
   }