diff --git a/pyrefly/lib/alt/narrow.rs b/pyrefly/lib/alt/narrow.rs index 344319421a..2c5f958e33 100644 --- a/pyrefly/lib/alt/narrow.rs +++ b/pyrefly/lib/alt/narrow.rs @@ -1285,6 +1285,22 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { }) } + /// Removes `None` from every element of concrete tuples. A tuple with an element that is + /// always `None` cannot satisfy the condition, so it becomes `Never`. + fn narrow_tuple_elements_not_none(&self, ty: &Type) -> Type { + self.distribute_over_union(ty, |t| match t { + Type::Tuple(Tuple::Concrete(elts)) => { + let narrowed = elts.map(|elt| self.narrow_is_not(elt, &Type::None)); + if narrowed.iter().any(Type::is_never) { + self.heap.mk_never() + } else { + self.heap.mk_tuple(Tuple::Concrete(narrowed)) + } + } + _ => t.clone(), + }) + } + /// Keep the union members consistent with `key` being present or absent. fn narrow_key_membership(&self, ty: &Type, key: &Name, present: bool) -> Type { self.distribute_over_union(ty, |member| { @@ -1661,6 +1677,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { }; self.narrow_length_less_than(ty, len + 1) } + AtomicNarrowOp::AnyElementIsNone => ty.clone(), + AtomicNarrowOp::NoElementIsNone => self.narrow_tuple_elements_not_none(ty), AtomicNarrowOp::IsSequence => { self.is_type_for_pattern(ty, |t| self.is_sequence_for_pattern(t)) } diff --git a/pyrefly/lib/binding/bindings.rs b/pyrefly/lib/binding/bindings.rs index 49fc7a787f..92fa6bdd83 100644 --- a/pyrefly/lib/binding/bindings.rs +++ b/pyrefly/lib/binding/bindings.rs @@ -1863,20 +1863,18 @@ impl<'a> BindingsBuilder<'a> { return SpecialExport::new(name); } Binding::Phi(_, branches) => { - // Check all branches for a consistent special export (e.g. try/except - // importing Literal from typing vs typing_extensions). + // Every branch must resolve to the same special export (e.g. try/except + // importing Literal from typing vs typing_extensions). A branch that is + // not a special export means the name may be an ordinary object. let mut result = None; for branch in branches { let branch_result = self.special_export_from_binding_idx( branch.value_key, visited_names, visited_keys, - ); - match (&result, &branch_result) { - (None, _) => result = branch_result, - (_, None) => {} - (Some(a), Some(b)) if a == b => {} - _ => return None, + )?; + if *result.get_or_insert(branch_result) != branch_result { + return None; } } return result; diff --git a/pyrefly/lib/binding/narrow.rs b/pyrefly/lib/binding/narrow.rs index 108e670bf0..6119d6df16 100644 --- a/pyrefly/lib/binding/narrow.rs +++ b/pyrefly/lib/binding/narrow.rs @@ -101,6 +101,10 @@ pub enum AtomicNarrowOp { LenLte(Expr), /// Narrowing for sequence pattern matching - confirms the subject is a sequence type IsSequence, + /// `any(x is None for x in )`. Does not narrow when true. + AnyElementIsNone, + /// `not any(x is None for x in )`: no element of the tuple subject is `None`. + NoElementIsNone, /// Negation of IsSequence - confirms the subject is NOT a sequence type IsNotSequence, /// Used to narrow to mapping types (types that extend typing.Mapping) @@ -237,6 +241,8 @@ impl DisplayWith for AtomicNarrowOp { AtomicNarrowOp::LenLt(expr) => write!(f, "LenLt({})", expr.display_with(ctx)), AtomicNarrowOp::LenLte(expr) => write!(f, "LenLte({})", expr.display_with(ctx)), AtomicNarrowOp::IsSequence => write!(f, "IsSequence"), + AtomicNarrowOp::AnyElementIsNone => write!(f, "AnyElementIsNone"), + AtomicNarrowOp::NoElementIsNone => write!(f, "NoElementIsNone"), AtomicNarrowOp::IsNotSequence => write!(f, "IsNotSequence"), AtomicNarrowOp::IsMapping => write!(f, "IsMapping"), AtomicNarrowOp::IsNotMapping => write!(f, "IsNotMapping"), @@ -348,6 +354,8 @@ impl AtomicNarrowOp { Self::LenGte(expr) => Some(format!("len({subject}) >= {}", snippet(expr.range())?)), Self::LenLt(expr) => Some(format!("len({subject}) < {}", snippet(expr.range())?)), Self::LenLte(expr) => Some(format!("len({subject}) <= {}", snippet(expr.range())?)), + Self::AnyElementIsNone => Some(format!("any(x is None for x in {subject})")), + Self::NoElementIsNone => Some(format!("not any(x is None for x in {subject})")), Self::IsSequence => Some(format!("isinstance({subject}, Sequence)")), Self::IsNotSequence => Some(format!("not isinstance({subject}, Sequence)")), Self::IsMapping => Some(format!("isinstance({subject}, Mapping)")), @@ -410,6 +418,8 @@ impl AtomicNarrowOp { Self::LenLte(v) => Self::LenGt(v.clone()), Self::LenLt(v) => Self::LenGte(v.clone()), Self::LenNotEq(v) => Self::LenEq(v.clone()), + Self::AnyElementIsNone => Self::NoElementIsNone, + Self::NoElementIsNone => Self::AnyElementIsNone, Self::IsSequence => Self::IsNotSequence, Self::IsNotSequence => Self::IsSequence, Self::IsMapping => Self::IsNotMapping, @@ -1240,6 +1250,16 @@ impl NarrowOps { test.range(), ) } + Expr::Call( + test @ ExprCall { + node_index: _, + range_start: _, + func, + arguments, + }, + ) if let Some((iterable, op)) = Self::element_none_check(builder, func, arguments) => { + Self::from_single_narrow_op(iterable, op, test.range()) + } Expr::Call( test @ ExprCall { node_index: _, @@ -1359,6 +1379,39 @@ impl NarrowOps { } } + /// Matches `any(x is None for x in S)` and `all(x is not None for x in S)`, where the + /// callee is the builtin and the generator is single, unconditional and synchronous. + /// Returns `S` and the op describing the call being true. + fn element_none_check<'a>( + builder: &BindingsBuilder, + func: &Expr, + arguments: &'a Arguments, + ) -> Option<(&'a Expr, AtomicNarrowOp)> { + let [Expr::Generator(generator)] = &*arguments.args else { + return None; + }; + let [comprehension] = &*generator.generators else { + return None; + }; + let Expr::Compare(ExprCompare { ops, operands, .. }) = &*generator.elt else { + return None; + }; + let (Expr::Name(elt), [op], [_, Expr::NoneLiteral(_)], Expr::Name(target)) = + (&operands[0], &**ops, &**operands, &comprehension.target) + else { + return None; + }; + if elt.id != target.id || !comprehension.ifs.is_empty() || comprehension.is_async { + return None; + } + let op = match (builder.as_special_export(func)?, op) { + (SpecialExport::Any, CmpOp::Is) => AtomicNarrowOp::AnyElementIsNone, + (SpecialExport::All, CmpOp::IsNot) => AtomicNarrowOp::NoElementIsNone, + _ => return None, + }; + Some((&comprehension.iter, op)) + } + fn get_original_binding<'a>( builder: &'a BindingsBuilder, name: &ExprName, @@ -1401,6 +1454,9 @@ impl NarrowOps { | AtomicNarrowOp::LenGte(..) | AtomicNarrowOp::LenLt(..) | AtomicNarrowOp::LenLte(..) + // These ops are only applied to concrete tuples, which are immutable. + | AtomicNarrowOp::AnyElementIsNone + | AtomicNarrowOp::NoElementIsNone // This is technically unsafe, because it marks arbitrary TypeGuard/TypeIs results // as still valid, but we need to allow this for `isinstance` and friends to work. | AtomicNarrowOp::Call(..) diff --git a/pyrefly/lib/export/special.rs b/pyrefly/lib/export/special.rs index e9917dfa9a..fd96ee855c 100644 --- a/pyrefly/lib/export/special.rs +++ b/pyrefly/lib/export/special.rs @@ -44,6 +44,8 @@ pub enum SpecialExport { Quit, OsExit, Len, + Any, + All, Range, Bool, BuiltinsType, @@ -137,6 +139,8 @@ impl SpecialExport { "quit" => Some(Self::Quit), "_exit" => Some(Self::OsExit), "len" => Some(Self::Len), + "any" => Some(Self::Any), + "all" => Some(Self::All), "range" => Some(Self::Range), "bool" => Some(Self::Bool), "type" => Some(Self::BuiltinsType), @@ -237,6 +241,8 @@ impl SpecialExport { Self::Enum | Self::StrEnum | Self::IntEnum => matches!(m.as_str(), "enum"), Self::Super | Self::Len + | Self::Any + | Self::All | Self::Range | Self::Quit | Self::Bool diff --git a/pyrefly/lib/test/narrow.rs b/pyrefly/lib/test/narrow.rs index f4c94f02f0..7d0b036cd2 100644 --- a/pyrefly/lib/test/narrow.rs +++ b/pyrefly/lib/test/narrow.rs @@ -5137,3 +5137,92 @@ def f(a: Fraction | str, b: UserString | int) -> None: assert_type(b, UserString) "#, ); + +testcase!( + test_narrow_tuple_elements_any_is_none, + r#" +from typing import assert_type + +def f(a: int | None, b: str | None): + values = (a, b) + if any(v is None for v in values): + assert_type(values, tuple[int | None, str | None]) + raise ValueError() + assert_type(values, tuple[int, str]) + +def g(a: int | None, b: str | None): + values = (a, b) + if all(v is not None for v in values): + assert_type(values, tuple[int, str]) + else: + assert_type(values, tuple[int | None, str | None]) + +def h(a: int | None, b: str | None): + values = (a, b) + if any(v is None for v in values if v != 0): + raise ValueError() + assert_type(values, tuple[int | None, str | None]) + "#, +); + +testcase!( + test_narrow_tuple_elements_none_check_drops_impossible_members, + r#" +from typing import assert_type + +def f(values: tuple[None, str] | tuple[int, int] | tuple[int | None]): + if any(v is None for v in values): + raise ValueError() + assert_type(values, tuple[int, int] | tuple[int]) + +def g(values: tuple[None, str] | tuple[int, int] | tuple[int | None]): + if all(v is not None for v in values): + assert_type(values, tuple[int, int] | tuple[int]) + "#, +); + +testcase!( + test_narrow_tuple_elements_none_check_alias_and_shadowing, + r#" +from typing import Iterable, assert_type + +def mixed_any(values: tuple[int | None, str | None], use_builtin: bool): + if use_builtin: + check = any + else: + def check(it: Iterable[bool]) -> bool: + return False + if check(v is None for v in values): + raise ValueError() + assert_type(values, tuple[int | None, str | None]) + +def mixed_all(values: tuple[int | None, str | None], use_builtin: bool): + if use_builtin: + check = all + else: + def check(it: Iterable[bool]) -> bool: + return True + if check(v is not None for v in values): + assert_type(values, tuple[int | None, str | None]) + +def shadowed(values: tuple[int | None, str | None]): + def any(it: Iterable[bool]) -> bool: + return False + if any(v is None for v in values): + raise ValueError() + assert_type(values, tuple[int | None, str | None]) + "#, +); + +testcase!( + test_narrow_tuple_elements_none_check_through_variable, + r#" +from typing import assert_type + +def f(a: int | None, b: str | None): + values = (a, b) + ok = all(v is not None for v in values) + if ok: + assert_type(values, tuple[int, str]) + "#, +);