diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md new file mode 100644 index 0000000000..5e1931f374 --- /dev/null +++ b/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: + +## Function syntax + +To do: + +## Exhaustiveness checking + +To do + +## References + +- Typing spec: +- Documentation: + +[class-private name]: https://docs.python.org/3/reference/lexical_analysis.html#reserved-classes-of-identifiers diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 8a1c6662b0..d10aa123e0 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/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; diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 18cb581773..771dec6d89 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/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( diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 6fadfa7df9..40a8b92ca9 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/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, }; diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs new file mode 100644 index 0000000000..53e534a9c3 --- /dev/null +++ b/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 { + 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> = FxHashSet::default(); + // TODO: handle `StrEnum` which uses lowercase names as values when using `auto()`. + let mut auto_counter = 0; + + let ignored_names: Option> = 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() +} diff --git a/crates/ty_python_semantic/src/types/function.rs b/crates/ty_python_semantic/src/types/function.rs index 189ea9dd12..b97c12960c 100644 --- a/crates/ty_python_semantic/src/types/function.rs +++ b/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 diff --git a/crates/ty_vendored/ty_extensions/ty_extensions.pyi b/crates/ty_vendored/ty_extensions/ty_extensions.pyi index 158f0dcbf4..f4a447d34c 100644 --- a/crates/ty_vendored/ty_extensions/ty_extensions.pyi +++ b/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: ...