From 61ad4c31096697fffabe7b4378930bd1d0a059c0 Mon Sep 17 00:00:00 2001 From: Lada Gagina Date: Tue, 31 Mar 2020 13:18:47 +0300 Subject: [PATCH] IDEA-CR-61508: PY-39703 Add TypedDict keys completion GitOrigin-RevId: b4e71d5edff69b23e1f1eb68ffe77f994a6a0293 --- .../python/PythonCommonCompletionTest.java | 85 +++++++++++++++++++ .../PyDictKeyNamesCompletionContributor.java | 26 +++++- 2 files changed, 109 insertions(+), 2 deletions(-) diff --git a/python/python-common-tests/com/jetbrains/python/PythonCommonCompletionTest.java b/python/python-common-tests/com/jetbrains/python/PythonCommonCompletionTest.java index de5bf725073c..f8f4cc10a728 100644 --- a/python/python-common-tests/com/jetbrains/python/PythonCommonCompletionTest.java +++ b/python/python-common-tests/com/jetbrains/python/PythonCommonCompletionTest.java @@ -1641,6 +1641,91 @@ public abstract class PythonCommonCompletionTest extends PythonCommonTestCase { assertContainsElements(suggested, "update", "clear", "pop", "popitem", "setdefault"); } + // PY-39703 + public void testTypedDictUsingConstructorKeysCompletion() { + final String test1 = "from typing import TypedDict\n" + + "class A(TypedDict):\n" + + " x: int\n" + + " y: int\n" + + "a = A(x=42, y=1)\n" + + "b = a['']"; + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> { + myFixture.configureByText(PythonFileType.INSTANCE, test1); + myFixture.completeBasic(); + myFixture.checkResult(test1); + + assertContainsElements(myFixture.getLookupElementStrings(), "'x'", "'y'"); + } + ); + } + + // PY-39703 + public void testTypedDictInTypedDictKeys() { + final String test2 = "from typing import TypedDict, Mapping\n" + + "class Movie(TypedDict):\n" + + " name: str\n" + + " year: int\n" + + "class A(TypedDict):\n" + + " film: Movie\n" + + " genre: str\n" + + "def get_back(m: A):\n" + + " a = m['film']['']"; + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> { + myFixture.configureByText(PythonFileType.INSTANCE, test2); + myFixture.completeBasic(); + myFixture.checkResult(test2); + + assertContainsElements(myFixture.getLookupElementStrings(), "'name'", "'year'"); + } + ); + } + + // PY-39703 + public void testTypedDictKeysUsingTypeAnnotation() { + final String test3 = "from typing import TypedDict\n" + + "class Movie(TypedDict, total=False):\n" + + " name: str\n" + + " year: int\n" + + "m: Movie = {}\n" + + "m['']"; + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> { + myFixture.configureByText(PythonFileType.INSTANCE, test3); + myFixture.completeBasic(); + myFixture.checkResult(test3); + + assertContainsElements(myFixture.getLookupElementStrings(), "'name'", "'year'"); + } + ); + } + + // PY-39703 + public void testTypedDictWithWrongKey() { + final String test4 = "from typing import TypedDict\n" + + "class Movie(TypedDict, total=False):\n" + + " name: str\n" + + " year: int\n" + + "m: Movie = {'name': 'Alien', 'year': 1979}\n" + + "m['wrong_key'] = 1\n" + + "m['']"; + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> { + myFixture.configureByText(PythonFileType.INSTANCE, test4); + myFixture.completeBasic(); + myFixture.checkResult(test4); + + assertContainsElements(myFixture.getLookupElementStrings(), "'name'", "'year'"); + assertDoesntContain(myFixture.getLookupElementStrings(), "'wrong_key'"); + } + ); + } + // PY-36008 public void testTypedDictDefinition() { final List suggested = doTestByText("from typing import TypedDict\n" + diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictKeyNamesCompletionContributor.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictKeyNamesCompletionContributor.java index c876380c8acc..de710c471d83 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictKeyNamesCompletionContributor.java +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyDictKeyNamesCompletionContributor.java @@ -17,6 +17,7 @@ import com.jetbrains.python.PyTokenTypes; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.resolve.PyResolveUtil; import com.jetbrains.python.psi.types.PyType; +import com.jetbrains.python.psi.types.PyTypedDictType; import com.jetbrains.python.psi.types.TypeEvalContext; import org.jetbrains.annotations.NotNull; @@ -46,9 +47,12 @@ public class PyDictKeyNamesCompletionContributor extends CompletionContributor { final PySubscriptionExpression subscription = PsiTreeUtil.getParentOfType(original, PySubscriptionExpression.class); if (subscription == null) return; - final PsiElement operand = subscription.getOperand(); + final PyExpression operand = subscription.getOperand(); + if (addCompletionIfOperandIsTypedDict(operand, dictCompletion)) { + return; + } if (operand instanceof PyReferenceExpression) { - final PsiElement resolvedElement = PyResolveUtil.fullResolveLocally((PyReferenceExpression)operand); + final PyExpression resolvedElement = PyResolveUtil.fullResolveLocally((PyReferenceExpression)operand); if (resolvedElement instanceof PyDictLiteralExpression) { addDictLiteralKeys((PyDictLiteralExpression)resolvedElement, dictCompletion); addAdditionalKeys(parameters.getOriginalFile(), operand, dictCompletion); @@ -63,6 +67,24 @@ public class PyDictKeyNamesCompletionContributor extends CompletionContributor { ); } + /** + * Add index expression completion if an operand is a TypedDict + * + * @return true if an operand is a TypedDict + */ + private static boolean addCompletionIfOperandIsTypedDict(@NotNull final PyExpression operand, + @NotNull final CompletionResultSet dictCompletion) { + final TypeEvalContext typeEvalContext = TypeEvalContext.codeCompletion(operand.getProject(), operand.getContainingFile()); + final PyType type = typeEvalContext.getType(operand); + if (type instanceof PyTypedDictType) { + for (String key : ((PyTypedDictType)type).getFields().keySet()) { + dictCompletion.addElement(createElement("'" + key + "'")); + } + return true; + } + return false; + } + /** * create completion result with prefix matcher if needed *