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

Code before:
            && let lhs_value_type = inference.expression_type(&*subscript.value)
            // Checking for `TypedDict`s up front isn't strictly necessary, since the iterator
            // below will yield nothing if it doesn't find any, but we want to do as little work as
            // possible in the common case.
            && is_typeddict_or_union_with_typeddicts(lhs_value_type, self.db)
            && let Some(subscript_place_expr) = place_expr(&subscript.value)
            && let Some(key_literal) = inference
                .expression_type(&*subscript.slice)
                .as_string_literal()
        {
            let field_name = Name::from(key_literal.value(self.db));
            let rhs_type = inference.expression_type(&comparators[0]);
            let mut intersection = IntersectionBuilder::new(self.db);
            // This iterator handles individual `TypedDict`, unions, intersections (the positive
            // members), and unions of intersections.
            for typed_dict_type in all_typeddicts_within_type_iter(lhs_value_type, self.db) {
                if let Some(field) = typed_dict_type.items(self.db).get(&field_name) {
                    let known_equality_result = match (field.declared_ty, rhs_type) {
                        (Type::StringLiteral(lhs), Type::StringLiteral(rhs)) => Some(lhs == rhs),
                        (Type::IntLiteral(lhs), Type::IntLiteral(rhs)) => Some(lhs == rhs),
                        // It's unlikely that the user will mix int and string literals in the same
                        // union, but go ahead and handle it.
                        (Type::StringLiteral(_), Type::IntLiteral(_))
                        | (Type::IntLiteral(_), Type::StringLiteral(_)) => Some(false),
                        _ => None,
                    };
                    if let Some(equality) = known_equality_result
                        && equality != is_positive
                    {
                        // We could synthesize a `TypedDict` with just this one specific field, but
                        // using the caller's own named types makes the resulting intersections
                        // easier to read, in the cases where they don't simplify out.
                        intersection = intersection.add_negative(Type::TypedDict(typed_dict_type));
                    }
                }
            }
            let place = self.expect_place(&subscript_place_expr);
            constraints.insert(place, intersection.build());
        }

        let mut last_rhs_ty: Option<Type> = None;

        for (op, (left, right)) in std::iter::zip(&**ops, comparator_tuples) {
            let lhs_ty = last_rhs_ty.unwrap_or_else(|| inference.expression_type(left));
            let rhs_ty = inference.expression_type(right);
            last_rhs_ty = Some(rhs_ty);

            match left {
                ast::Expr::Name(_)
                | ast::Expr::Attribute(_)
                | ast::Expr::Subscript(_)

Review comment (bug):
another important case to handle here is enum-literal types.

However, there's some additional complications with enum literals, so I think it might be good to postpone enums to a followup. Namely, an enum literal can't be used as a "tag" if the enum class has a custom `__eq__` method. That's true if the enum class inherits from `IntEnum` or `StrEnum` (we handle this incorrectly in [some other situations](https://github.com/astral-sh/ty/issues/1454) right now).

Code after:
        //         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
                // 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.
                let field_type = rhs_type.negate_if(self.db, constrain_with_equality);
                // Create the synthesized `TypedDict` with that (possibly negated) field. We don't
                // want to constrain the mutability or required-ness of the field, so the most
                // compatible form is not-required and read-only.
                let field = TypedDictFieldBuilder::new(field_type)
                    .required(false)
                    .read_only(true)
                    .build();