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

Code before:
        //     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
            && let lhs_value_type = inference.expression_type(&*subscript.value)
            // Checking for `TypedDict`s up front isn't strictly necessary, since the intersection
            // we're going to build is compatible with non-`TypedDict` types, but we don't want to
            // do the work to build it and intersect it (or for that matter, let the user see it)
            // in the common case where there are no `TypedDict`s.
            && is_typeddict_or_union_with_typeddicts(self.db, lhs_value_type)
            && let Some(subscript_place_expr) = place_expr(&subscript.value)
            && let Type::StringLiteral(key_literal) = inference.expression_type(&*subscript.slice)
            && let rhs_type = inference.expression_type(&comparators[0])
            && matches!(rhs_type, Type::StringLiteral(_) | Type::IntLiteral(_))
        {
            let constrain_with_equality = match (ops[0], is_positive) {
                (ast::CmpOp::Eq, true) | (ast::CmpOp::NotEq, false) => true,
                (ast::CmpOp::Eq, false) | (ast::CmpOp::NotEq, true) => false,
                _ => unreachable!("we checked for Eq/NotEq above"),
            };
            // If we have an equality constraint (either `==` on the `if` side, or `!=` on the
            // `else` side), we have to be careful. If all the matching fields in all the
            // `TypedDict`s here have literal types, then yes, equality is as good as a type check.
            // However, if any of them are e.g. `int` or `str` or some random class, then we can't
            // narrow their type at all, because subclasses of those types can implement `__eq__`
            // in any perverse way they like. On the other hand, if this is an *inequality*
            // constraint, then we can go ahead and assert "you can't be this exact literal type"
            // without worrying about what types are present.
            if !constrain_with_equality
                || all_matching_typeddict_fields_have_literal_types(
                    self.db,
                    lhs_value_type,
                    key_literal.value(self.db),
                )
            {
                let field_name = Name::from(key_literal.value(self.db));
                let rhs_type = inference.expression_type(&comparators[0]);
                let field_type = match (ops[0], is_positive) {
                    // To avoid excluding non-`TypedDict` types, our constraints are always
                    // expressed as a negative intersection (i.e. "you're *not* this kind of
                    // `TypedDict`"). In these two cases, the whole constraint is going to be a
                    // double negative, i.e. "you're *not* a `TypedDict` *without* this literal
                    // field". As the first step of building that, we negate the right hand side.
                    (ast::CmpOp::Eq, true) | (ast::CmpOp::NotEq, false) => {
                        IntersectionBuilder::new(self.db)

Review comment (suggestion):
```suggestion
            let constrain_with_equality = is_positive == (ops[0] == ast::CmpOp);
```

Code after:
        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 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
            && let lhs_value_type = inference.expression_type(&*subscript.value)
            // Checking for `TypedDict`s up front isn't strictly necessary, since the intersection
            // we're going to build is compatible with non-`TypedDict` types, but we don't want to
            // do the work to build it and intersect it (or for that matter, let the user see it)
            // in the common case where there are no `TypedDict`s.
            && is_typeddict_or_union_with_typeddicts(self.db, lhs_value_type)
            && let Some(subscript_place_expr) = place_expr(&subscript.value)
            && let Type::StringLiteral(key_literal) = inference.expression_type(&*subscript.slice)
            && let rhs_type = inference.expression_type(&comparators[0])
            && is_supported_typeddict_tag_literal(rhs_type)
        {
            // If we have an equality constraint (either `==` on the `if` side, or `!=` on the
            // `else` side), we have to be careful. If all the matching fields in all the
            // `TypedDict`s here have literal types, then yes, equality is as good as a type check.
            // However, if any of them are e.g. `int` or `str` or some random class, then we can't
            // narrow their type at all, because subclasses of those types can implement `__eq__`
            // in any perverse way they like. On the other hand, if this is an *inequality*
            // constraint, then we can go ahead and assert "you can't be this exact literal type"
            // without worrying about what other types might be present.
            let constrain_with_equality = is_positive == (ops[0] == ast::CmpOp::Eq);
            if !constrain_with_equality
                || all_matching_typeddict_fields_have_literal_types(
                    self.db,
                    lhs_value_type,
                    key_literal.value(self.db),
                )
            {
                let field_name = Name::from(key_literal.value(self.db));
                let rhs_type = inference.expression_type(&comparators[0]);
                // To avoid excluding non-`TypedDict` types, our constraints are always expressed
                // as a negative intersection (i.e. "you're *not* this kind of `TypedDict`"). If
                // `constrain_with_equality` is true, the whole constraint is going to be a double