diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java index 5e0734030a16..d98567f40c43 100644 --- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java +++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java @@ -60,7 +60,7 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { if (type != null) { return type; } - type = getEnumType(referenceTarget, context); + type = getEnumType(referenceTarget, context, anchor); if (type != null) { return type; } @@ -68,7 +68,8 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { } @Nullable - private static PyType getEnumType(@NotNull PsiElement referenceTarget, @NotNull TypeEvalContext context) { + private static PyType getEnumType(@NotNull PsiElement referenceTarget, @NotNull TypeEvalContext context, + @Nullable PsiElement anchor) { if (referenceTarget instanceof PyTargetExpression) { final PyTargetExpression target = (PyTargetExpression)referenceTarget; final ScopeOwner owner = ScopeUtil.getScopeOwner(target); @@ -85,6 +86,34 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase { } } } + if (referenceTarget instanceof PyQualifiedNameOwner) { + final PyQualifiedNameOwner qualifiedNameOwner = (PyQualifiedNameOwner)referenceTarget; + final String name = qualifiedNameOwner.getQualifiedName(); + if ("enum.Enum.name".equals(name)) { + return PyBuiltinCache.getInstance(referenceTarget).getStrType(); + } + else if ("enum.Enum.value".equals(name) && anchor instanceof PyReferenceExpression && context.maySwitchToAST(anchor)) { + final PyReferenceExpression anchorExpr = (PyReferenceExpression)anchor; + final PyExpression qualifier = anchorExpr.getQualifier(); + if (qualifier instanceof PyReferenceExpression) { + 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(); + if (value != null) { + return context.getType(value); + } + } + } + } + } + else if ("enum.EnumMeta.__members__".equals(name)) { + return PyTypeParser.getTypeByName(referenceTarget, "dict[str, unknown]"); + } + } return null; } diff --git a/python/testData/inspections/PyUnresolvedReferencesInspection3K/EnumMemberAttributes/a.py b/python/testData/inspections/PyUnresolvedReferencesInspection3K/EnumMemberAttributes/a.py index 36d36b0ec620..9d4ffeb8d729 100644 --- a/python/testData/inspections/PyUnresolvedReferencesInspection3K/EnumMemberAttributes/a.py +++ b/python/testData/inspections/PyUnresolvedReferencesInspection3K/EnumMemberAttributes/a.py @@ -8,5 +8,11 @@ class Color(Enum): print(Color.red.name, Color.red.name.upper()) +print(Color.red.name.foo) print(Color.red.value, Color.red.value.real) +print(Color.red.value.foo) print(Color.red.foo) + + +print(Color.__members__.items()) +print(Color.__members__.foo)