IDEA-CR-61508: PY-39703 Add TypedDict keys completion

GitOrigin-RevId: b4e71d5edff69b23e1f1eb68ffe77f994a6a0293
This commit is contained in:
Lada Gagina
2020-04-16 20:10:49 +00:00
committed by intellij-monorepo-bot
parent a05f39db52
commit 61ad4c3109
2 changed files with 109 additions and 2 deletions
@@ -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['<caret>']";
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']['<caret>']";
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['<caret>']";
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['<caret>']";
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<String> suggested = doTestByText("from typing import TypedDict\n" +
@@ -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
*