fix(enums): reject incompatible Enum value-lookup arguments

pyrefly/lib/alt/call.rs

@@ -1120,6 +1120,237 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         )
     }
 
+    /// `EnumMeta.__call__` types its `value` parameter as `Any`. For a plain value lookup
+    /// (`MyEnum(x)` / `MyEnum(value=x)`), return the type the lookup argument should be
+    /// checked against, or `None` when we cannot soundly refine it.
+    ///
+    /// Value lookup goes `EnumMeta.__call__` -> `Enum.__new__` -> existing member,
+    /// `_missing_(value)`, or `ValueError`. The enum's own `__init__`/`__new__` only run at
+    /// class-definition time, so they are deliberately not consulted here. We bail out when
+    /// the call contract may differ from stock `EnumMeta`: a custom metaclass, a user-defined
+    /// `_missing_`, the functional API, or an `Any` value type (e.g. a `__new__` that
+    /// rewrites `_value_` without a data-type mixin).
+    fn enum_value_lookup_expected_arg_type(
+        &self,
+        cls: &ClassType,
+        constructor_kind: &ConstructorKind,
+        keywords: &[CallKeyword],
+    ) -> Option<Type> {
+        // Do not restrict lookup on `type[Self]` calls, or on extensible memberless enum bases
+        // where subclasses can define members or dynamic lookup contracts like `_missing_`.
+        if *constructor_kind == ConstructorKind::TypeOfSelf
+            || self.get_enum_members(cls.class_object()).is_empty()
+        {
+            return None;
+        }
+        if keywords
+            .iter()
+            .any(|kw| kw.arg.is_none_or(|id| id.as_str() != "value"))
+        {
+            return None;
+        }
+        let metadata = self.get_metadata_for_class(cls.class_object());
+        let enum_metadata = metadata.enum_metadata()?;
+        let enum_module = ModuleName::enum_();
+        let is_stock_metaclass = metadata.custom_metaclass().is_some_and(|m| {
+            m.class_object()
+                .has_toplevel_qname(enum_module.as_str(), "EnumMeta")
+                || m.class_object()
+                    .has_toplevel_qname(enum_module.as_str(), "EnumType")
+        });
+        if !is_stock_metaclass {
+            return None;
+        }
+        if let Some(missing) = self.get_class_member_with_defining_class(
+            cls.class_object(),
+            &Name::new_static("_missing_"),
+        ) && missing.defining_class.module_name() != enum_module
+        {
+            return None;
+        }
+        // A user-defined `__new__` can store a `_value_` of any type, regardless of mixin.
+        if let Some(new) =
+            self.get_class_member_with_defining_class(cls.class_object(), &dunder::NEW)
+            && new.defining_class.module_name() != enum_module
+            && new.defining_class.module_name() != ModuleName::builtins()
+        {
+            return None;
+        }
+        let value_ty = self.enum_value_lookup_on_class(cls, enum_metadata);
+        // Stub enums often write `X = ...`; that is a placeholder, not the real value type.
+        let mut has_ellipsis = false;
+        value_ty.universe(&mut |t| has_ellipsis |= t.is_ellipsis_value());
+        if value_ty.is_any() || has_ellipsis {
+            return None;
+        }
+        // For containers (tuples, lists, etc.), cross-numeric equality across nested elements
+        // makes strict nominal lookup checking unsound without complex contextual inference.
+        // Bail out conservatively rather than rejecting equality-compatible container lookups.
+        let mut has_container = false;
+        value_ty.universe(&mut |t| {
+            has_container |= matches!(t, Type::Tuple(..))
+                || matches!(t, Type::ClassType(c) if {
+                    let obj = c.class_object();
+                    obj == self.stdlib.tuple_object()
+                        || obj == self.stdlib.list_object()
+                        || obj == self.stdlib.set_object()
+                        || obj == self.stdlib.frozenset_object()
+                        || obj == self.stdlib.dict_object()
+                });
+        });
+        if has_container {
+            return None;
+        }
+        // Member literals say which values exist, but lookup accepts any value of the base
+        // type and raises at runtime when it is missing, so promote literals.
+        let promoted = value_ty.promote_implicit_literals(self.stdlib);
+        let mut options = vec![promoted.clone(), Type::ClassType(cls.clone())];
+        // Lookup uses equality, so numeric values compare equal across int/float/complex
+        // (`Color(1.0)` finds a member whose value is `1`, `B(1)` finds `True`).
+        let int = self.heap.mk_class_type(self.stdlib.int().clone());
+        let float = self.heap.mk_class_type(self.stdlib.float().clone());
+        let complex = self.heap.mk_class_type(self.stdlib.complex().clone());
+        let is_numeric = self.is_subset_eq(&promoted, &int)
+            || self.is_subset_eq(&int, &promoted)
+            || self.is_subset_eq(&promoted, &float)
+            || self.is_subset_eq(&float, &promoted)
+            || self.is_subset_eq(&promoted, &complex)
+            || self.is_subset_eq(&complex, &promoted);
+        if is_numeric {
+            options.push(int);
+            options.push(float);
+            options.push(complex);
+        }
+        Some(self.unions(options))
+    }
+
+    fn is_numeric_equality_compatible(&self, ty: &Type) -> bool {
+        match ty {
+            Type::None
+            | Type::LiteralString(..)
+            | Type::Tuple(..)
+            | Type::TypedDict(..)
+            | Type::PartialTypedDict(..) => false,
+            Type::Literal(lit) => match &lit.value {
+                Lit::Str(..) | Lit::Bytes(..) => false,
+                Lit::Int(..) | Lit::Bool(..) => true,
+                Lit::Enum(..) => false,
+            },
+            Type::ClassType(cls) => {
+                let obj = cls.class_object();
+                if obj == self.stdlib.str().class_object()
+                    || obj == self.stdlib.bytes().class_object()
+                    || obj == self.stdlib.bytearray().class_object()
+                    || obj == self.stdlib.tuple_object()
+                    || obj == self.stdlib.list_object()
+                    || obj == self.stdlib.set_object()
+                    || obj == self.stdlib.frozenset_object()
+                    || obj == self.stdlib.dict_object()
+                    || obj == self.stdlib.mapping_object()
+                    || obj == self.stdlib.none_type().class_object()
+                    || obj == self.stdlib.object().class_object()
+                {
+                    return false;
+                }
+                if self.get_metadata_for_class(obj).is_enum()
+                    && !self.is_subset_eq(ty, &self.heap.mk_class_type(self.stdlib.int().clone()))
+                {
+                    return false;
+                }
+                true
+            }
+            Type::Union(u) => u
+                .members
+                .iter()
+                .all(|t| self.is_numeric_equality_compatible(t)),
+            Type::Intersect(inner) => {
+                let (tys, _) = inner.as_ref();
+                tys.iter().any(|t| self.is_numeric_equality_compatible(t))
+            }
+            _ => false,
+        }
+    }
+
+    fn is_enum_lookup_compatible(&self, got: &Type, want: &Type) -> bool {
+        if self.is_subset_eq(got, want) {
+            return true;
+        }
+        // Arguments typed `object` or `Any` are accepted because enum lookup
+        // performs equality comparison on arbitrary objects at runtime.
+        if got.is_object() || got.is_any() {
+            return true;
+        }
+        let int = self.heap.mk_class_type(self.stdlib.int().clone());
+        let is_numeric_target = self.is_subset_eq(&int, want);
+        if is_numeric_target && self.is_numeric_equality_compatible(got) {
+            return true;
+        }
+        match got {
+            Type::Union(u) => u
+                .members
+                .iter()
+                .all(|t| self.is_enum_lookup_compatible(t, want)),
+            Type::Intersect(inner) => {
+                let (tys, _) = inner.as_ref();
+                tys.iter().any(|t| self.is_enum_lookup_compatible(t, want))
+            }
+            _ => false,
+        }
+    }
+
+    fn check_enum_lookup_arg_compatibility(
+        &self,
+        got: &Type,
+        want: &Type,
+        range: TextRange,
+        errors: &ErrorCollector,
+        context: Option<&dyn Fn() -> ErrorContext>,
+    ) {
+        if self.is_enum_lookup_compatible(got, want) {
+            return;
+        }
+        self.check_type_as_call_argument(got, want, range, errors, &|| {
+            let mut tcc = TypeCheckContext::of_kind(TypeCheckKind::CallArgument(
+                Some(Name::new_static("value")),
+                None,
+            ));
+            if let Some(ctx) = context {
+                tcc = tcc.with_context(Some(ctx()));
+            }
+            tcc
+        });
+    }
+
+    /// Check the value-lookup argument of an enum call against the expected lookup type.
+    /// Accepts both expressions and already-flattened types; expressions are inferred into
+    /// `arg_errors` so errors do not trigger retry without a hint. Returns rewritten
+    /// args/keywords carrying the inferred type so metaclass checking avoids re-inference.
+    fn check_enum_value_lookup_arg<'b>(
+        &self,
+        cls: &ClassType,
+        constructor_kind: &ConstructorKind,
+        args: &[CallArg<'b>],
+        keywords: &[CallKeyword<'b>],
+        errors: &ErrorCollector,
+        arg_errors: &ErrorCollector,
+        context: Option<&dyn Fn() -> ErrorContext>,
+    ) -> Option<(Type, TextRange, bool)> {
+        let (value_arg, positional) = match (args, keywords) {
+            ([CallArg::Arg(arg)], []) => (arg, true),
+            ([], [kw]) if kw.arg.as_ref().is_some_and(|id| id.as_str() == "value") => {
+                (&kw.value, false)
+            }
+            _ => return None,
+        };
+        let want = self.enum_value_lookup_expected_arg_type(cls, constructor_kind, keywords)?;
+        let (got, range) = match value_arg {
+            TypeOrExpr::Expr(e) => (self.expr_infer(e, arg_errors), e.range()),
+            TypeOrExpr::Type(ty, r) => ((*ty).clone(), *r),
+        };
+        self.check_enum_lookup_arg_compatibility(&got, &want, range, errors, context);
+        Some((got, range, positional))
+    }
+
     /// If the metaclass defines a custom `__call__`, call it. If the `__call__` comes from `type`, ignore
     /// it because `type.__call__` behavior is baked into our constructor logic.
     fn call_metaclass(
@@ -1281,6 +1512,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         context: Option<&dyn Fn() -> ErrorContext>,
         hint: Option<HintRef>,
     ) -> Type {
+        let class_metadata = self.get_metadata_for_class(cls.class_object());
         // Classes that override `__new__` and also define `__init__` check/infer args
         // twice, potentially leading to exponential check costs: `O(2^depth)`.
         let checks_args_twice = |cls: &ClassType, preserve_self| {
@@ -1299,19 +1531,33 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             })
             .chain(keywords.iter().map(|k| &k.value))
             .any(is_deeply_nested)
-            && checks_args_twice(&cls, constructor_kind == ConstructorKind::TypeOfSelf);
+            && (checks_args_twice(&cls, constructor_kind == ConstructorKind::TypeOfSelf)
+                || class_metadata.is_enum());
+
+        let flatten_enum = class_metadata.is_enum()
+            && match (args, keywords) {
+                ([CallArg::Arg(TypeOrExpr::Expr(_))], []) => true,
+                ([], [kw]) => {
+                    kw.arg.as_ref().is_some_and(|id| id.as_str() == "value")
+                        && matches!(kw.value, TypeOrExpr::Expr(_))
+                }
+                _ => false,
+            };
 
+        let flatten = flatten_nested || flatten_enum;
         let call = CallWithTypes::new();
-        let flattened = flatten_nested.then(|| {
+        let flattened = flatten.then(|| {
             (
                 args.map(|a| match a {
-                    CallArg::Arg(x) | CallArg::Star(x, _) if is_deeply_nested(x) => {
+                    CallArg::Arg(x) | CallArg::Star(x, _)
+                        if (flatten_nested && is_deeply_nested(x)) || flatten_enum =>
+                    {
                         call.call_arg(a, self, arg_errors)
                     }
                     _ => a.clone(),
                 }),
                 keywords.map(|k| {
-                    if is_deeply_nested(&k.value) {
+                    if (flatten_nested && is_deeply_nested(&k.value)) || flatten_enum {
                         call.call_keyword(k, self, arg_errors)
                     } else {
                         k.clone()
@@ -1389,6 +1635,38 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         let return_errors = self.error_collector();
         // The solutions the constructor call settled on.
         let mut ctor_table = None;
+        // For enum value lookup, check the lookup argument and hand the inferred type
+        // (not the expression) to the metaclass `__call__` check.
+        let enum_lookup = if class_metadata.is_enum() {
+            self.check_enum_value_lookup_arg(
+                &cls,
+                &constructor_kind,
+                args,
+                keywords,
+                &errors,
+                &arg_errors,
+                context,
+            )
+        } else {
+            None
+        };
+        let enum_args;
+        let enum_keywords;
+        let (args, keywords) = match &enum_lookup {
+            Some((ty, range, true)) => {
+                enum_args = [CallArg::Arg(TypeOrExpr::Type(ty, *range))];
+                (&enum_args[..], keywords)
+            }
+            Some((ty, range, false)) => {
+                enum_keywords = [CallKeyword {
+                    range: keywords[0].range,
+                    arg: keywords[0].arg,
+                    value: TypeOrExpr::Type(ty, *range),
+                }];
+                (args, &enum_keywords[..])
+            }
+            None => (args, keywords),
+        };
         if let Some(ret) = self.call_metaclass(
             &cls,
             arguments_range,

pyrefly/lib/alt/class/enums.rs

@@ -286,7 +286,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
     /// Look up the `_value_` attribute for an enum type (not a specific member).
     /// Whether `_value_` should be read-write is unspecified, but we need to allow assigning
     /// it in `__init__` so we make it read-write.
-    fn enum_value_lookup_on_class(&self, class: &ClassType, enum_metadata: &EnumMetadata) -> Type {
+    pub(crate) fn enum_value_lookup_on_class(
+        &self,
+        class: &ClassType,
+        enum_metadata: &EnumMetadata,
+    ) -> Type {
         let mixed_in = self.mixed_in_enum_data_type(class.class_object());
         let has_new = self
             .get_class_fields(class.class_object())
@@ -337,9 +341,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
     /// If this enum mixes in a data type by inheriting from it, return the mixed-in type.
     /// Searches all bases, not just the first, to handle cases like
     /// `IntegerChoices(Choices, IntEnum)` where the data type comes from `IntEnum`.
-    /// A non-enum base is only treated as a data type if it has `__new__`
+    /// A non-enum base is treated as a data type if it is a dataclass or has `__new__`
     /// inherited from a class other than `object`, which is how Python
-    /// distinguishes data type mixins (`int`, `str`, `float`, and their
+    /// distinguishes data type mixins (`int`, `str`, `float`, dataclasses, and their
     /// subclasses like `MyStr(str)`) from regular method mixins.
     fn mixed_in_enum_data_type(&self, class: &Class) -> Option<Type> {
         let bases = self.get_base_types_for_class(class);
@@ -352,13 +356,18 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     return Some(ty);
                 }
             } else {
-                let is_data_type = self
-                    .get_class_member_with_defining_class(base.class_object(), &dunder::NEW)
-                    .is_some_and(|field| {
-                        !field
-                            .defining_class
-                            .has_toplevel_qname("builtins", "object")
-                    });
+                let is_dataclass = self
+                    .get_metadata_for_class(base.class_object())
+                    .dataclass_metadata()
+                    .is_some();
+                let is_data_type = is_dataclass
+                    || self
+                        .get_class_member_with_defining_class(base.class_object(), &dunder::NEW)
+                        .is_some_and(|field| {
+                            !field
+                                .defining_class
+                                .has_toplevel_qname("builtins", "object")
+                        });
                 if is_data_type {
                     return Some(self.heap.mk_class_type(base.clone()));
                 }