diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index f6f1b4eec4d8e..c1379a3887b37 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -2095,6 +2095,416 @@ def _(x: EnumWithSubclassOfEnumMetaMetaclass): reveal_type(x._name_) # revealed: Literal["NO", "YES"] ``` +## Flags + +`Flag` members represent integer masks. Zero-valued and multi-bit names are valid attributes but are +not included when the class is iterated. `auto()` assigns the next power of two, including when an +earlier member used a larger explicit value: + +```toml +[environment] +python-version = "3.13" +``` + +```py +from enum import CONFORM, EJECT, EnumMeta, Flag, FlagBoundary, IntFlag, auto, member +from typing import Any, Literal +from ty_extensions import enum_members + +class Permission(Flag): + NONE = 0 + READ = auto() + WRITE = auto() + EXECUTE = auto() + READ_WRITE = READ | WRITE + READ_ALIAS = READ + +# Zero, composites, and aliases are not canonical members. +# revealed: tuple[Literal["READ"], Literal["WRITE"], Literal["EXECUTE"]] +reveal_type(enum_members(Permission)) + +reveal_type(Permission.EXECUTE.value) # revealed: Literal[4] +reveal_type(Permission.READ_WRITE.value) # revealed: Literal[3] +reveal_type(Permission.READ_ALIAS) # revealed: Literal[Permission.READ] +reveal_type(Permission(Permission.READ)) # revealed: Literal[Permission.READ] +reveal_type(Permission(value=1)) # revealed: Literal[Permission.READ] + +class OutOfOrder(Flag): + HIGH = 8 + LOW = 1 + NEXT = auto() + +reveal_type(OutOfOrder.NEXT.value) # revealed: Literal[16] + +class WrappedAuto(Flag): + FIRST = member(auto()) + SECOND = auto() + +reveal_type(WrappedAuto.SECOND.value) # revealed: Literal[2] + +class BooleanValue(Flag): + TRUE = True + ONE = 1 + +reveal_type(BooleanValue.TRUE.value) # revealed: Literal[True] +reveal_type(BooleanValue.ONE) # revealed: Literal[BooleanValue.TRUE] +reveal_type(enum_members(BooleanValue)) # revealed: tuple[Literal["TRUE"]] + +# A custom integer mixin can change member values during class creation. +class Offset(int): + def __new__(cls, value: int): + return int.__new__(cls, value + 1) + +class TransformedValues(Offset, Flag): + A = 1 + B = 2 + +reveal_type(enum_members(TransformedValues)) # revealed: Unknown + +def unknown_values(value: int): + class Dynamic(Flag): + A = value + B = auto() + + reveal_type(enum_members(Dynamic)) # revealed: Unknown + reveal_type(Dynamic.B.value) # revealed: int +``` + +Bitwise operations use the member masks without enumerating every possible combination. A result is +a literal when it has a declared name and otherwise remains the nominal Flag type: + +```py +reveal_type(Permission.READ | Permission.WRITE) # revealed: Literal[Permission.READ_WRITE] +reveal_type(Permission.READ & Permission.WRITE) # revealed: Literal[Permission.NONE] +reveal_type(Permission.READ ^ Permission.WRITE) # revealed: Literal[Permission.READ_WRITE] +reveal_type(~Permission.READ) # revealed: Permission + +reveal_type(bool(Permission.NONE)) # revealed: Literal[False] +reveal_type(bool(Permission.READ_WRITE)) # revealed: Literal[True] +reveal_type(len(Permission.READ_WRITE)) # revealed: Literal[2] + +# revealed: tuple[Literal[Permission.READ], Literal[Permission.WRITE]] +reveal_type(tuple(Permission.READ_WRITE)) + +class MissingSingleBit(Flag): + A = 1 + AB = 3 + +# CPython yields `None` for the unnamed single-bit value `2`. +reveal_type(tuple(MissingSingleBit.AB)) # revealed: tuple[MissingSingleBit | None, ...] + +reveal_type(Permission.READ in Permission.READ_WRITE) # revealed: Literal[True] +reveal_type(Permission.EXECUTE in Permission.READ_WRITE) # revealed: Literal[False] + +def attributes(value: Permission): + reveal_type(value.name) # revealed: str | None + reveal_type(value.value) # revealed: int + +def gradual(left: Permission, right: Any): + if isinstance(right, Permission): + reveal_type(left | right) # revealed: Permission + +class IntValue(int): + pass + +class MixedFlag(IntValue, Flag): + A = 1 + B = 2 + +# A Flag with a custom integer member type accepts that type, not every integer. +reveal_type(MixedFlag.A | 2) # revealed: int +reveal_type(MixedFlag.A | IntValue(2)) # revealed: MixedFlag +``` + +The boundary policy determines what happens to undeclared bits. `CONFORM` removes them, `EJECT` +returns an integer, and the default `IntFlag` policy keeps them: + +```py +class Strict(Flag): + A = 1 + C = 4 + +class Conforming(Flag, boundary=CONFORM): + NONE = 0 + A = 1 + C = 4 + +class Ejecting(IntFlag, boundary=EJECT): + A = 1 + C = 4 + +class CustomMissing(IntFlag, boundary=EJECT): + A = 1 + + @classmethod + def _missing_(cls, value: object) -> Any: + return 42 + +class SignedEjecting(Flag, boundary=EJECT): + NEGATIVE = -1 + A = 1 + B = 2 + +class Keeping(IntFlag): + A = 1 + C = 4 + NOT_A = 6 + +class OtherIntFlag(IntFlag): + B = 2 + +class CrossLeft(IntFlag): + A = 1 + B = 2 + AB = 3 + C = 4 + +class CrossConforming(IntFlag, boundary=CONFORM): + B = 2 + +class CustomReflected(IntFlag): + B = 2 + + def __ror__(self, other: int) -> Any: + return 4 + +reveal_type(~Strict.A) # revealed: Literal[Strict.C] +reveal_type(Conforming(2)) # revealed: Literal[Conforming.NONE] +reveal_type(Ejecting(2)) # revealed: Literal[2] +reveal_type(Ejecting(value=2)) # revealed: Literal[2] +reveal_type(CustomMissing(2)) # revealed: int +reveal_type(CustomMissing(CustomMissing.A)) # revealed: Literal[CustomMissing.A] +reveal_type(Ejecting.A | 2) # revealed: Literal[3] +reveal_type(CustomMissing.A | 2) # revealed: int +reveal_type(2 | Ejecting.A) # revealed: Literal[3] +reveal_type(True | Ejecting.A) # revealed: int +reveal_type(Ejecting.A | OtherIntFlag.B) # revealed: OtherIntFlag +reveal_type(OtherIntFlag.B | Ejecting.A) # revealed: OtherIntFlag +reveal_type(CrossLeft.A | CrossConforming.B) # revealed: Literal[CrossLeft.B] +reveal_type(CrossLeft.A | CustomReflected.B) # revealed: CrossLeft +reveal_type(~Ejecting.A) # revealed: Literal[-2] +reveal_type(SignedEjecting(-1)) # revealed: Literal[SignedEjecting.NEGATIVE] +reveal_type(SignedEjecting(3)) # revealed: Literal[3] +reveal_type(Keeping.A | 8) # revealed: Keeping +reveal_type(~Keeping.A) # revealed: Literal[Keeping.NOT_A] +reveal_type(Keeping(-12)) # revealed: Literal[Keeping.C] +Keeping.A | "invalid" # error: [unsupported-operator] + +def dynamic_boundary(value: int): + reveal_type(Ejecting(value)) # revealed: int + reveal_type(Ejecting.A | value) # revealed: int + +class EjectBase(Flag, boundary=EJECT): + pass + +class InheritedEject(EjectBase): + A = 1 + +reveal_type(InheritedEject(2)) # revealed: Literal[2] + +def unknown_boundary(boundary: FlagBoundary): + class Configured(Flag, boundary=boundary): + A = 1 + + reveal_type(Configured(0)) # revealed: Configured + reveal_type(Configured(2)) # revealed: Configured | int +``` + +The same rules apply to classes created with the functional syntax: + +```py +Functional = Flag("Functional", {"HIGH": 8, "LOW": 1, "NEXT": auto()}) +reveal_type(Functional.NEXT.value) # revealed: Literal[16] + +FunctionalEject = IntFlag("FunctionalEject", {"A": 1, "C": 4}, boundary=EJECT) +reveal_type(FunctionalEject(2)) # revealed: Literal[2] +``` + +User-defined operators retain their declared return types: + +```py +class CustomOperator(Flag): + A = 1 + + # error: [invalid-method-override] + def __or__(self, other: "CustomOperator") -> str: + return "custom" + +reveal_type(CustomOperator.A | CustomOperator.A) # revealed: str +``` + +On Python 3.11 and later, `EnumMeta` replaces inherited operator overrides with the standard `Flag` +implementations on concrete Flag classes: + +```py +class ModernOperatorBase(Flag): + # error: [invalid-method-override] + def __or__(self, other: "ModernOperatorBase") -> str: + return "custom" + +class CopiedOperator(ModernOperatorBase): + A = 1 + B = 2 + AB = 3 + +reveal_type(CopiedOperator.A | CopiedOperator.B) # revealed: Literal[CopiedOperator.AB] +``` + +A custom metaclass `__call__` also controls results that the standard Flag methods construct: + +```py +class CallingMeta(EnumMeta): + def __call__(cls, *args: Any, **kwargs: Any) -> Any: + return "custom" + +class CustomCall(Flag, metaclass=CallingMeta): + A = 1 + B = 2 + AB = 3 + +reveal_type(CustomCall(1)) # revealed: Any +reveal_type(CustomCall.A | CustomCall.B) # revealed: Any +reveal_type(~CustomCall.A) # revealed: Any +``` + +Inherited behavioral methods and custom iteration hooks also take precedence over the standard Flag +behavior: + +```py +class BehaviorBase(Flag): + def __bool__(self) -> Literal[False]: + return False + + def __len__(self) -> Literal[7]: + return 7 + + def __contains__(self, other: Flag) -> Literal[False]: + return False + +class CustomBehavior(BehaviorBase): + A = 1 + +reveal_type(bool(CustomBehavior.A)) # revealed: Literal[False] +reveal_type(len(CustomBehavior.A)) # revealed: Literal[7] +reveal_type(CustomBehavior.A in CustomBehavior.A) # revealed: Literal[False] + +class IterationBase(Flag): + @classmethod + def _iter_member_(cls, value: int): + return iter(()) + +class CustomIteration(IterationBase): + A = 1 + +reveal_type(tuple(CustomIteration.A)) # revealed: tuple[CustomIteration, ...] +``` + +When canonical members are declared out of value order, `EnumMeta` generates `_iter_member_` from +the two helper methods. Overrides of either helper can change iteration: + +```py +class CustomDefinitionIteration(Flag): + @classmethod + def _iter_member_by_def_(cls, value: int): + return iter(()) + + HIGH = 2 + LOW = 1 + BOTH = HIGH | LOW + +reveal_type(tuple(CustomDefinitionIteration.BOTH)) # revealed: tuple[CustomDefinitionIteration, ...] + +class CustomValueIteration(Flag): + @classmethod + def _iter_member_by_value_(cls, value: int): + return iter(()) + + HIGH = 2 + LOW = 1 + BOTH = HIGH | LOW + +reveal_type(tuple(CustomValueIteration.BOTH)) # revealed: tuple[CustomValueIteration, ...] +``` + +## Flag `auto()` before Python 3.11 + +Before Python 3.11, `Flag._generate_next_value_` uses the most recently assigned value rather than +the largest value assigned so far. + +```toml +[environment] +python-version = "3.10" +``` + +```py +from enum import Flag, auto + +LegacyOrder = Flag("LegacyOrder", {"HIGH": 8, "LOW": 1, "NEXT": auto()}) + +reveal_type(LegacyOrder.NEXT.value) # revealed: Literal[2] +``` + +## IntFlag negative values before Python 3.11 + +Before Python 3.11, `IntFlag` preserves negative pseudo-member values instead of normalizing them to +the class's positive mask. + +```toml +[environment] +python-version = "3.10" +``` + +```py +from enum import IntFlag + +class LegacyIntFlag(IntFlag): + A = 1 + NOT_A = 6 + +reveal_type(LegacyIntFlag(-2)) # revealed: LegacyIntFlag +reveal_type(~LegacyIntFlag.A) # revealed: LegacyIntFlag +``` + +## Flag behavior before Python 3.11 + +Before Python 3.11, class iteration includes every non-alias `Flag` name, including zero and +composite values. Inversion uses that member set. + +```toml +[environment] +python-version = "3.10" +``` + +```py +from enum import Flag +from ty_extensions import enum_members + +class LegacyMembers(Flag): + ZERO = 0 + A = 1 + AB = 3 + +# revealed: tuple[Literal["ZERO"], Literal["A"], Literal["AB"]] +reveal_type(enum_members(LegacyMembers)) +reveal_type(~LegacyMembers.ZERO) # revealed: Literal[LegacyMembers.AB] +``` + +Before Python 3.11, inherited operators are not replaced with the standard `Flag` implementations: + +```py +class OperatorBase(Flag): + # error: [invalid-method-override] + def __or__(self, other: "OperatorBase") -> str: + return "custom" + +class InheritedOperator(OperatorBase): + A = 1 + B = 2 + AB = 3 + +reveal_type(InheritedOperator.A | InheritedOperator.B) # revealed: str +``` + ## Function syntax ### String names (positional) diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 4d91cdbf2775e..84e625551c495 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -4243,6 +4243,13 @@ impl<'db> Type<'db> { } } + if let Some(length) = self + .as_enum_literal() + .and_then(|literal| enums::flag_literal_len(db, literal)) + { + return Some(Type::int_literal(length)); + } + let return_ty = match self.try_call_dunder( db, "__len__", @@ -5144,6 +5151,7 @@ impl<'db> Type<'db> { .to_class_literal(db) .to_class_type(db) .is_some_and(|enum_class| class.is_subclass_of(db, enum_class)) + && enums::enum_uses_standard_metaclass_call(db, class_literal) { return fallback_bindings(); } diff --git a/crates/ty_python_semantic/src/types/bool.rs b/crates/ty_python_semantic/src/types/bool.rs index e5385a6677256..49eb9ee84b92b 100644 --- a/crates/ty_python_semantic/src/types/bool.rs +++ b/crates/ty_python_semantic/src/types/bool.rs @@ -7,7 +7,8 @@ use crate::{ CallArguments, CallDunderError, ClassType, CycleDetector, KnownClass, KnownInstanceType, LiteralValueTypeKind, SubclassOfInner, Type, TypeContext, TypeVarBoundOrConstraints, UnionType, call::CallErrorKind, constraints::ConstraintSetBuilder, context::InferContext, - diagnostic::UNSUPPORTED_BOOL_CONVERSION, typed_dict::TypedDictField, + diagnostic::UNSUPPORTED_BOOL_CONVERSION, enums::flag_literal_truthiness, + typed_dict::TypedDictField, }, }; use ty_python_core::Truthiness; @@ -306,9 +307,17 @@ impl<'db> Type<'db> { Type::LiteralValue(literal) => match literal.kind() { LiteralValueTypeKind::LiteralString => Truthiness::Ambiguous, - LiteralValueTypeKind::Enum(enum_type) => enum_type - .enum_class_instance(db) - .try_bool_impl(db, allow_short_circuit, visitor)?, + LiteralValueTypeKind::Enum(enum_type) => { + if let Some(truthiness) = flag_literal_truthiness(db, enum_type) { + Truthiness::from(truthiness) + } else { + enum_type.enum_class_instance(db).try_bool_impl( + db, + allow_short_circuit, + visitor, + )? + } + } LiteralValueTypeKind::Int(num) => Truthiness::from(num.as_i64() != 0), LiteralValueTypeKind::Bool(bool) => Truthiness::from(bool), diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 5e360954d28ba..7450330101d74 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -2065,12 +2065,12 @@ impl<'db> Bindings<'db> { if let [Some(ty)] = overload.parameter_types() { let return_ty = match ty { Type::ClassLiteral(class) => { - if let Some(metadata) = enums::enum_metadata(db, *class) { + if let Some(metadata) = enums::enum_metadata(db, *class) + && let Some(members) = metadata.canonical_member_names() + { Type::heterogeneous_tuple( db, - metadata - .members - .keys() + members .map(|member| Type::string_literal(db, member)), ) } else { diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 56ee09f550378..4b6d26921687a 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -23,7 +23,7 @@ use crate::types::callable::{CallableFunctionProvenance, CallableTypeKind}; use crate::types::constraints::{ ConstraintSet, ConstraintSetBuilder, IteratorConstraintsExtension, }; -use crate::types::enums::enum_metadata; +use crate::types::enums::{enum_metadata, enum_uses_standard_metaclass_call}; use crate::types::function::{AbstractMethodKind, DataclassTransformerParams}; use crate::types::generics::{ GenericContext, InferableTypeVars, Specialization, walk_specialization, @@ -1985,7 +1985,7 @@ impl<'db> ClassType<'db> { // `Color("red")`, instead of the overloaded signature of `EnumMeta.__call__` which also accounts // for dynamic Enum creation. let is_actual_enum = enum_metadata(db, self.class_literal(db)).is_some(); - if !is_actual_enum { + if !is_actual_enum || !enum_uses_standard_metaclass_call(db, self.class_literal(db)) { return CallableTypes::one(metaclass_dunder_call_function.into_callable_type(db)); } } diff --git a/crates/ty_python_semantic/src/types/class/enum_literal.rs b/crates/ty_python_semantic/src/types/class/enum_literal.rs index a62b36e241222..7e355ee258145 100644 --- a/crates/ty_python_semantic/src/types/class/enum_literal.rs +++ b/crates/ty_python_semantic/src/types/class/enum_literal.rs @@ -10,6 +10,7 @@ use crate::types::Type; use crate::types::class::known::KnownClass; use crate::types::class::{ClassLiteral, ClassType, MemberLookupPolicy}; use crate::types::class_base::ClassBase; +use crate::types::enums::FlagBoundary; use crate::types::member::Member; use crate::types::mro::{DynamicMroError, Mro}; use ty_python_core::definition::Definition; @@ -21,6 +22,7 @@ pub struct EnumSpec<'db> { #[returns(deref)] pub(crate) members: Box<[(Name, Type<'db>)]>, pub(crate) has_known_members: bool, + pub(crate) boundary: FlagBoundary, } impl<'db> EnumSpec<'db> { @@ -40,7 +42,12 @@ impl<'db> EnumSpec<'db> { }) .collect::>>()?; - Some(Self::new(db, members, self.has_known_members(db))) + Some(Self::new( + db, + members, + self.has_known_members(db), + self.boundary(db), + )) } } diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 3adbaeda745f0..bba67a3fbfa78 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -1,3 +1,4 @@ +use itertools::Either; use ruff_db::parsed::parsed_module; use ruff_python_ast::name::Name; use rustc_hash::{FxHashMap, FxHashSet}; @@ -23,6 +24,14 @@ use crate::{ }; use ty_python_core::{definition::DefinitionKind, place_table, scope::ScopeId, use_def_map}; +mod flag; + +pub(crate) use flag::{ + FlagAutoValueState, FlagBoundary, FlagIteration, flag_binary_result, flag_constructor_result, + flag_invert_result, flag_literal_iteration, flag_literal_len, flag_literal_truthiness, + flag_membership_result, +}; + /// A resolved enum method, retaining both whether it is analyzable and who defines it. /// /// Standard-library methods and user-defined methods are both callable functions, but callers need @@ -153,6 +162,8 @@ pub(crate) struct EnumMetadata<'db> { pub(crate) members: FxIndexMap>, pub(crate) aliases: FxHashMap, + pub(crate) flag: Option>, + /// Members whose values were defined using `auto()`. pub(crate) auto_members: FxHashSet, @@ -198,10 +209,11 @@ pub(super) fn class_defines_property<'db>( false } -/// An enum class literal with its canonical members, value types, and aliases. +/// An enum class literal with its declared member objects, value types, and aliases. /// /// This keeps the enum-specific information used by enum literals and enum complements alongside -/// the underlying class literal. +/// the underlying class literal. For a `Flag`, zero-valued and multi-bit named members are retained +/// here so attribute lookup can resolve them even though iteration excludes them. #[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)] pub struct EnumClassLiteral<'db> { pub(crate) class_literal: ClassLiteral<'db>, @@ -209,7 +221,7 @@ pub struct EnumClassLiteral<'db> { pub(crate) members: Box<[(Name, Type<'db>)]>, #[returns(ref)] pub(crate) aliases: Box<[(Name, Name)]>, - /// Whether the canonical members exhaust the runtime values of this enum class. + /// Whether the declared member objects exhaust the runtime values of this enum class. /// /// `Flag` classes, transforming metaclasses, and enums with a custom `_missing_` method can /// create runtime members beyond those declared in the class body, so their declared members @@ -396,6 +408,7 @@ impl<'db> EnumMetadata<'db> { EnumMetadata { members: FxIndexMap::default(), aliases: FxHashMap::default(), + flag: None, auto_members: FxHashSet::default(), value_annotation: None, value_construction: EnumValueConstruction::default(), @@ -431,6 +444,10 @@ impl<'db> EnumMetadata<'db> { }; if let Some(EnumValueAnnotation::StandardLibrary(annotation)) = self.value_annotation + && !self + .flag + .as_ref() + .is_some_and(flag::FlagMetadata::preserves_assigned_value_type) && !known_constructor_preserves_value_type(db, value, annotation) { Some(annotation) @@ -487,6 +504,13 @@ impl<'db> EnumMetadata<'db> { if self.members.is_empty() { return None; } + if self.flag.is_some() { + return Some(UnionType::from_two_elements( + db, + KnownClass::Str.to_instance(db), + Type::none(db), + )); + } let union = self .members .keys() @@ -495,6 +519,16 @@ impl<'db> EnumMetadata<'db> { .build(); Some(union) } + + pub(crate) fn canonical_member_names(&self) -> Option> { + match &self.flag { + Some(flag) if flag.canonical_members_are_known() => Some(Either::Left( + flag.canonical_members().iter().map(|(name, _)| name), + )), + Some(_) => None, + None => Some(Either::Right(self.members.keys())), + } + } } /// A compact representation of an enum type with excluded members. @@ -805,6 +839,26 @@ fn try_register_alias<'db>( false } +/// Register aliases for a `Flag`, where integer values that compare equal identify the same +/// member even if their literal types differ (for example, `True` and `1`). +fn try_register_flag_alias( + value_ty: Type<'_>, + name: &Name, + flag_values: &mut FxHashMap, + aliases: &mut FxHashMap, +) -> bool { + let Some(value) = value_ty.as_int_like_literal() else { + return false; + }; + if let Some(canonical) = flag_values.get(&value) { + aliases.insert(name.clone(), canonical.clone()); + true + } else { + flag_values.insert(value, name.clone()); + false + } +} + /// List all members of an enum. #[salsa::tracked(returns(as_ref), cycle_initial=|_, _, _| Some(EnumMetadata::empty()), heap_size=ruff_memory_usage::heap_size)] pub(crate) fn enum_metadata<'db>( @@ -834,21 +888,38 @@ pub(crate) fn enum_metadata<'db>( let mut members = FxIndexMap::default(); let mut aliases = FxHashMap::default(); let mut enum_values: FxHashMap, Name> = FxHashMap::default(); + let is_flag = matches!( + enum_lit.base_class(db), + KnownClass::Flag | KnownClass::IntFlag + ); + let mut flag_values = FxHashMap::default(); for (name, ty) in spec.members(db) { - if try_register_alias(*ty, name, &mut enum_values, &mut aliases) { + if (is_flag && try_register_flag_alias(*ty, name, &mut flag_values, &mut aliases)) + || (!is_flag && try_register_alias(*ty, name, &mut enum_values, &mut aliases)) + { continue; } members.insert(name.clone(), *ty); } members.shrink_to_fit(); - return Some(EnumMetadata { + let mut metadata = EnumMetadata { members, aliases, + flag: None, auto_members: FxHashSet::default(), value_annotation: None, value_construction: EnumValueConstruction::default(), - }); + }; + if is_flag { + metadata.flag = Some(flag::FlagMetadata::from_enum_metadata( + db, + ClassLiteral::DynamicEnum(enum_lit), + &metadata, + spec.boundary(db), + )); + } + return Some(metadata); } }; @@ -865,10 +936,16 @@ pub(crate) fn enum_metadata<'db>( } let scope_id = class.body_scope(db); + let module = parsed_module(db, class.file(db)).load(db); let use_def_map = use_def_map(db, scope_id); let table = place_table(db, scope_id); let mut enum_values: FxHashMap, Name> = FxHashMap::default(); + let is_flag = Type::ClassLiteral(ClassLiteral::Static(class)) + .is_subtype_of(db, KnownClass::Flag.to_subclass_of(db)); + let mut flag_values: FxHashMap = FxHashMap::default(); + let mut flag_member_values: FxHashMap = FxHashMap::default(); + let mut flag_auto_values = FlagAutoValueState::new(); let mut auto_counter = 0; let mut auto_members = FxHashSet::default(); let mut prev_value_was_non_literal_int = false; @@ -923,10 +1000,30 @@ pub(crate) fn enum_metadata<'db>( return None; } + let flag_definition_and_expression = is_flag + .then(|| { + let mut definitions = bindings + .clone() + .filter_map(|binding| binding.binding.definition()); + let definition = definitions.next()?; + if definitions.next().is_some() { + return None; + } + let expression = match definition.kind(db) { + DefinitionKind::Assignment(assignment) => assignment.value(&module), + DefinitionKind::AnnotatedAssignment(assignment) => { + assignment.value(&module)? + } + _ => return None, + }; + Some((definition, expression)) + }) + .flatten(); + let inferred = place_from_bindings(db, bindings).place; let mut explicit_member_wrapper = false; - let value_ty = match inferred { + let mut value_ty = match inferred { Place::Undefined => { return None; } @@ -953,8 +1050,10 @@ pub(crate) fn enum_metadata<'db>( // enum.auto Some(KnownClass::Auto) => { - auto_counter += 1; - auto_members.insert(name.clone()); + if !is_flag { + auto_counter += 1; + auto_members.insert(name.clone()); + } // `StrEnum`s have different `auto()` behaviour to enums inheriting from `(str, Enum)` let auto_value_ty = @@ -962,6 +1061,8 @@ pub(crate) fn enum_metadata<'db>( .is_subtype_of(db, KnownClass::StrEnum.to_subclass_of(db)) { Type::string_literal(db, &name.to_lowercase()) + } else if is_flag { + flag_auto_values.next_value(db) } else { let custom_mixins: SmallVec<[Option; 1]> = class @@ -1038,10 +1139,55 @@ pub(crate) fn enum_metadata<'db>( } }; + if is_flag && explicit_member_wrapper && value_ty.is_instance_of(db, KnownClass::Auto) { + value_ty = flag_auto_values.next_value(db); + auto_members.insert(name.clone()); + } + + if let Some((definition, expression)) = flag_definition_and_expression { + let next_auto_value = value_construction + .alias_detection_value(db, flag_auto_values.next_value(db), true) + .and_then(Type::as_int_like_literal); + let (symbolic_value, is_direct_auto) = flag::flag_member_expression_value( + db, + definition, + expression, + &flag_member_values, + next_auto_value, + ); + if is_direct_auto { + auto_members.insert(name.clone()); + } + if let Some(symbolic_value) = symbolic_value + && value_ty.as_int_like_literal() != Some(symbolic_value) + { + value_ty = Type::int_literal(symbolic_value); + } + } + let alias_value_ty = value_construction.alias_detection_value(db, value_ty, auto_members.contains(name)); + if is_flag { + flag_auto_values.observe(alias_value_ty.unwrap_or_else(Type::any)); + if let Some(value) = alias_value_ty.and_then(Type::as_int_like_literal) { + flag_member_values.insert(name.clone(), value); + } + } if let Some(alias_value_ty) = alias_value_ty - && try_register_alias(alias_value_ty, name, &mut enum_values, &mut aliases) + && ((is_flag + && try_register_flag_alias( + alias_value_ty, + name, + &mut flag_values, + &mut aliases, + )) + || (!is_flag + && try_register_alias( + alias_value_ty, + name, + &mut enum_values, + &mut aliases, + ))) { return None; } @@ -1054,9 +1200,7 @@ pub(crate) fn enum_metadata<'db>( !matches!( declaration.kind(db), DefinitionKind::AnnotatedAssignment(assignment) - if assignment - .value(&parsed_module(db, declaration.file(db)).load(db)) - .is_some() + if assignment.value(&module).is_some() ) }) }) @@ -1097,13 +1241,23 @@ pub(crate) fn enum_metadata<'db>( members.shrink_to_fit(); - Some(EnumMetadata { + let mut metadata = EnumMetadata { members, aliases, + flag: None, auto_members, value_annotation, value_construction, - }) + }; + if is_flag { + metadata.flag = Some(flag::FlagMetadata::from_enum_metadata( + db, + ClassLiteral::Static(class), + &metadata, + flag::static_flag_boundary(db, class), + )); + } + Some(metadata) } /// Returns whether an enum's metaclass has a custom value-transforming method before the stdlib @@ -1133,6 +1287,27 @@ fn enum_metaclass_may_transform_values<'db>( }) } +/// Return whether enum construction uses the standard `EnumMeta.__call__` implementation. +pub(crate) fn enum_uses_standard_metaclass_call<'db>( + db: &'db dyn Db, + class: ClassLiteral<'db>, +) -> bool { + let ClassLiteral::Static(class) = class else { + return matches!(class, ClassLiteral::DynamicEnum(_)); + }; + let Some(metaclass) = class.metaclass(db).to_class_type(db) else { + return false; + }; + + !metaclass + .class_literal(db) + .iter_mro(db) + .filter_map(ClassBase::into_class) + .filter_map(|base| base.class_literal(db).as_static()) + .take_while(|base| base.known(db) != Some(KnownClass::EnumType)) + .any(|base| custom_enum_method(db, base.body_scope(db), "__call__").is_some()) +} + /// Iterates over parent enum classes in the MRO, skipping known enum /// infrastructure classes but including `IntEnum`, `Flag`, and `IntFlag` /// which declare `_value_` annotations that normally should be inherited. diff --git a/crates/ty_python_semantic/src/types/enums/flag.rs b/crates/ty_python_semantic/src/types/enums/flag.rs new file mode 100644 index 0000000000000..1fdac77c503ec --- /dev/null +++ b/crates/ty_python_semantic/src/types/enums/flag.rs @@ -0,0 +1,1078 @@ +use ruff_db::parsed::parsed_module; +use ruff_python_ast::{self as ast, PythonVersion, name::Name}; +use rustc_hash::FxHashMap; +use ty_module_resolver::{KnownModule, file_to_module}; + +use crate::Db; +use crate::types::call::CallArguments; +use crate::types::{ + ClassLiteral, EnumLiteralType, KnownClass, LiteralValueTypeKind, Program, Type, UnionType, + definition_expression_type, +}; +use ty_python_core::definition::Definition; + +use super::{ + EnumClassLiteral, EnumMetadata, custom_enum_method, enum_uses_standard_metaclass_call, +}; + +/// The policy used when a `Flag` constructor receives bits that are not declared by the class. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update, get_size2::GetSize)] +pub enum FlagBoundary { + Strict, + Conform, + Eject, + Keep, + Unknown, +} + +impl FlagBoundary { + pub(crate) const fn default_for_base(base: KnownClass) -> Self { + match base { + KnownClass::IntFlag => Self::Keep, + KnownClass::Flag => Self::Strict, + _ => Self::Unknown, + } + } + + /// Resolve a value of the standard-library `FlagBoundary` enum. + pub(crate) fn from_type(db: &dyn Db, ty: Type<'_>) -> Option { + let LiteralValueTypeKind::Enum(boundary) = + ty.resolve_type_alias(db).as_literal_value_kind()? + else { + return None; + }; + let ClassLiteral::Static(boundary_class) = boundary.enum_class(db) else { + return None; + }; + if boundary_class.name(db) != "FlagBoundary" + || file_to_module(db, boundary_class.file(db)).and_then(|module| module.known(db)) + != Some(KnownModule::Enum) + { + return None; + } + match boundary.name(db).as_str() { + "STRICT" => Some(Self::Strict), + "CONFORM" => Some(Self::Conform), + "EJECT" => Some(Self::Eject), + "KEEP" => Some(Self::Keep), + _ => None, + } + } +} + +bitflags::bitflags! { + #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, salsa::Update)] + struct FlagMetadataFlags: u8 { + const PRESERVES_ASSIGNED_VALUE_TYPE = 1 << 0; + const PRESERVES_NEGATIVE_VALUES = 1 << 1; + const USES_LEGACY_INVERSION = 1 << 2; + const CANONICAL_MEMBERS_ARE_KNOWN = 1 << 3; + const CANONICAL_MEMBERS_ARE_IN_VALUE_ORDER = 1 << 4; + } +} + +impl get_size2::GetSize for FlagMetadataFlags {} + +/// Cached integer semantics for a `Flag` class. +/// +/// A profile stores only masks and direct lookup tables. It never enumerates the combinations of +/// declared flags, so its size and construction cost are linear in the number of declared names. +#[derive(Debug, PartialEq, Eq, salsa::Update, get_size2::GetSize)] +pub(crate) struct FlagMetadata<'db> { + boundary: FlagBoundary, + member_type: Option>, + flags: FlagMetadataFlags, + member_values: FxHashMap, + named_values: FxHashMap, + canonical_members: Box<[(Name, i64)]>, + flag_mask: Option, + singles_mask: Option, +} + +struct FlagMemberType<'db> { + ty: Option>, + values_are_known: bool, +} + +fn flag_member_type<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> FlagMemberType<'db> { + match class { + ClassLiteral::Static(class) => { + for base in class + .iter_mro(db, None) + .skip(1) + .filter_map(crate::types::ClassBase::into_class) + .filter_map(|base| base.class_literal(db).as_static()) + { + match base.known(db) { + Some(KnownClass::Flag) => { + return FlagMemberType { + ty: None, + values_are_known: true, + }; + } + Some(KnownClass::IntFlag | KnownClass::Int) => { + return FlagMemberType { + ty: Some(KnownClass::Int.to_instance(db)), + values_are_known: true, + }; + } + _ if Type::ClassLiteral(ClassLiteral::Static(base)) + .is_subtype_of(db, KnownClass::Flag.to_subclass_of(db)) => {} + _ => { + let base = ClassLiteral::Static(base); + return FlagMemberType { + ty: Type::ClassLiteral(base) + .is_subtype_of(db, KnownClass::Int.to_subclass_of(db)) + .then(|| base.to_non_generic_instance(db)), + values_are_known: false, + }; + } + } + } + FlagMemberType { + ty: None, + values_are_known: false, + } + } + ClassLiteral::DynamicEnum(enum_lit) => match enum_lit.mixin_type(db) { + Some(mixin) + if mixin + .to_class_type(db) + .is_some_and(|mixin| mixin.known(db) == Some(KnownClass::Int)) => + { + FlagMemberType { + ty: Some(KnownClass::Int.to_instance(db)), + values_are_known: true, + } + } + Some(_) => FlagMemberType { + ty: None, + values_are_known: false, + }, + None => FlagMemberType { + ty: (enum_lit.base_class(db) == KnownClass::IntFlag) + .then(|| KnownClass::Int.to_instance(db)), + values_are_known: true, + }, + }, + ClassLiteral::Dynamic(_) + | ClassLiteral::DynamicNamedTuple(_) + | ClassLiteral::DynamicTypedDict(_) => FlagMemberType { + ty: None, + values_are_known: false, + }, + } +} + +impl<'db> FlagMetadata<'db> { + pub(super) fn from_enum_metadata( + db: &'db dyn Db, + class: ClassLiteral<'db>, + metadata: &EnumMetadata<'db>, + mut boundary: FlagBoundary, + ) -> Self { + let mut member_values = FxHashMap::default(); + let mut named_values = FxHashMap::default(); + let mut canonical_members = Vec::new(); + let mut flag_mask = 0_i64; + let mut singles_mask = 0_i64; + let member_type = flag_member_type(db, class); + let mut all_values_are_known = member_type.values_are_known; + let preserves_assigned_value_type = + member_type.ty.is_none() && member_type.values_are_known; + let before_python_311 = Program::get(db).python_version(db) < PythonVersion::PY311; + let is_int_flag = + Type::ClassLiteral(class).is_subtype_of(db, KnownClass::IntFlag.to_subclass_of(db)); + let preserves_negative_values = before_python_311 && is_int_flag; + let uses_legacy_inversion = before_python_311 && !is_int_flag; + + if member_type.values_are_known { + for name in metadata.members.keys() { + if metadata.member_value_may_be_transformed(name) { + all_values_are_known = false; + continue; + } + let value_ty = metadata.value_type(db, name); + let value = value_ty.and_then(Type::as_int_like_literal).or_else(|| { + if value_ty?.is_assignable_to(db, KnownClass::Int.to_instance(db)) { + metadata.members.get(name)?.as_int_like_literal() + } else { + None + } + }); + let Some(value) = value else { + all_values_are_known = false; + continue; + }; + + member_values.insert(name.clone(), value); + named_values.entry(value).or_insert_with(|| name.clone()); + flag_mask |= value; + if is_positive_power_of_two(value) { + singles_mask |= value; + } + if before_python_311 || is_positive_power_of_two(value) { + canonical_members.push((name.clone(), value)); + } + } + } + + if metadata.value_construction.metaclass_may_transform_values { + all_values_are_known = false; + boundary = FlagBoundary::Unknown; + } + + member_values.shrink_to_fit(); + named_values.shrink_to_fit(); + let canonical_members_are_in_value_order = canonical_members + .windows(2) + .all(|members| members[0].1 < members[1].1); + let mut flags = FlagMetadataFlags::empty(); + flags.set( + FlagMetadataFlags::PRESERVES_ASSIGNED_VALUE_TYPE, + preserves_assigned_value_type, + ); + flags.set( + FlagMetadataFlags::PRESERVES_NEGATIVE_VALUES, + preserves_negative_values, + ); + flags.set( + FlagMetadataFlags::USES_LEGACY_INVERSION, + uses_legacy_inversion, + ); + flags.set( + FlagMetadataFlags::CANONICAL_MEMBERS_ARE_KNOWN, + all_values_are_known, + ); + flags.set( + FlagMetadataFlags::CANONICAL_MEMBERS_ARE_IN_VALUE_ORDER, + canonical_members_are_in_value_order, + ); + + Self { + boundary, + member_type: member_type.ty, + flags, + member_values, + named_values, + canonical_members: canonical_members.into_boxed_slice(), + flag_mask: all_values_are_known.then_some(flag_mask), + singles_mask: all_values_are_known.then_some(singles_mask), + } + } + + pub(crate) const fn boundary(&self) -> FlagBoundary { + self.boundary + } + + pub(super) const fn preserves_assigned_value_type(&self) -> bool { + self.flags + .contains(FlagMetadataFlags::PRESERVES_ASSIGNED_VALUE_TYPE) + } + + fn accepts_operand(&self, db: &dyn Db, operand: Type<'db>) -> bool { + self.member_type + .is_some_and(|member_type| operand.is_assignable_to(db, member_type)) + } + + pub(crate) const fn canonical_members_are_known(&self) -> bool { + self.flags + .contains(FlagMetadataFlags::CANONICAL_MEMBERS_ARE_KNOWN) + } + + pub(crate) fn canonical_members(&self) -> &[(Name, i64)] { + &self.canonical_members + } + + const fn canonical_members_are_in_value_order(&self) -> bool { + self.flags + .contains(FlagMetadataFlags::CANONICAL_MEMBERS_ARE_IN_VALUE_ORDER) + } + + fn value_has_missing_single_bit_member(&self, value: i64) -> Option { + Some(value & (self.flag_mask? & !self.singles_mask?) != 0) + } + + fn member_value(&self, name: &Name) -> Option { + self.member_values.get(name).copied() + } + + fn named_member(&self, value: i64) -> Option<&Name> { + self.named_values.get(&value) + } + + fn all_bits(&self) -> Option { + let flag_mask = self.flag_mask?; + if flag_mask == 0 { + return Some(0); + } + let bit_length = i64::BITS - flag_mask.unsigned_abs().leading_zeros(); + if bit_length == 63 { + Some(i64::MAX) + } else { + 1_i64.checked_shl(bit_length).map(|value| value - 1) + } + } + + fn value_is_out_of_range(&self, value: i64) -> Option { + let flag_mask = self.flag_mask?; + let all_bits = self.all_bits()?; + Some(value < !all_bits || value > all_bits || value & (all_bits ^ flag_mask) != 0) + } + + fn normalize_negative(&self, value: i64) -> Option { + if value >= 0 { + Some(value) + } else { + i64::try_from(i128::from(self.all_bits()?) + 1 + i128::from(value)).ok() + } + } + + fn normalize_kept_negative(&self, value: i64) -> Option { + let all_bits = i128::from(self.all_bits()?); + let bit_length = i64::BITS - value.unsigned_abs().leading_zeros(); + let modulus = 1_i128.checked_shl(bit_length)?; + i64::try_from((all_bits + 1).max(modulus) + i128::from(value)).ok() + } + + fn construct(&self, value: i64) -> FlagConstruction { + // Enum value lookup returns an existing named member before `Flag._missing_` applies the + // boundary policy. This matters for explicitly declared negative values that would not + // otherwise satisfy the class's effective mask. + if self.named_member(value).is_some() { + return FlagConstruction::Flag(value); + } + if self + .flags + .contains(FlagMetadataFlags::PRESERVES_NEGATIVE_VALUES) + && value < 0 + { + return FlagConstruction::Flag(value); + } + + let Some(out_of_range) = self.value_is_out_of_range(value) else { + return FlagConstruction::Unknown; + }; + + let value = match (self.boundary, out_of_range) { + (FlagBoundary::Unknown, true) => return FlagConstruction::Unknown, + (FlagBoundary::Eject, true) => return FlagConstruction::Ejected(value), + (FlagBoundary::Conform, true) => value & self.flag_mask.unwrap_or_default(), + (FlagBoundary::Keep, true) if value < 0 => { + return self + .normalize_kept_negative(value) + .map_or(FlagConstruction::Unknown, FlagConstruction::Flag); + } + // A strict construction raises at runtime. The caller retains the nominal result type + // because ty does not generally use constructor exceptions to infer `Never`. + (FlagBoundary::Strict, true) => return FlagConstruction::Raises, + (FlagBoundary::Keep | FlagBoundary::Strict | FlagBoundary::Conform, _) + | (FlagBoundary::Unknown | FlagBoundary::Eject, false) => value, + }; + + self.normalize_negative(value) + .map_or(FlagConstruction::Unknown, FlagConstruction::Flag) + } + + fn legacy_invert(&self, value: i64) -> FlagConstruction { + if !self.canonical_members_are_known() { + return FlagConstruction::Unknown; + } + let inverted = self + .canonical_members + .iter() + .filter(|(_, member)| member & value == 0) + .fold(0, |inverted, (_, member)| inverted | member); + self.construct(inverted) + } +} + +#[derive(Clone, Copy)] +enum FlagConstruction { + Flag(i64), + Ejected(i64), + Raises, + Unknown, +} + +/// Values used by the standard `Flag._generate_next_value_` implementation. +#[derive(Clone, Copy, Debug)] +pub(crate) struct FlagAutoValueState { + maximum: Option, + last: Option, + all_values_are_known: bool, + has_values: bool, +} + +impl FlagAutoValueState { + pub(crate) const fn new() -> Self { + Self { + maximum: None, + last: None, + all_values_are_known: true, + has_values: false, + } + } + + pub(crate) fn observe(&mut self, value: Type<'_>) { + self.has_values = true; + if let Some(value) = value.as_int_like_literal() { + self.maximum = Some(self.maximum.map_or(value, |maximum| maximum.max(value))); + self.last = Some(value); + } else { + self.last = None; + self.all_values_are_known = false; + } + } + + pub(crate) fn next_value<'db>(&self, db: &'db dyn Db) -> Type<'db> { + if !self.has_values { + return Type::int_literal(1); + } + let value = if Program::get(db).python_version(db) < PythonVersion::PY311 { + self.last + } else { + self.maximum.filter(|_| self.all_values_are_known) + }; + let Some(value) = value else { + return KnownClass::Int.to_instance(db); + }; + let bit_length = i64::BITS - value.unsigned_abs().leading_zeros(); + 1_u64 + .checked_shl(bit_length) + .and_then(|value| i64::try_from(value).ok()) + .map(Type::int_literal) + .unwrap_or_else(|| KnownClass::Int.to_instance(db)) + } +} + +/// Evaluate the integer value of a Flag member expression using values established by earlier +/// members in the same class body. +/// +/// This is needed for expressions such as `READ_WRITE = READ | WRITE`: during ordinary class-body +/// inference, names assigned from `auto()` still have the placeholder `auto` type. +pub(crate) fn flag_member_expression_value<'db>( + db: &'db dyn Db, + definition: Definition<'db>, + expression: &ast::Expr, + previous_values: &FxHashMap, + auto_value: Option, +) -> (Option, bool) { + fn evaluate<'db>( + db: &'db dyn Db, + definition: Definition<'db>, + expression: &ast::Expr, + previous_values: &FxHashMap, + ) -> Option { + match expression { + ast::Expr::Name(name) => previous_values.get(&name.id).copied().or_else(|| { + definition_expression_type(db, definition, expression).as_int_like_literal() + }), + ast::Expr::BinOp(binary) => { + let left = evaluate(db, definition, &binary.left, previous_values)?; + let right = evaluate(db, definition, &binary.right, previous_values)?; + match binary.op { + ast::Operator::BitOr => Some(left | right), + ast::Operator::BitAnd => Some(left & right), + ast::Operator::BitXor => Some(left ^ right), + _ => { + definition_expression_type(db, definition, expression).as_int_like_literal() + } + } + } + ast::Expr::UnaryOp(unary) => { + let value = evaluate(db, definition, &unary.operand, previous_values)?; + match unary.op { + ast::UnaryOp::Invert => Some(!value), + ast::UnaryOp::UAdd => Some(value), + ast::UnaryOp::USub => value.checked_neg(), + ast::UnaryOp::Not => None, + } + } + _ => definition_expression_type(db, definition, expression).as_int_like_literal(), + } + } + + let is_direct_auto = matches!(expression, ast::Expr::Call(_)) + && definition_expression_type(db, definition, expression) + .is_instance_of(db, KnownClass::Auto); + if is_direct_auto { + return (auto_value, true); + } + + (evaluate(db, definition, expression, previous_values), false) +} + +/// Return the effective boundary inherited by a static `Flag` class. +pub(super) fn static_flag_boundary<'db>( + db: &'db dyn Db, + class: crate::types::StaticClassLiteral<'db>, +) -> FlagBoundary { + for base in class + .iter_mro(db, None) + .filter_map(crate::types::ClassBase::into_class) + .filter_map(|base| base.class_literal(db).as_static()) + { + if !Type::ClassLiteral(ClassLiteral::Static(base)) + .is_subtype_of(db, KnownClass::Flag.to_subclass_of(db)) + { + continue; + } + if let Some(boundary) = explicit_flag_boundary(db, base) { + return boundary; + } + match base.known(db) { + Some(KnownClass::IntFlag) => return FlagBoundary::Keep, + Some(KnownClass::Flag) => return FlagBoundary::Strict, + _ => {} + } + } + FlagBoundary::Unknown +} + +fn explicit_flag_boundary<'db>( + db: &'db dyn Db, + class: crate::types::StaticClassLiteral<'db>, +) -> Option { + if Program::get(db).python_version(db) < PythonVersion::PY311 { + return None; + } + let module = parsed_module(db, class.file(db)).load(db); + let definition = class.definition(db); + let keyword = definition + .kind(db) + .as_class()? + .node(&module) + .arguments + .as_ref()? + .find_keyword("boundary")?; + let ty = definition_expression_type(db, class.definition(db), &keyword.value); + if ty.is_none(db) { + None + } else { + Some(FlagBoundary::from_type(db, ty).unwrap_or(FlagBoundary::Unknown)) + } +} + +fn flag_operand<'db>( + db: &'db dyn Db, + ty: Type<'db>, +) -> Option<(EnumClassLiteral<'db>, Option>)> { + match ty.resolve_type_alias(db) { + Type::LiteralValue(literal) => { + let enum_literal = literal.as_enum()?; + let enum_class = enum_literal.enum_class_literal(db); + enum_metadata_and_flag(db, enum_class)?; + Some((enum_class, Some(enum_literal))) + } + Type::NominalInstance(instance) => { + let enum_class = instance.class_literal(db).into_enum_class(db)?; + enum_metadata_and_flag(db, enum_class)?; + Some((enum_class, None)) + } + Type::Intersection(intersection) => { + let mut flags = intersection + .positive(db) + .iter() + .filter_map(|positive| flag_operand(db, *positive)); + let (enum_class, mut literal) = flags.next()?; + for (other_class, other_literal) in flags { + if other_class != enum_class { + return None; + } + if literal != other_literal { + literal = None; + } + } + Some((enum_class, literal)) + } + _ => None, + } +} + +fn enum_metadata_and_flag<'db>( + db: &'db dyn Db, + enum_class: EnumClassLiteral<'db>, +) -> Option<(&'db EnumMetadata<'db>, &'db FlagMetadata<'db>)> { + let metadata = super::enum_metadata(db, enum_class.class_literal(db))?; + Some((metadata, metadata.flag.as_ref()?)) +} + +fn literal_value<'db>( + db: &'db dyn Db, + enum_class: EnumClassLiteral<'db>, + literal: EnumLiteralType<'db>, +) -> Option { + let (metadata, flag) = enum_metadata_and_flag(db, enum_class)?; + let name = metadata + .aliases + .get(literal.name(db)) + .unwrap_or(literal.name(db)); + flag.member_value(name) +} + +fn class_uses_standard_flag_method( + db: &dyn Db, + enum_class: EnumClassLiteral<'_>, + name: &str, +) -> bool { + let is_copied_onto_concrete_flag = Program::get(db).python_version(db) >= PythonVersion::PY311 + && matches!( + name, + "__or__" | "__and__" | "__xor__" | "__ror__" | "__rand__" | "__rxor__" | "__invert__" + ); + match enum_class.class_literal(db) { + ClassLiteral::Static(class) if is_copied_onto_concrete_flag => { + custom_enum_method(db, class.body_scope(db), name).is_none() + } + ClassLiteral::Static(class) => !class + .iter_mro(db, None) + .filter_map(crate::types::ClassBase::into_class) + .filter_map(|base| base.class_literal(db).as_static()) + .take_while(|base| { + !matches!(base.known(db), Some(KnownClass::Flag | KnownClass::IntFlag)) + }) + .any(|base| custom_enum_method(db, base.body_scope(db), name).is_some()), + ClassLiteral::DynamicEnum(_) if is_copied_onto_concrete_flag => true, + ClassLiteral::DynamicEnum(enum_lit) => enum_lit.mixin_type(db).is_none_or(|mixin| { + let Some(mixin) = mixin.to_class_type(db) else { + return false; + }; + !mixin + .iter_mro(db) + .filter_map(crate::types::ClassBase::into_class) + .filter_map(|base| base.class_literal(db).as_static()) + .any(|base| custom_enum_method(db, base.body_scope(db), name).is_some()) + }), + ClassLiteral::Dynamic(_) + | ClassLiteral::DynamicNamedTuple(_) + | ClassLiteral::DynamicTypedDict(_) => false, + } +} + +fn class_uses_standard_flag_construction(db: &dyn Db, enum_class: EnumClassLiteral<'_>) -> bool { + class_uses_standard_flag_method(db, enum_class, "_missing_") +} + +fn class_uses_standard_flag_operation( + db: &dyn Db, + enum_class: EnumClassLiteral<'_>, + method: &str, +) -> bool { + class_uses_standard_flag_method(db, enum_class, method) + && class_uses_standard_flag_method(db, enum_class, "_get_value") +} + +fn possible_flag_or_int<'db>(db: &'db dyn Db, enum_class: EnumClassLiteral<'db>) -> Type<'db> { + UnionType::from_two_elements( + db, + enum_class.class_literal(db).to_non_generic_instance(db), + KnownClass::Int.to_instance(db), + ) +} + +fn nonstandard_flag_construction_result<'db>( + db: &'db dyn Db, + enum_class: EnumClassLiteral<'db>, + argument: Type<'db>, +) -> Option> { + let class = enum_class.class_literal(db); + if !enum_uses_standard_metaclass_call(db, class) { + let arguments = CallArguments::positional([argument]); + return Some(match Type::ClassLiteral(class).try_call(db, &arguments) { + Ok(bindings) => bindings.return_type(db), + Err(error) => error.return_type(db), + }); + } + if !class_uses_standard_flag_construction(db, enum_class) { + let (_, flag) = enum_metadata_and_flag(db, enum_class)?; + return Some(construction_type( + db, + enum_class, + flag, + FlagConstruction::Unknown, + None, + )); + } + None +} + +fn construction_type<'db>( + db: &'db dyn Db, + enum_class: EnumClassLiteral<'db>, + flag: &FlagMetadata<'db>, + construction: FlagConstruction, + ejected_type: Option>, +) -> Type<'db> { + let nominal = enum_class.class_literal(db).to_non_generic_instance(db); + match construction { + FlagConstruction::Flag(value) => flag.named_member(value).map_or(nominal, |name| { + Type::enum_literal(EnumLiteralType::new(db, enum_class, name.clone())) + }), + FlagConstruction::Ejected(value) => { + ejected_type.unwrap_or_else(|| Type::int_literal(value)) + } + FlagConstruction::Unknown + if matches!(flag.boundary(), FlagBoundary::Eject | FlagBoundary::Unknown) => + { + possible_flag_or_int(db, enum_class) + } + FlagConstruction::Raises | FlagConstruction::Unknown => nominal, + } +} + +fn flag_integer_binary_result<'db>( + db: &'db dyn Db, + enum_class: EnumClassLiteral<'db>, + literal: Option>, + operand: Type<'db>, + method: &str, + operation: fn(i64, i64) -> i64, +) -> Option> { + let (_, flag) = enum_metadata_and_flag(db, enum_class)?; + if !flag.accepts_operand(db, operand) + || !class_uses_standard_flag_operation(db, enum_class, method) + { + return None; + } + let exact = literal + .and_then(|literal| literal_value(db, enum_class, literal)) + .zip(operand.as_int_like_literal()) + .map(|(left, right)| operation(left, right)); + if let Some(result) = nonstandard_flag_construction_result( + db, + enum_class, + exact + .map(Type::int_literal) + .unwrap_or_else(|| KnownClass::Int.to_instance(db)), + ) { + return Some(result); + } + + Some(match exact { + Some(value) => construction_type(db, enum_class, flag, flag.construct(value), None), + None => construction_type(db, enum_class, flag, FlagConstruction::Unknown, None), + }) +} + +fn is_builtin_int_operand(db: &dyn Db, ty: Type<'_>) -> bool { + match ty.resolve_type_alias(db) { + Type::LiteralValue(literal) => literal.as_int().is_some(), + Type::NominalInstance(instance) => instance.known_class(db) == Some(KnownClass::Int), + _ => false, + } +} + +/// Infer a standard-library Flag bitwise operation without expanding possible combinations. +pub(crate) fn flag_binary_result<'db>( + db: &'db dyn Db, + left: Type<'db>, + right: Type<'db>, + op: ruff_python_ast::Operator, +) -> Option> { + let operation: fn(i64, i64) -> i64 = match op { + ruff_python_ast::Operator::BitOr => |left, right| left | right, + ruff_python_ast::Operator::BitAnd => |left, right| left & right, + ruff_python_ast::Operator::BitXor => |left, right| left ^ right, + _ => return None, + }; + + let left_flag = flag_operand(db, left); + let right_flag = flag_operand(db, right); + match (left_flag, right_flag) { + (Some((left_class, left_literal)), Some((right_class, right_literal))) + if left_class == right_class => + { + if !class_uses_standard_flag_operation(db, left_class, op.dunder()) { + return None; + } + let value = left_literal.zip(right_literal).and_then(|(left, right)| { + Some(operation( + literal_value(db, left_class, left)?, + literal_value(db, right_class, right)?, + )) + }); + if let Some(result) = nonstandard_flag_construction_result( + db, + left_class, + value + .map(Type::int_literal) + .unwrap_or_else(|| KnownClass::Int.to_instance(db)), + ) { + return Some(result); + } + let (_, flag) = enum_metadata_and_flag(db, left_class)?; + Some(value.map_or_else( + || left_class.class_literal(db).to_non_generic_instance(db), + |value| { + construction_type( + db, + left_class, + flag, + flag.construct(value), + Some(Type::int_literal(value)), + ) + }, + )) + } + (Some((left_class, left_literal)), Some((right_class, right_literal))) => { + let (_, left_flag) = enum_metadata_and_flag(db, left_class)?; + if !left_flag.accepts_operand(db, right) + || !class_uses_standard_flag_operation(db, left_class, op.dunder()) + { + return None; + } + + let exact = left_literal.zip(right_literal).and_then(|(left, right)| { + if !class_uses_standard_flag_operation(db, right_class, op.reflected_dunder()) + || !class_uses_standard_flag_construction(db, right_class) + || !enum_uses_standard_metaclass_call(db, right_class.class_literal(db)) + { + return None; + } + let left = literal_value(db, left_class, left)?; + let right = literal_value(db, right_class, right)?; + let (_, right_flag) = enum_metadata_and_flag(db, right_class)?; + let raw_value = operation(left, right); + let construction = right_flag.construct(raw_value); + let value = match construction { + FlagConstruction::Flag(value) | FlagConstruction::Ejected(value) => value, + FlagConstruction::Raises | FlagConstruction::Unknown => return None, + }; + Some(( + value, + construction_type( + db, + right_class, + right_flag, + construction, + Some(Type::int_literal(raw_value)), + ), + )) + }); + if let Some(result) = nonstandard_flag_construction_result( + db, + left_class, + exact + .map(|(_, ty)| ty) + .unwrap_or_else(|| KnownClass::Int.to_instance(db)), + ) { + return Some(result); + } + let (exact_value, ejected_type) = exact.unzip(); + Some(match exact_value { + Some(value) => construction_type( + db, + left_class, + left_flag, + left_flag.construct(value), + ejected_type, + ), + None => { + construction_type(db, left_class, left_flag, FlagConstruction::Unknown, None) + } + }) + } + (Some((enum_class, literal)), None) | (None, Some((enum_class, literal))) => { + let flag_on_left = left_flag.is_some(); + let integer = if flag_on_left { right } else { left }; + let method = if flag_on_left { + op.dunder() + } else { + op.reflected_dunder() + }; + let standard_dispatch = flag_on_left || is_builtin_int_operand(db, integer); + if !standard_dispatch { + return None; + } + flag_integer_binary_result(db, enum_class, literal, integer, method, operation) + } + _ => None, + } +} + +/// Infer `~flag` for the standard `Flag.__invert__` implementation. +pub(crate) fn flag_invert_result<'db>(db: &'db dyn Db, operand: Type<'db>) -> Option> { + let (enum_class, literal) = flag_operand(db, operand)?; + if !class_uses_standard_flag_operation(db, enum_class, "__invert__") { + return None; + } + let (_, flag) = enum_metadata_and_flag(db, enum_class)?; + let value = literal.and_then(|literal| literal_value(db, enum_class, literal)); + if let Some(result) = nonstandard_flag_construction_result( + db, + enum_class, + value + .map(|value| Type::int_literal(!value)) + .unwrap_or_else(|| KnownClass::Int.to_instance(db)), + ) { + return Some(result); + } + let Some(value) = value else { + return Some(construction_type( + db, + enum_class, + flag, + FlagConstruction::Unknown, + None, + )); + }; + let construction = if flag + .flags + .contains(FlagMetadataFlags::USES_LEGACY_INVERSION) + { + flag.legacy_invert(value) + } else { + match flag.boundary() { + FlagBoundary::Strict | FlagBoundary::Conform => flag + .singles_mask + .map(|mask| FlagConstruction::Flag(mask & !value)) + .unwrap_or(FlagConstruction::Unknown), + FlagBoundary::Eject | FlagBoundary::Keep => flag.construct(!value), + FlagBoundary::Unknown => FlagConstruction::Unknown, + } + }; + Some(construction_type( + db, + enum_class, + flag, + construction, + Some(Type::int_literal(!value)), + )) +} + +/// Adjust the return type of a call to a concrete `Flag` class for its boundary policy. +pub(crate) fn flag_constructor_result<'db>( + db: &'db dyn Db, + class: ClassLiteral<'db>, + argument: Option>, +) -> Option> { + let enum_class = class.into_enum_class(db)?; + let (_, flag) = enum_metadata_and_flag(db, enum_class)?; + if !enum_uses_standard_metaclass_call(db, class) { + return None; + } + if let Some((argument_class, literal)) = argument.and_then(|ty| flag_operand(db, ty)) + && argument_class == enum_class + { + return Some(literal.map_or_else(|| class.to_non_generic_instance(db), Type::enum_literal)); + } + if !class_uses_standard_flag_construction(db, enum_class) { + return Some(construction_type( + db, + enum_class, + flag, + FlagConstruction::Unknown, + None, + )); + } + Some(match argument.and_then(Type::as_int_like_literal) { + Some(value) => construction_type(db, enum_class, flag, flag.construct(value), argument), + None => construction_type(db, enum_class, flag, FlagConstruction::Unknown, None), + }) +} + +/// Return the truthiness of an exact Flag member when the standard `__bool__` is in use. +pub(crate) fn flag_literal_truthiness(db: &dyn Db, literal: EnumLiteralType<'_>) -> Option { + let enum_class = literal.enum_class_literal(db); + class_uses_standard_flag_method(db, enum_class, "__bool__") + .then(|| literal_value(db, enum_class, literal).map(|value| value != 0)) + .flatten() +} + +/// Return the length of an exact Flag member when the standard `__len__` is in use. +pub(crate) fn flag_literal_len(db: &dyn Db, literal: EnumLiteralType<'_>) -> Option { + if Program::get(db).python_version(db) < PythonVersion::PY311 { + return None; + } + let enum_class = literal.enum_class_literal(db); + class_uses_standard_flag_method(db, enum_class, "__len__") + .then(|| { + literal_value(db, enum_class, literal) + .map(|value| i64::from(value.unsigned_abs().count_ones())) + }) + .flatten() +} + +/// Return the constituent canonical members yielded by iterating an exact Flag member. +pub(crate) fn flag_literal_iteration<'db>( + db: &'db dyn Db, + literal: EnumLiteralType<'db>, +) -> Option> { + if Program::get(db).python_version(db) < PythonVersion::PY311 { + return None; + } + let enum_class = literal.enum_class_literal(db); + if !class_uses_standard_flag_method(db, enum_class, "__iter__") { + return None; + } + let (_, flag) = enum_metadata_and_flag(db, enum_class)?; + let nominal = enum_class.class_literal(db).to_non_generic_instance(db); + if !class_uses_standard_flag_method(db, enum_class, "_iter_member_") + || !flag.canonical_members_are_known() + { + return Some(FlagIteration::Unknown(nominal)); + } + let generated_iteration_hook_is_custom = !flag.canonical_members_are_in_value_order() + && ["_iter_member_by_def_", "_iter_member_by_value_"] + .into_iter() + .any(|name| !class_uses_standard_flag_method(db, enum_class, name)); + if generated_iteration_hook_is_custom { + return Some(FlagIteration::Unknown(nominal)); + } + let Some(value) = literal_value(db, enum_class, literal) else { + return Some(FlagIteration::Unknown(nominal)); + }; + if value < 0 { + return Some(FlagIteration::Unknown(nominal)); + } + if flag + .value_has_missing_single_bit_member(value) + .unwrap_or(true) + { + return Some(FlagIteration::Unknown(UnionType::from_two_elements( + db, + nominal, + Type::none(db), + ))); + } + Some(FlagIteration::Exact( + flag.canonical_members() + .iter() + .filter(|(_, member)| value & member != 0) + .map(|(name, _)| name) + .map(|name| Type::enum_literal(EnumLiteralType::new(db, enum_class, name.clone()))) + .collect(), + )) +} + +pub(crate) enum FlagIteration<'db> { + Exact(Vec>), + Unknown(Type<'db>), +} + +const fn is_positive_power_of_two(value: i64) -> bool { + value > 0 && value & (value - 1) == 0 +} + +/// Evaluate exact Flag subset membership (`member in flags`). +pub(crate) fn flag_membership_result( + db: &dyn Db, + member: Type<'_>, + flags: Type<'_>, +) -> Option { + let (member_class, Some(member)) = flag_operand(db, member)? else { + return None; + }; + let (flags_class, Some(flags)) = flag_operand(db, flags)? else { + return None; + }; + if member_class != flags_class + || !class_uses_standard_flag_method(db, flags_class, "__contains__") + { + return None; + } + let member = literal_value(db, member_class, member)?; + let flags = literal_value(db, flags_class, flags)?; + Some(flags & member == member) +} diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index c1b65c192fe65..921b3e2ea60c2 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -68,7 +68,9 @@ use crate::types::diagnostic::{ report_too_many_positional_patterns_for_callable_class_pattern, report_unsupported_augmented_assignment, report_unsupported_comparison, }; -use crate::types::enums::{enum_ignored_names, is_enum_class_by_inheritance}; +use crate::types::enums::{ + enum_ignored_names, flag_constructor_result, flag_invert_result, is_enum_class_by_inheritance, +}; use crate::types::function::{ FunctionDecorators, FunctionType, KnownFunction, report_revealed_type, same_module_uncached_raw_signature, @@ -8929,6 +8931,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let db = self.db(); let scope = self.scope(); let return_ty = bindings.return_type(db); + let return_ty = callable_type + .as_class_literal() + .and_then(|class| { + flag_constructor_result( + db, + class, + arguments + .find_argument_value("value", 0) + .map(|argument| self.expression_type(argument)), + ) + }) + .unwrap_or(return_ty); let find_narrowed_place = |argument_index: usize| match arguments.args.get(argument_index) { None => { @@ -10348,6 +10362,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { self.infer_unary_expression_type(op, alias.value_type(self.db()), unary) } + (ast::UnaryOp::Invert, ty) if let Some(result) = flag_invert_result(self.db(), ty) => { + result + } + (ast::UnaryOp::UAdd, Type::LiteralValue(literal)) => match literal.kind() { LiteralValueTypeKind::Int(value) => Type::int_literal(value.as_i64()), LiteralValueTypeKind::Bool(value) => Type::int_literal(i64::from(value)), diff --git a/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs b/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs index 642b7d857c533..7e04becfc606b 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs @@ -8,6 +8,7 @@ use crate::types::cyclic::CycleDetector; use crate::types::diagnostic::{ DIVISION_BY_ZERO, report_unsupported_augmented_assignment, report_unsupported_binary_operation, }; +use crate::types::enums::flag_binary_result; use crate::types::set_theoretic::RecursivelyDefined; use crate::types::typevar::TypeVarConstraints; use crate::types::{ @@ -548,6 +549,12 @@ impl<'db> TypeInferenceBuilder<'db, '_> { (Type::Never, _, _) | (_, Type::Never, _) => Some(Type::Never), + ( + left, + right, + op @ (ast::Operator::BitOr | ast::Operator::BitAnd | ast::Operator::BitXor), + ) if let Some(result) = flag_binary_result(db, left, right, op) => Some(result), + (Type::LiteralValue(left), Type::LiteralValue(right), _) => { let recursively_defined = if left.recursively_defined().is_yes() || right.recursively_defined().is_yes() diff --git a/crates/ty_python_semantic/src/types/infer/builder/enum_call.rs b/crates/ty_python_semantic/src/types/infer/builder/enum_call.rs index 88bb549453d61..12f43ff4b91ed 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/enum_call.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/enum_call.rs @@ -14,6 +14,7 @@ use crate::{ INVALID_ARGUMENT_TYPE, INVALID_BASE, MISSING_ARGUMENT, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS, UNKNOWN_ARGUMENT, report_mismatched_type_name, }, + enums::{FlagAutoValueState, FlagBoundary}, infer::TypeInferenceBuilder, infer::builder::dynamic_class::report_mro_error_kind, subclass_of::SubclassOfType, @@ -186,32 +187,18 @@ fn next_auto_value<'db>( base_class: KnownClass, name: &str, last_int_value: Option, + flag_values: &FlagAutoValueState, ) -> Type<'db> { match base_class { KnownClass::StrEnum => Type::string_literal(db, &name.to_lowercase()), + KnownClass::Flag | KnownClass::IntFlag => flag_values.next_value(db), _ => { let Some(last) = last_int_value else { return KnownClass::Int.to_instance(db); }; - match base_class { - KnownClass::Flag | KnownClass::IntFlag => { - // next power of two after highest bit in the last value - if last <= 0 { - Type::int_literal(1) - } else { - let shift = i64::BITS - last.leading_zeros(); - 1_u64 - .checked_shl(shift) - .and_then(|value| i64::try_from(value).ok()) - .map(Type::int_literal) - .unwrap_or_else(|| KnownClass::Int.to_instance(db)) - } - } - _ => last - .checked_add(1) - .map(Type::int_literal) - .unwrap_or_else(|| KnownClass::Int.to_instance(db)), - } + last.checked_add(1) + .map(Type::int_literal) + .unwrap_or_else(|| KnownClass::Int.to_instance(db)) } } } @@ -224,14 +211,16 @@ fn enum_members_from_names( ) -> Vec<(Name, Type<'_>)> { let mut members = Vec::with_capacity(names.len()); let mut last_int_value = None; + let mut flag_values = FlagAutoValueState::new(); for (index, name) in names.into_iter().enumerate() { let value = if index == 0 { first_enum_auto_value(db, base_class, name.as_str(), start) } else { - next_auto_value(db, base_class, name.as_str(), last_int_value) + next_auto_value(db, base_class, name.as_str(), last_int_value, &flag_values) }; last_int_value = value.as_int_literal(); + flag_values.observe(value); members.push((name, value)); } @@ -332,6 +321,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let names_kw = call_expr.arguments.find_keyword("names"); let start_kw = call_expr.arguments.find_keyword("start"); let type_kw = call_expr.arguments.find_keyword("type"); + let boundary_kw = call_expr.arguments.find_keyword("boundary"); let has_positional_keyword_conflict = (!args.is_empty() && value_kw.is_some()) || (args.len() >= 2 && names_kw.is_some()); @@ -402,6 +392,21 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let (mixin_type, valid_mixin_type) = type_kw.map_or((None, true), |kw| { self.infer_enum_mixin_argument(&kw.value, base_class) }); + let boundary = if matches!(base_class, KnownClass::Flag | KnownClass::IntFlag) { + boundary_kw.map_or_else( + || FlagBoundary::default_for_base(base_class), + |keyword| { + let ty = self.expression_type(&keyword.value); + if ty.is_none(db) { + FlagBoundary::default_for_base(base_class) + } else { + FlagBoundary::from_type(db, ty).unwrap_or(FlagBoundary::Unknown) + } + }, + ) + } else { + FlagBoundary::Unknown + }; // Only 1 extra positional arg is allowed (the `names` parameter). // `Enum("Color", "RED", "GREEN")` is invalid at runtime. @@ -443,6 +448,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { base_class, mixin_type, has_too_many_positional || has_positional_keyword_conflict || !valid_mixin_type, + boundary, ); // Non-literal names use the ordinary `type[EnumSubclass]` overload result @@ -573,6 +579,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { base_class: KnownClass, mixin_type: Option>, has_invalid_arguments: bool, + boundary: FlagBoundary, ) -> EnumSpec<'db> { let db = self.db(); let (members, has_known_members) = if has_invalid_arguments { @@ -622,7 +629,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { } }; - EnumSpec::new(db, members.into_boxed_slice(), has_known_members) + EnumSpec::new(db, members.into_boxed_slice(), has_known_members, boundary) } fn create_dynamic_enum_anchor( @@ -773,13 +780,15 @@ impl<'db> TypeInferenceBuilder<'db, '_> { // Explicit-value forms match static enums: `start` is ignored and a leading `auto()` // still begins from the default seed of `1`. let mut last_int_value = Some(0); + let mut flag_values = FlagAutoValueState::new(); for (name, value) in explicit_members { let value = if value.is_instance_of(db, KnownClass::Auto) { - next_auto_value(db, base_class, name.as_str(), last_int_value) + next_auto_value(db, base_class, name.as_str(), last_int_value, &flag_values) } else { value }; last_int_value = value.as_int_like_literal(); + flag_values.observe(value); members.push((name, value)); } EnumMembersArgParseResult::Known(KnownEnumMembers { @@ -801,6 +810,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let db = self.db(); let mut members = Vec::with_capacity(dict.items.len()); let mut last_int_value = Some(0); + let mut flag_values = FlagAutoValueState::new(); let mut has_opaque_keys = false; for item in &dict.items { let Some(key) = &item.key else { @@ -819,11 +829,12 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let name = Name::new(string_lit.value(db)); let raw_value = self.expression_type(&item.value); let value = if raw_value.is_instance_of(db, KnownClass::Auto) { - next_auto_value(db, base_class, name.as_str(), last_int_value) + next_auto_value(db, base_class, name.as_str(), last_int_value, &flag_values) } else { raw_value }; last_int_value = value.as_int_like_literal(); + flag_values.observe(value); members.push((name, value)); } if has_opaque_keys { diff --git a/crates/ty_python_semantic/src/types/infer/comparisons.rs b/crates/ty_python_semantic/src/types/infer/comparisons.rs index adc1551ca6039..5c4eaa220d0a7 100644 --- a/crates/ty_python_semantic/src/types/infer/comparisons.rs +++ b/crates/ty_python_semantic/src/types/infer/comparisons.rs @@ -7,6 +7,7 @@ use crate::types::call::{CallArguments, CallDunderError}; use crate::types::constraints::ConstraintSetBuilder; use crate::types::context::InferContext; use crate::types::cyclic::CycleDetector; +use crate::types::enums::flag_membership_result; use crate::types::tuple::TupleSpec; use crate::types::{ DynamicType, IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, @@ -913,6 +914,12 @@ fn infer_membership_test_comparison<'db>( range: TextRange, ) -> Result, UnsupportedComparisonError<'db>> { let db = context.db(); + if let Some(result) = flag_membership_result(db, left, right) { + return Ok(Type::bool_literal(match op { + MembershipTestCompareOperator::In => result, + MembershipTestCompareOperator::NotIn => !result, + })); + } let compare_result_opt = match right.try_call_dunder( db, "__contains__", diff --git a/crates/ty_python_semantic/src/types/iteration.rs b/crates/ty_python_semantic/src/types/iteration.rs index 1dcc7fb331aef..e0c7b32bf9ecf 100644 --- a/crates/ty_python_semantic/src/types/iteration.rs +++ b/crates/ty_python_semantic/src/types/iteration.rs @@ -7,6 +7,7 @@ use crate::{ call::CallErrorKind, context::InferContext, diagnostic::NOT_ITERABLE, + enums::{FlagIteration, flag_literal_iteration}, todo_type, tuple::{TupleSpec, TupleSpecBuilder}, }, @@ -143,6 +144,16 @@ impl<'db> Type<'db> { LiteralValueTypeKind::LiteralString => { Some(Cow::Owned(TupleSpec::homogeneous(ty))) } + LiteralValueTypeKind::Enum(literal) => { + flag_literal_iteration(db, literal).map(|iteration| { + Cow::Owned(match iteration { + FlagIteration::Exact(members) => { + TupleSpec::heterogeneous(members) + } + FlagIteration::Unknown(member) => TupleSpec::homogeneous(member), + }) + }) + } _ => None } Type::Never => {