Repository: astral-sh/ruff
File: crates/ty_python_semantic/src/types/narrow.rs (Rust)

Code before:
                .all(|c| !is_narrowing_target_candidate(c))
        {
            return None;
        }

        if !is_positive && comparators.len() > 1 {
            // We can't negate a constraint made by a multi-comparator expression, since we can't
            // know which comparison part is the one being negated.
            // For example, the negation of  `x is 1 is y is 2`, would be `(x is not 1) or (y is not 1) or (y is not 2)`
            // and that requires cross-symbol constraints, which we don't support yet.
            return None;
        }

        let inference = infer_expression_types(self.db, expression, TypeContext::default());

        let comparator_tuples = std::iter::once(&**left)
            .chain(comparators)
            .tuple_windows::<(&ruff_python_ast::Expr, &ruff_python_ast::Expr)>();
        let mut constraints = NarrowingConstraints::default();

        // Narrow unions of tuples based on element checks. For example:
        //
        //     def _(t: tuple[int, int] | tuple[None, None]):
        //         if t[0] is not None:
        //             reveal_type(t)  # tuple[int, int]
        if matches!(&**ops, [ast::CmpOp::Is | ast::CmpOp::IsNot])
            && let ast::Expr::Subscript(subscript) = &**left
            && let Type::Union(union) = inference.expression_type(&*subscript.value)
            && let Some(subscript_place_expr) = place_expr(&subscript.value)
            && let Type::IntLiteral(index) = inference.expression_type(&*subscript.slice)
            && let Ok(index) = i32::try_from(index)
            && let rhs_ty = inference.expression_type(&comparators[0])
            && rhs_ty.is_singleton(self.db)
        {
            let dominated = is_positive == (ops[0] == ast::CmpOp::Is);
            let filtered: Vec<_> = union
                .elements(self.db)
                .iter()
                .filter(|elem| {
                    elem.as_nominal_instance()
                        .and_then(|inst| inst.tuple_spec(self.db))
                        .and_then(|spec| spec.py_index(self.db, index).ok())
                        .filter(|el_ty| !el_ty.is_union())
                        .is_none_or(|el_ty| el_ty.is_disjoint_from(self.db, rhs_ty) != dominated)
                })
                .copied()
                .collect();
            if filtered.len() < union.elements(self.db).len() {
                let place = self.expect_place(&subscript_place_expr);
                constraints.insert(place, UnionType::from_elements(self.db, filtered));
            }

Review comment (suggestion):
At some point we may want to do this for `Eq` and `NotEq` too, but that is a bit more complex to get right, can be left to a separate PR.

Code after:
                .all(|c| !is_narrowing_target_candidate(c))
        {
            return None;
        }

        if !is_positive && comparators.len() > 1 {
            // We can't negate a constraint made by a multi-comparator expression, since we can't
            // know which comparison part is the one being negated.
            // For example, the negation of  `x is 1 is y is 2`, would be `(x is not 1) or (y is not 1) or (y is not 2)`
            // and that requires cross-symbol constraints, which we don't support yet.
            return None;
        }

        let inference = infer_expression_types(self.db, expression, TypeContext::default());

        let comparator_tuples = std::iter::once(&**left)
            .chain(comparators)
            .tuple_windows::<(&ruff_python_ast::Expr, &ruff_python_ast::Expr)>();
        let mut constraints = NarrowingConstraints::default();

        // Narrow unions of tuples based on element checks. For example:
        //
        //     def _(t: tuple[int, int] | tuple[None, None]):
        //         if t[0] is not None:
        //             reveal_type(t)  # tuple[int, int]
        if matches!(&**ops, [ast::CmpOp::Is | ast::CmpOp::IsNot])
            && let ast::Expr::Subscript(subscript) = &**left
            && let Type::Union(union) = inference.expression_type(&*subscript.value)
            && let Some(subscript_place_expr) = place_expr(&subscript.value)
            && let Type::IntLiteral(index) = inference.expression_type(&*subscript.slice)
            && let Ok(index) = i32::try_from(index)
            && let rhs_ty = inference.expression_type(&comparators[0])
            && rhs_ty.is_singleton(self.db)
        {
            let is_positive_check = is_positive == (ops[0] == ast::CmpOp::Is);
            let filtered: Vec<_> = union
                .elements(self.db)
                .iter()
                .filter(|elem| {
                    elem.as_nominal_instance()
                        .and_then(|inst| inst.tuple_spec(self.db))
                        .and_then(|spec| spec.py_index(self.db, index).ok())
                        .is_none_or(|el_ty| {
                            if is_positive_check {
                                // `is X` context: keep tuples where element could be X
                                !el_ty.is_disjoint_from(self.db, rhs_ty)
                            } else {
                                // `is not X` context: keep tuples where element is not always X
                                !el_ty.is_subtype_of(self.db, rhs_ty)
                            }
                        })