diff --git a/pyrefly/lib/alt/call.rs b/pyrefly/lib/alt/call.rs index 567a7be856..7273a2a993 100644 --- a/pyrefly/lib/alt/call.rs +++ b/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 { + // 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, ) -> 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, diff --git a/pyrefly/lib/alt/class/enums.rs b/pyrefly/lib/alt/class/enums.rs index 0aa7e98823..c19dd881a1 100644 --- a/pyrefly/lib/alt/class/enums.rs +++ b/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 { 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())); } diff --git a/pyrefly/lib/test/enums.rs b/pyrefly/lib/test/enums.rs index 49a23ca272..4efa8b542b 100644 --- a/pyrefly/lib/test/enums.rs +++ b/pyrefly/lib/test/enums.rs @@ -1394,6 +1394,255 @@ assert_type(x, Literal["q"]) "#, ); +testcase!( + test_enum_call_rejects_none_for_int_members, + r#" +from enum import Enum + +class Color(Enum): + RED = 1 + GREEN = 2 + +Color(1) +Color(value=2) +Color(Color.RED) +Color(None) # E: Argument `None` is not assignable to parameter `value` with type `Color | complex | float | int` +Color(value=None) # E: Argument `None` is not assignable to parameter `value` with type `Color | complex | float | int` + "#, +); + +testcase!( + test_enum_call_cross_numeric_equality, + r#" +from decimal import Decimal +from enum import Enum +from fractions import Fraction + +class Color(Enum): + RED = 1 + GREEN = 2 + +# Numeric values and equality domains succeed at runtime. +Color(1.0) +Color(True) +Color(Decimal(1)) +Color(Fraction(1)) + +class B(Enum): + YES = True + +B(1) +B(1.0) +B(1 + 0j) +B(None) # E: Argument `None` is not assignable to parameter `value` with type `B | bool | complex | float | int` + +class TupleEnum(Enum): + PAIR = (1, 2) + +TupleEnum((1.0, 2.0)) + +class ListEnum(Enum): + ITEMS = [1, 2] + +ListEnum([1.0, 2.0]) + "#, +); + +testcase!( + test_enum_call_ignores_multi_param_init, + r#" +from enum import Enum + +class Planet(Enum): + MERCURY = (3.303e+23, 2.4397e6) + EARTH = (5.976e+24, 6.37814e6) + + def __init__(self, mass: float, radius: float) -> None: + self.mass = mass + self.radius = radius + +# `__init__` runs at class definition, not at lookup. +Planet((5.976e+24, 6.37814e6)) +Planet(Planet.EARTH) +Planet(value=(3.303e+23, 2.4397e6)) + "#, +); + +testcase!( + test_enum_call_init_param_names_irrelevant, + r#" +from enum import Enum + +class Category(Enum): + GAMES = "GAMES", "0_0" + MUSIC = "MUSIC", "1_1" + + def __init__(self, label: str, id: str) -> None: + self._value_ = label + self.id = id + +Category(Category.GAMES) +Category(value=("GAMES", "0_0")) + "#, +); + +testcase!( + test_enum_call_new_changes_value_type, + r#" +from enum import Enum + +class Coordinate(Enum): + def __new__(cls, value: int, label: str): + obj = object.__new__(cls) + obj._value_ = label + return obj + + PX = (0, "P.X") + PY = (1, "P.Y") + +# `_value_` is the label string, not the `(int, str)` creation arguments. +Coordinate("P.X") +Coordinate(Coordinate.PX) + +class IntCoordinate(int, Enum): + def __new__(cls, value: int, label: str): + obj = int.__new__(cls, value) + obj._value_ = value + return obj + + PX = (0, "P.X") + +IntCoordinate(0) + +class Untyped(Enum): + def __new__(cls, raw: str): + obj = object.__new__(cls) + obj._value_ = len(raw) + return obj + + A = "a" + BB = "bb" + +Untyped(1) +Untyped("anything") + "#, +); + +testcase!( + test_enum_call_custom_metaclass_and_missing, + r#" +from enum import Enum, EnumMeta +from typing import Any, Self + +class LenientMeta(EnumMeta): + def __call__(cls, value: Any, *args: Any, **kwargs: Any) -> Any: + return super().__call__(value, *args, **kwargs) + +class WithMeta(Enum, metaclass=LenientMeta): + A = 1 + +WithMeta(None) +WithMeta("a") + +class WithMissing(Enum): + A = 1 + + @classmethod + def _missing_(cls, value: object) -> Self: + return cls.A + +WithMissing(None) +WithMissing("a") + "#, +); + +testcase!( + test_enum_call_deeply_nested, + r#" +from enum import Enum + +class Color(Enum): + RED = 1 + +Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(Color(1))))))))))))))))))))) + "#, +); + +testcase!( + test_enum_call_deeply_nested_error_linear_time, + r#" +from enum import Enum + +class E(Enum): + A = 1 + +def identity[T](x: T) -> T: + return x + +identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(identity(E(1 + "x")))))))))))))))))))))))))))))))))))))))) # E: `+` is not supported between `Literal[1]` and `Literal['x']` + "#, +); + +testcase!( + test_enum_call_extensible_base_self_type, + r#" +from enum import Enum +from typing import Self + +class Base(str, Enum): + @classmethod + def parse(cls, raw: object) -> Self: + return cls(raw) + +class Color(Base): + RED = "red" + + @classmethod + def _missing_(cls, value: object) -> "Color": + return cls.RED + +assert Color.parse(None) is Color.RED + "#, +); + +testcase!( + test_enum_call_dataclass_mixin, + r#" +from dataclasses import dataclass +from enum import Enum + +@dataclass +class Data: + n: int + +class E(Data, Enum): + A = (1,) + +assert E(Data(1)) is E.A + "#, +); + +testcase!( + test_enum_call_pre_inferred_arguments, + r#" +from enum import Enum + +class Color(Enum): + RED = 1 + + def __init__(self, value: int) -> None: + self._value_ = value + +def identity[T](x: T) -> T: + return x + +Color(None) # E: Argument `None` is not assignable to parameter `value` with type `Color | complex | float | int` +Color(identity(None)) # E: Argument `None` is not assignable to parameter `value` with type `Color | complex | float | int` +Color(identity(identity(identity(identity(None))))) # E: Argument `None` is not assignable to parameter `value` with type `Color | complex | float | int` +Color(value=identity(identity(identity(identity(None))))) # E: Argument `None` is not assignable to parameter `value` with type `Color | complex | float | int` + "#, +); + // https://github.com/facebook/pyrefly/issues/5279 testcase!( test_enum_member_keyword_value, @@ -1412,3 +1661,71 @@ class Format(Enum): return formatter(value) "#, ); + +fn env_enum_stub_ellipsis() -> TestEnv { + let mut env = TestEnv::new(); + env.add_with_path( + "res", + "res.pyi", + r#" +from enum import Enum + +class Resolution(Enum): + RESO_NS = ... + RESO_US = ... +"#, + ); + env +} + +testcase!( + test_enum_lookup_stub_ellipsis_values_unchecked, + env_enum_stub_ellipsis(), + r#" +import res +res.Resolution(7) +"#, +); + +testcase!( + test_enum_call_object_and_union_arguments, + r#" +from enum import Enum +from typing import Any + +class IntColor(Enum): + RED = 1 + GREEN = 2 + +class StrColor(Enum): + HOSTED = "hosted" + LOCAL = "local" + +def test_object_and_any( + obj: object, + any_val: Any, + opt_obj: object | None, + opt_str: str | None, + zulip_union: bool | int | str, + meson_union: bool | int | list[str] | str, +) -> None: + # `object` and `Any` are accepted for both numeric and string enums because + # lookup uses equality at runtime. + IntColor(obj) + IntColor(any_val) + StrColor(obj) + StrColor(any_val) + + # Nullable types must be rejected unless the enum accepts None or defines _missing_. + IntColor(opt_obj) # E: Argument `object | None` is not assignable to parameter `value` with type `IntColor | complex | float | int` + StrColor(opt_obj) # E: Argument `object | None` is not assignable to parameter `value` with type `StrColor | str` + StrColor(opt_str) # E: Argument `str | None` is not assignable to parameter `value` with type `StrColor | str` + + # Incompatible unions where some branches do not match the target must be rejected. + IntColor(zulip_union) # E: Argument `bool | int | str` is not assignable to parameter `value` with type `IntColor | complex | float | int` + StrColor(meson_union) # E: Argument `bool | int | list[str] | str` is not assignable to parameter `value` with type `StrColor | str` + + # Disjoint type (str passed to int enum) must be rejected. + IntColor("1") # E: Argument `Literal['1']` is not assignable to parameter `value` with type `IntColor | complex | float | int` +"#, +);