[ty] List all enum members (#19283)

## Summary

Adds a way to list all members of an `Enum` and implements almost all of
the mechanisms by which members are distinguished from non-members
([spec](https://typing.python.org/en/latest/spec/enums.html#defining-members)).
This has no effect on actual enums, so far.

## Test Plan

New Markdown tests using `ty_extensions.enum_members`.

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -0,0 +1,439 @@
+# Enums
+
+## Basic
+
+```py
+from enum import Enum
+
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+reveal_type(Color.RED)  # revealed: @Todo(Attribute access on enum classes)
+reveal_type(Color.RED.value)  # revealed: @Todo(Attribute access on enum classes)
+
+# TODO: Should be `Color` or `Literal[Color.RED]`
+reveal_type(Color["RED"])  # revealed: Unknown
+
+# TODO: Could be `Literal[Color.RED]` to be more precise
+reveal_type(Color(1))  # revealed: Color
+
+reveal_type(Color.RED in Color)  # revealed: bool
+```
+
+## Enum members
+
+### Basic
+
+Simple enums with integer or string values:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class ColorInt(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(ColorInt))
+
+class ColorStr(Enum):
+    RED = "red"
+    GREEN = "green"
+    BLUE = "blue"
+
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(ColorStr))
+```
+
+### When deriving from `IntEnum`
+
+```py
+from enum import IntEnum
+from ty_extensions import enum_members
+
+class ColorInt(IntEnum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(ColorInt))
+```
+
+### Declared non-member attributes
+
+Attributes on the enum class that are declared are not considered members of the enum:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    YES = 1
+    NO = 2
+
+    non_member_1: int
+
+    # TODO: this could be considered an error:
+    non_member_1: str = "some value"
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+### Non-member attributes with disallowed type
+
+Methods, callables, descriptors (including properties), and nested classes that are defined in the
+class are not treated as enum members:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+from typing import Callable, Literal
+
+def identity(x) -> int:
+    return x
+
+class Descriptor:
+    def __get__(self, instance, owner):
+        return 0
+
+class Answer(Enum):
+    YES = 1
+    NO = 2
+
+    def some_method(self) -> None: ...
+    @staticmethod
+    def some_static_method() -> None: ...
+    @classmethod
+    def some_class_method(cls) -> None: ...
+
+    some_callable = lambda x: 0
+    declared_callable: Callable[[int], int] = identity
+    function_reference = identity
+
+    some_descriptor = Descriptor()
+
+    @property
+    def some_property(self) -> str:
+        return ""
+
+    class NestedClass: ...
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+### `enum.property`
+
+Enum attributes that are defined using `enum.property` are not considered members:
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from enum import Enum, property as enum_property
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    YES = 1
+    NO = 2
+
+    @enum_property
+    def some_property(self) -> str:
+        return "property value"
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+### `types.DynamicClassAttribute`
+
+Attributes defined using `types.DynamicClassAttribute` are not considered members:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+from types import DynamicClassAttribute
+
+class Answer(Enum):
+    YES = 1
+    NO = 2
+
+    @DynamicClassAttribute
+    def dynamic_property(self) -> str:
+        return "dynamic value"
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+### In stubs
+
+Stubs can optionally use `...` for the actual value:
+
+```pyi
+from enum import Enum
+from ty_extensions import enum_members
+from typing import cast
+
+class Color(Enum):
+    RED = ...
+    GREEN = cast(int, ...)
+    BLUE = 3
+
+# revealed: tuple[Literal["RED"], Literal["GREEN"], Literal["BLUE"]]
+reveal_type(enum_members(Color))
+```
+
+### Aliases
+
+Enum members can have aliases, which are not considered separate members:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    YES = 1
+    NO = 2
+
+    DEFINITELY = YES
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+### Using `auto()`
+
+```py
+from enum import Enum, auto
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    YES = auto()
+    NO = auto()
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+Combining aliases with `auto()`:
+
+```py
+from enum import Enum, auto
+
+class Answer(Enum):
+    YES = auto()
+    NO = auto()
+
+    DEFINITELY = YES
+
+# TODO: This should ideally be `tuple[Literal["YES"], Literal["NO"]]`
+# revealed: tuple[Literal["YES"], Literal["NO"], Literal["DEFINITELY"]]
+reveal_type(enum_members(Answer))
+```
+
+### `member` and `nonmember`
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from enum import Enum, auto, member, nonmember
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    YES = member(1)
+    NO = member(2)
+    OTHER = nonmember(17)
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+`member` can also be used as a decorator:
+
+```py
+from enum import Enum, member
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    yes = member(1)
+    no = member(2)
+
+    @member
+    def maybe(self) -> None:
+        return
+
+# revealed: tuple[Literal["yes"], Literal["no"], Literal["maybe"]]
+reveal_type(enum_members(Answer))
+```
+
+### Class-private names
+
+An attribute with a [class-private name] (beginning with, but not ending in, a double underscore) is
+treated as a non-member:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    YES = 1
+    NO = 2
+
+    __private_member = 3
+    __maybe__ = 4
+
+# revealed: tuple[Literal["YES"], Literal["NO"], Literal["__maybe__"]]
+reveal_type(enum_members(Answer))
+```
+
+### Ignored names
+
+An enum class can define a class symbol named `_ignore_`. This can be a string containing a
+whitespace-delimited list of names:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    _ignore_ = "IGNORED _other_ignored       also_ignored"
+
+    YES = 1
+    NO = 2
+
+    IGNORED = 3
+    _other_ignored = "test"
+    also_ignored = "test2"
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+`_ignore_` can also be a list of names:
+
+```py
+class Answer2(Enum):
+    _ignore_ = ["MAYBE", "_other"]
+
+    YES = 1
+    NO = 2
+
+    MAYBE = 3
+    _other = "test"
+
+# TODO: This should be `tuple[Literal["YES"], Literal["NO"]]`
+# revealed: tuple[Literal["YES"], Literal["NO"], Literal["MAYBE"], Literal["_other"]]
+reveal_type(enum_members(Answer2))
+```
+
+### Special names
+
+Make sure that special names like `name` and `value` can be used for enum members (without
+conflicting with `Enum.name` and `Enum.value`):
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class Answer(Enum):
+    name = 1
+    value = 2
+
+# revealed: tuple[Literal["name"], Literal["value"]]
+reveal_type(enum_members(Answer))
+
+# TODO: These should be `Answer` or `Literal[Answer.name]`/``Literal[Answer.value]`
+reveal_type(Answer.name)  # revealed: @Todo(Attribute access on enum classes)
+reveal_type(Answer.value)  # revealed: @Todo(Attribute access on enum classes)
+```
+
+## Iterating over enum members
+
+```py
+from enum import Enum
+
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+for color in Color:
+    # TODO: Should be `Color`
+    reveal_type(color)  # revealed: Unknown
+
+# TODO: Should be `list[Color]`
+reveal_type(list(Color))  # revealed: list[Unknown]
+```
+
+## Properties of enum types
+
+### Implicitly final
+
+An enum with one or more defined members cannot be subclassed. They are implicitly "final".
+
+```py
+from enum import Enum
+
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+# TODO: This should emit an error
+class ExtendedColor(Color):
+    YELLOW = 4
+
+def f(color: Color):
+    if isinstance(color, int):
+        # TODO: This should be `Never`
+        reveal_type(color)  # revealed: Color & int
+```
+
+An `Enum` subclass without any defined members can be subclassed:
+
+```py
+from enum import Enum
+from ty_extensions import enum_members
+
+class MyEnum(Enum):
+    def some_method(self) -> None:
+        pass
+
+class Answer(MyEnum):
+    YES = 1
+    NO = 2
+
+# revealed: tuple[Literal["YES"], Literal["NO"]]
+reveal_type(enum_members(Answer))
+```
+
+## Custom enum types
+
+To do: <https://typing.python.org/en/latest/spec/enums.html#enum-definition>
+
+## Function syntax
+
+To do: <https://typing.python.org/en/latest/spec/enums.html#enum-definition>
+
+## Exhaustiveness checking
+
+To do
+
+## References
+
+- Typing spec: <https://typing.python.org/en/latest/spec/enums.html>
+- Documentation: <https://docs.python.org/3/library/enum.html>
+
+[class-private name]: https://docs.python.org/3/reference/lexical_analysis.html#reserved-classes-of-identifiers

crates/ty_python_semantic/src/types.rs

@@ -69,6 +69,7 @@ mod context;
 mod cyclic;
 mod diagnostic;
 mod display;
+mod enums;
 mod function;
 mod generics;
 pub(crate) mod ide_support;

crates/ty_python_semantic/src/types/call/bind.rs

@@ -31,7 +31,7 @@ use crate::types::tuple::TupleType;
 use crate::types::{
     BoundMethodType, ClassLiteral, DataclassParams, KnownClass, KnownInstanceType,
     MethodWrapperKind, PropertyInstanceType, SpecialFormType, TypeMapping, UnionType,
-    WrapperDescriptorKind, ide_support, todo_type,
+    WrapperDescriptorKind, enums, ide_support, todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, Severity, SubDiagnostic};
 use ruff_python_ast as ast;
@@ -661,6 +661,22 @@ impl<'db> Bindings<'db> {
                             }
                         }
 
+                        Some(KnownFunction::EnumMembers) => {
+                            if let [Some(ty)] = overload.parameter_types() {
+                                let return_ty = match ty {
+                                    Type::ClassLiteral(class) => TupleType::from_elements(
+                                        db,
+                                        enums::enum_members(db, *class)
+                                            .into_iter()
+                                            .map(|member| Type::string_literal(db, &member)),
+                                    ),
+                                    _ => Type::unknown(),
+                                };
+
+                                overload.set_return_type(return_ty);
+                            }
+                        }
+
                         Some(KnownFunction::AllMembers) => {
                             if let [Some(ty)] = overload.parameter_types() {
                                 overload.set_return_type(TupleType::from_elements(

crates/ty_python_semantic/src/types/class.rs

@@ -2365,6 +2365,9 @@ pub enum KnownClass {
     Super,
     // enum
     Enum,
+    Auto,
+    Member,
+    Nonmember,
     // abc
     ABCMeta,
     // Types
@@ -2485,6 +2488,9 @@ impl KnownClass {
             | Self::Deque
             | Self::Float
             | Self::Enum
+            | Self::Auto
+            | Self::Member
+            | Self::Nonmember
             | Self::ABCMeta
             | Self::Iterable
             // Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue
@@ -2563,6 +2569,9 @@ impl KnownClass {
             Self::ABCMeta
             | Self::Any
             | Self::Enum
+            | Self::Auto
+            | Self::Member
+            | Self::Nonmember
             | Self::ChainMap
             | Self::Exception
             | Self::ExceptionGroup
@@ -2643,6 +2652,9 @@ impl KnownClass {
             | Self::Deque
             | Self::OrderedDict
             | Self::Enum
+            | Self::Auto
+            | Self::Member
+            | Self::Nonmember
             | Self::ABCMeta
             | Self::Super
             | Self::StdlibAlias
@@ -2708,6 +2720,9 @@ impl KnownClass {
             Self::Deque => "deque",
             Self::OrderedDict => "OrderedDict",
             Self::Enum => "Enum",
+            Self::Auto => "auto",
+            Self::Member => "member",
+            Self::Nonmember => "nonmember",
             Self::ABCMeta => "ABCMeta",
             Self::Super => "super",
             Self::Iterable => "Iterable",
@@ -2929,7 +2944,7 @@ impl KnownClass {
             | Self::Property => KnownModule::Builtins,
             Self::VersionInfo => KnownModule::Sys,
             Self::ABCMeta => KnownModule::Abc,
-            Self::Enum => KnownModule::Enum,
+            Self::Enum | Self::Auto | Self::Member | Self::Nonmember => KnownModule::Enum,
             Self::GenericAlias
             | Self::ModuleType
             | Self::FunctionType
@@ -3042,6 +3057,9 @@ impl KnownClass {
             | Self::ParamSpecKwargs
             | Self::TypeVarTuple
             | Self::Enum
+            | Self::Auto
+            | Self::Member
+            | Self::Nonmember
             | Self::ABCMeta
             | Self::Super
             | Self::NamedTuple
@@ -3110,6 +3128,9 @@ impl KnownClass {
             | Self::ParamSpecKwargs
             | Self::TypeVarTuple
             | Self::Enum
+            | Self::Auto
+            | Self::Member
+            | Self::Nonmember
             | Self::ABCMeta
             | Self::Super
             | Self::UnionType
@@ -3182,6 +3203,9 @@ impl KnownClass {
             "_NoDefaultType" => Self::NoDefaultType,
             "SupportsIndex" => Self::SupportsIndex,
             "Enum" => Self::Enum,
+            "auto" => Self::Auto,
+            "member" => Self::Member,
+            "nonmember" => Self::Nonmember,
             "ABCMeta" => Self::ABCMeta,
             "super" => Self::Super,
             "_version_info" => Self::VersionInfo,
@@ -3243,6 +3267,9 @@ impl KnownClass {
             | Self::MethodType
             | Self::MethodWrapperType
             | Self::Enum
+            | Self::Auto
+            | Self::Member
+            | Self::Nonmember
             | Self::ABCMeta
             | Self::Super
             | Self::NotImplementedType
@@ -3762,6 +3789,7 @@ mod tests {
                 KnownClass::BaseExceptionGroup | KnownClass::ExceptionGroup => PythonVersion::PY311,
                 KnownClass::GenericAlias => PythonVersion::PY39,
                 KnownClass::KwOnly => PythonVersion::PY310,
+                KnownClass::Member | KnownClass::Nonmember => PythonVersion::PY311,
                 _ => PythonVersion::PY37,
             };
 

crates/ty_python_semantic/src/types/enums.rs

@@ -0,0 +1,145 @@
+use rustc_hash::FxHashSet;
+
+use crate::{
+    Db,
+    place::{Place, place_from_bindings, place_from_declarations},
+    semantic_index::{place_table, use_def_map},
+    types::{ClassLiteral, KnownClass, MemberLookupPolicy, Type},
+};
+
+/// List all members of an enum.
+pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Vec<String> {
+    let scope_id = class.body_scope(db);
+    let use_def_map = use_def_map(db, scope_id);
+    let table = place_table(db, scope_id);
+
+    let mut enum_values: FxHashSet<Type<'db>> = FxHashSet::default();
+    // TODO: handle `StrEnum` which uses lowercase names as values when using `auto()`.
+    let mut auto_counter = 0;
+
+    let ignored_names: Option<Vec<&str>> = if let Some(ignore) = table.place_id_by_name("_ignore_")
+    {
+        let ignore_bindings = use_def_map.all_reachable_bindings(ignore);
+        let ignore_place = place_from_bindings(db, ignore_bindings);
+
+        match ignore_place {
+            Place::Type(Type::StringLiteral(ignored_names), _) => {
+                Some(ignored_names.value(db).split_ascii_whitespace().collect())
+            }
+            // TODO: support the list-variant of `_ignore_`.
+            _ => None,
+        }
+    } else {
+        None
+    };
+
+    use_def_map
+        .all_end_of_scope_bindings()
+        .filter_map(|(place_id, bindings)| {
+            let name = table
+                .place_expr(place_id)
+                .as_name()
+                .map(ToString::to_string)?;
+
+            if name.starts_with("__") && !name.ends_with("__") {
+                // Skip private attributes
+                return None;
+            }
+
+            if name == "_ignore_"
+                || ignored_names
+                    .as_ref()
+                    .is_some_and(|names| names.contains(&name.as_str()))
+            {
+                // Skip ignored attributes
+                return None;
+            }
+
+            let inferred = place_from_bindings(db, bindings);
+            let value_ty = match inferred {
+                Place::Unbound => {
+                    return None;
+                }
+                Place::Type(ty, _) => {
+                    match ty {
+                        Type::Callable(_) | Type::FunctionLiteral(_) => {
+                            // Some types are specifically disallowed for enum members.
+                            return None;
+                        }
+                        // enum.nonmember
+                        Type::NominalInstance(instance)
+                            if instance.class.is_known(db, KnownClass::Nonmember) =>
+                        {
+                            return None;
+                        }
+                        // enum.member
+                        Type::NominalInstance(instance)
+                            if instance.class.is_known(db, KnownClass::Member) =>
+                        {
+                            ty.member(db, "value")
+                                .place
+                                .ignore_possibly_unbound()
+                                .unwrap_or(Type::unknown())
+                        }
+                        // enum.auto
+                        Type::NominalInstance(instance)
+                            if instance.class.is_known(db, KnownClass::Auto) =>
+                        {
+                            auto_counter += 1;
+                            Type::IntLiteral(auto_counter)
+                        }
+                        _ => {
+                            let dunder_get = ty
+                                .member_lookup_with_policy(
+                                    db,
+                                    "__get__".into(),
+                                    MemberLookupPolicy::NO_INSTANCE_FALLBACK,
+                                )
+                                .place;
+
+                            match dunder_get {
+                                Place::Unbound | Place::Type(Type::Dynamic(_), _) => ty,
+
+                                Place::Type(_, _) => {
+                                    // Descriptors are not considered members.
+                                    return None;
+                                }
+                            }
+                        }
+                    }
+                }
+            };
+
+            // Duplicate values are aliases that are not considered separate members. This check is only
+            // performed if we can infer a precise literal type for the enum member. If we only get `int`,
+            // we don't know if it's a duplicate or not.
+            if matches!(
+                value_ty,
+                Type::IntLiteral(_) | Type::StringLiteral(_) | Type::BytesLiteral(_)
+            ) && !enum_values.insert(value_ty)
+            {
+                return None;
+            }
+
+            let declarations = use_def_map.end_of_scope_declarations(place_id);
+            let declared = place_from_declarations(db, declarations);
+
+            match declared.map(|d| d.place) {
+                Ok(Place::Unbound) => {
+                    // Undeclared attributes are considered members
+                }
+                Ok(Place::Type(Type::NominalInstance(instance), _))
+                    if instance.class.is_known(db, KnownClass::Member) =>
+                {
+                    // If the attribute is specifically declared with `enum.member`, it is considered a member
+                }
+                _ => {
+                    // Declared attributes are considered non-members
+                    return None;
+                }
+            }
+
+            Some(name)
+        })
+        .collect()
+}

crates/ty_python_semantic/src/types/function.rs

@@ -957,6 +957,8 @@ pub enum KnownFunction {
     GenericContext,
     /// `ty_extensions.dunder_all_names`
     DunderAllNames,
+    /// `ty_extensions.enum_members`
+    EnumMembers,
     /// `ty_extensions.all_members`
     AllMembers,
     /// `ty_extensions.top_materialization`
@@ -1025,6 +1027,7 @@ impl KnownFunction {
             | Self::BottomMaterialization
             | Self::GenericContext
             | Self::DunderAllNames
+            | Self::EnumMembers
             | Self::StaticAssert
             | Self::AllMembers => module.is_ty_extensions(),
             Self::ImportModule => module.is_importlib(),
@@ -1288,6 +1291,7 @@ pub(crate) mod tests {
                 | KnownFunction::IsSubtypeOf
                 | KnownFunction::GenericContext
                 | KnownFunction::DunderAllNames
+                | KnownFunction::EnumMembers
                 | KnownFunction::StaticAssert
                 | KnownFunction::IsDisjointFrom
                 | KnownFunction::IsSingleValued

crates/ty_vendored/ty_extensions/ty_extensions.pyi

@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Any, LiteralString, _SpecialForm
 
 # Special operations
@@ -42,6 +43,9 @@ def generic_context(type: Any) -> Any: ...
 # either the module does not have `__all__` or it has invalid elements.
 def dunder_all_names(module: Any) -> Any: ...
 
+# List all members of an enum.
+def enum_members[E: type[Enum]](enum: E) -> tuple[str, ...]: ...
+
 # Returns the type that's an upper bound of materializing the given (gradual) type.
 def top_materialization(type: Any) -> Any: ...