Narrow tuple elements after any()/all() None checks

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))
             }

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;

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 <subject>)`. Does not narrow when true.
+    AnyElementIsNone,
+    /// `not any(x is None for x in <subject>)`: 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<ModuleInfo> 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(..)

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

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])
+    "#,
+);