PY-54503 Provide the result type for Enum[...].value and Enum(...).value

Precise types can be inferred only over AST if it's accessible.

I had to move PyStdlibTypeProvider higher in the provider's hierarchy
so that it could override types coming from Typeshed, otherwise we
infer enum.property type for the "value" attribute.

GitOrigin-RevId: 8727e080cfc06d0edda13eccfd601601dc661da9
This commit is contained in:
Mikhail Golubev
2022-09-02 20:51:03 +00:00
committed by intellij-monorepo-bot
parent 0fc1981b5c
commit 59e79c3c7c
3 changed files with 49 additions and 5 deletions
@@ -523,7 +523,8 @@
<inspectionExtension implementation="com.jetbrains.python.inspections.stdlib.PyStdlibInspectionExtension"/>
<typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyDataclassTypeProvider"/>
<typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyNamedTupleTypeProvider"/>
<typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibTypeProvider"/>
<!-- This provider should be able to override the results on anything else, including types coming from .pyi stubs -->
<typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibTypeProvider" order="first"/>
<pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyOverridingClassDunderMembersProvider"/>
<pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider"/>
<pyModuleMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibOverridingModuleMembersProvider"/>
@@ -110,10 +110,23 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase {
final PyReferenceExpression qualifierExpr = (PyReferenceExpression)qualifier;
final PsiElement resolvedQualifier = qualifierExpr.getReference().resolve();
if (resolvedQualifier instanceof PyTargetExpression) {
final PyTargetExpression qualifierTarget = (PyTargetExpression)resolvedQualifier;
// Requires switching to AST, we cannot use getType(qualifierTarget) here, because its type is overridden by this type provider
if (context.maySwitchToAST(qualifierTarget)) {
final PyExpression value = qualifierTarget.findAssignedValue();
final PyTargetExpression enumItem = (PyTargetExpression)resolvedQualifier;
// Requires switching to AST, we cannot use getType(enumItem) here, because its type is overridden by this type provider
if (context.maySwitchToAST(enumItem)) {
final PyExpression value = enumItem.findAssignedValue();
if (value != null) {
return context.getType(value);
}
}
}
}
else if (qualifier != null) {
PyClassType enumType = as(context.getType(qualifier), PyClassType.class);
if (enumType != null) {
PyClass enumClass = enumType.getPyClass();
PyTargetExpression firstEnumItem = ContainerUtil.getFirstItem(enumClass.getClassAttributes());
if (firstEnumItem != null && context.maySwitchToAST(firstEnumItem)) {
final PyExpression value = firstEnumItem.findAssignedValue();
if (value != null) {
return context.getType(value);
}
@@ -3979,6 +3979,36 @@ public class PyTypeTest extends PyTestCase {
);
}
// PY-54503
public void testEnumGetItemResultValueAttribute() {
runWithLanguageLevel(
LanguageLevel.getLatest(),
() -> doTest("int",
"import enum\n" +
"\n" +
"class MyEnum(enum.Enum):\n" +
" ONE = 1\n" +
" TWO = 2\n" +
"\n" +
"expr = MyEnum['ONE'].value")
);
}
// PY-54503
public void testEnumDunderCallResultValueAttribute() {
runWithLanguageLevel(
LanguageLevel.getLatest(),
() -> doTest("int",
"import enum\n" +
"\n" +
"class MyEnum(enum.Enum):\n" +
" ONE = 1\n" +
" TWO = 2\n" +
"\n" +
"expr = MyEnum(1).value")
);
}
private static List<TypeEvalContext> getTypeEvalContexts(@NotNull PyExpression element) {
return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(),
TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());