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

Code before:
        //         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)
                            }
                        })
                })
                .copied()
                .collect();
            if filtered.len() < union.elements(self.db).len() {
                let place = self.expect_place(&subscript_place_expr);
                constraints.insert(
                    place,
                    NarrowingConstraint::regular(UnionType::from_elements(self.db, filtered)),
                );
            }
        }

        // Narrow tagged unions of `TypedDict`s with `Literal` keys, for example:
        //
        //     class Foo(TypedDict):
        //         tag: Literal["foo"]
        //     class Bar(TypedDict):
        //         tag: Literal["bar"]
        //     def _(union: Foo | Bar):
        //         if union["tag"] == "foo":
        //             reveal_type(union)  # Foo
        //
        // Importantly, `my_typeddict_union["tag"]` isn't the place we're going to constraint.

Review comment (suggestion):
Maybe we should rename this to not be specific to `TypeGuard`?

But not in this PR.

Code after:
        {
            let is_positive_check = is_positive == (ops[0] == ast::CmpOp::Is);
            let union_elements = union.elements(self.db);
            let filtered: Vec<_> = union_elements
                .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)
                            }
                        })
                })
                .copied()
                .collect();
            if filtered.len() < union_elements.len() {
                let place = self.expect_place(&subscript_place_expr);
                constraints.insert(
                    place,
                    NarrowingConstraint::typeguard(UnionType::from_elements(self.db, filtered)),
                );
            }
        }

        // Narrow tagged unions of `TypedDict`s with `Literal` keys, for example:
        //
        //     class Foo(TypedDict):
        //         tag: Literal["foo"]
        //     class Bar(TypedDict):
        //         tag: Literal["bar"]
        //     def _(union: Foo | Bar):
        //         if union["tag"] == "foo":
        //             reveal_type(union)  # Foo
        //
        // Importantly, `my_typeddict_union["tag"]` isn't the place we're going to constraint.
        // Instead, we're going to constrain `my_typeddict_union` itself.
        if matches!(&**ops, [ast::CmpOp::Eq | ast::CmpOp::NotEq])
            && let ast::Expr::Subscript(subscript) = &**left
        {
            // For `==`, we use equality semantics on the `if` branch (is_positive=true).
            // For `!=`, we use equality semantics on the `else` branch (is_positive=false).
            let constrain_with_equality = is_positive == (ops[0] == ast::CmpOp::Eq);
            if let Some((place, constraint)) = self.narrow_typeddict_subscript(
                inference.expression_type(&*subscript.value),
                &subscript.value,