diff --git a/pyrefly/lib/solver/subset.rs b/pyrefly/lib/solver/subset.rs index 0a586bc87e..061efd1b09 100644 --- a/pyrefly/lib/solver/subset.rs +++ b/pyrefly/lib/solver/subset.rs @@ -2783,6 +2783,16 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> { Ok(()) } (Type::QuantifiedValue(_), Type::ClassType(cls)) if is_int_class_type(cls) => Ok(()), + (Type::ClassType(got_cls), Type::Literal(lit)) + if let Lit::Enum(enum_lit) = &lit.value + && got_cls.class_object() == enum_lit.class.class_object() + && self + .type_order + .get_enum_member_count(got_cls.class_object()) + == Some(1) => + { + Ok(()) + } (Type::Literal(l_lit), Type::Literal(u_lit)) => { ok_or(l_lit.value == u_lit.value, SubsetError::Other) } diff --git a/pyrefly/lib/test/enums.rs b/pyrefly/lib/test/enums.rs index 927737a400..44618ea355 100644 --- a/pyrefly/lib/test/enums.rs +++ b/pyrefly/lib/test/enums.rs @@ -438,6 +438,7 @@ def f(test: bool, e1: E1, e2: E2): x: Literal[E1.X, E1.Y] = e1 y: Literal[E1.X, E1.Y, 1] = e1 z: Literal[E2.X, E2.Y] = e2 # E: `E2` is not assignable to `Literal[E2.X, E2.Y]` + single: Literal[E2.X] = e2 # E: `E2` is not assignable to `Literal[E2.X]` "#, ); diff --git a/pyrefly/lib/test/type_alias.rs b/pyrefly/lib/test/type_alias.rs index bfd56e5b15..d8204596ae 100644 --- a/pyrefly/lib/test/type_alias.rs +++ b/pyrefly/lib/test/type_alias.rs @@ -160,6 +160,78 @@ def f(x: X[int]): "#, ); +testcase!( + test_generic_alias_union_with_enum_literal, + r#" +from enum import Enum +from typing import Literal, TypeAlias, TypeVar, assert_type + +T = TypeVar("T") + +class Unset(Enum): + _UNSET = "UNSET" + +UNSET: Literal[Unset._UNSET] = Unset._UNSET +UnsetType: TypeAlias = Literal[Unset._UNSET] +Missing: TypeAlias = UnsetType | T | None + +def takes(x: Missing[Literal["a", "b"]]) -> None: ... + +def or_unset[U](v: U | None) -> U | Literal[Unset._UNSET]: + return UNSET if v is None else v + +def f(s: Literal["a"] | None) -> None: + assert_type(or_unset(s), Literal["a"] | Literal[Unset._UNSET]) + assert_type(s, Literal["a"] | None) + takes(or_unset(s)) +"#, +); + +testcase!( + test_generic_alias_union_with_enum_literal_direct_argument, + r#" +from enum import Enum +from typing import Literal, TypeAlias, TypeVar + +T = TypeVar("T") + +class Unset(Enum): + _UNSET = "UNSET" + +UnsetType: TypeAlias = Literal[Unset._UNSET] +Missing: TypeAlias = UnsetType | T | None + +def takes(x: Missing[Literal["a", "b"]]) -> None: ... + +def f(x: Literal["a"] | Literal[Unset._UNSET] | None) -> None: + takes(x) +"#, +); + +testcase!( + test_generic_alias_union_with_inline_enum_literal, + r#" +from enum import Enum +from typing import Literal, TypeAlias, TypeVar + +T = TypeVar("T") + +class Unset(Enum): + _UNSET = "UNSET" + +UNSET: Literal[Unset._UNSET] = Unset._UNSET +Missing: TypeAlias = Literal[Unset._UNSET] | T | None + +def takes(x: Missing[Literal["a", "b"]]) -> None: ... + +def or_unset[U](v: U | None) -> U | Literal[Unset._UNSET]: + return UNSET if v is None else v + +def f(s: Literal["a"] | None) -> None: + takes(or_unset(s)) +"#, +); + testcase!( test_ordinary_alias_union_retains_precise_members, r#" @@ -1396,3 +1468,21 @@ X: TypeAlias = int | str X() # E: Expected a callable "#, ); + +testcase!( + test_enum_literal_union_member_compatibility, + r#" +from enum import Enum +from typing import Literal + +class Unset(Enum): + _UNSET = "UNSET" + +def takes_member(x: Literal["a", "b", Unset._UNSET] | None) -> None: ... +def takes_class(x: Literal["a", "b"] | Unset | None) -> None: ... + +def f(x: Literal["a"] | Literal[Unset._UNSET] | None) -> None: + takes_member(x) + takes_class(x) +"#, +);