Reduce remove calls

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -2124,20 +2124,26 @@ shows up in a subset of the union members) is present, but that isn't generally
 field, it could be *assigned to* with another `TypedDict` that does:
 
 ```py
+from typing_extensions import Literal
+
 class Foo(TypedDict):
     foo: int
 
 class Bar(TypedDict):
     bar: int
 
-def disappointment(u: Foo | Bar):
+def disappointment(u: Foo | Bar, v: Literal["foo"]):
     if "foo" in u:
         # We can't narrow the union here...
         reveal_type(u)  # revealed: Foo | Bar
     else:
         # ...(even though we *can* narrow it here)...
-        # TODO: This should narrow to `Bar`, because "foo" is required in `Foo`.
+        reveal_type(u)  # revealed: Bar
+
+    if v in u:
         reveal_type(u)  # revealed: Foo | Bar
+    else:
+        reveal_type(u)  # revealed: Bar
 
 # ...because `u` could turn out to be one of these.
 class FooBar(TypedDict):
@@ -2148,6 +2154,39 @@ static_assert(is_assignable_to(FooBar, Foo))
 static_assert(is_assignable_to(FooBar, Bar))
 ```
 
+`not in` works in the opposite way to `in`: we can narrow in the positive case, but we cannot narrow
+in the negative case. The following snippet also tests our narrowing behaviour for intersections
+that contain `TypedDict`s, and unions that contain intersections that contain `TypedDict`s:
+
+```py
+from typing_extensions import Literal, Any
+from ty_extensions import Intersection, is_assignable_to, static_assert
+
+def _(t: Bar, u: Foo | Intersection[Bar, Any], v: Intersection[Bar, Any], w: Literal["bar"]):
+    reveal_type(u)  # revealed: Foo | (Bar & Any)
+    reveal_type(v)  # revealed: Bar & Any
+
+    if "bar" not in t:
+        reveal_type(t)  # revealed: Never
+    else:
+        reveal_type(t)  # revealed: Bar
+
+    if "bar" not in u:
+        reveal_type(u)  # revealed: Foo
+    else:
+        reveal_type(u)  # revealed: Foo | (Bar & Any)
+
+    if "bar" not in v:
+        reveal_type(v)  # revealed: Never
+    else:
+        reveal_type(v)  # revealed: Bar & Any
+
+    if w not in u:
+        reveal_type(u)  # revealed: Foo
+    else:
+        reveal_type(u)  # revealed: Foo | (Bar & Any)
+```
+
 TODO: The narrowing that we didn't do above will become possible when we add support for
 `closed=True`. This is [one of the main use cases][closed] that motivated the `closed` feature.
 

crates/ty_python_semantic/src/types.rs

@@ -7267,10 +7267,7 @@ impl<'db> Type<'db> {
             }
 
             (Some(Place::Defined(new_method, ..)), Place::Defined(init_method, ..)) => {
-                let callable = UnionBuilder::new(db)
-                    .add(*new_method)
-                    .add(*init_method)
-                    .build();
+                let callable = UnionType::from_elements(db, [new_method, init_method]);
 
                 let new_method_bindings = new_method
                     .bindings(db)
@@ -10758,11 +10755,7 @@ fn walk_type_var_constraints<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
 
 impl<'db> TypeVarConstraints<'db> {
     fn as_type(self, db: &'db dyn Db) -> Type<'db> {
-        let mut builder = UnionBuilder::new(db);
-        for ty in self.elements(db) {
-            builder = builder.add(*ty);
-        }
-        builder.build()
+        UnionType::from_elements(db, self.elements(db))
     }
 
     fn to_instance(self, db: &'db dyn Db) -> Option<TypeVarConstraints<'db>> {

crates/ty_python_semantic/src/types/cyclic.rs

@@ -24,9 +24,8 @@ use std::cmp::Eq;
 use std::hash::Hash;
 use std::marker::PhantomData;
 
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
-use crate::FxIndexSet;
 use crate::types::Type;
 
 /// Maximum recursion depth for cycle detection.
@@ -64,7 +63,7 @@ pub struct CycleDetector<Tag, T, R> {
     /// If the type we're visiting is present in `seen`, it indicates that we've hit a cycle (due
     /// to a recursive type); we need to immediately short circuit the whole operation and return
     /// the fallback value. That's why we pop items off the end of `seen` after we've visited them.
-    seen: RefCell<FxIndexSet<T>>,
+    seen: RefCell<FxHashSet<T>>,
 
     /// Unlike `seen`, this field is a pure performance optimisation (and an essential one). If the
     /// type we're trying to normalize is present in `cache`, it doesn't necessarily mean we've hit
@@ -86,7 +85,7 @@ pub struct CycleDetector<Tag, T, R> {
 impl<Tag, T: Hash + Eq + Clone, R: Clone> CycleDetector<Tag, T, R> {
     pub fn new(fallback: R) -> Self {
         CycleDetector {
-            seen: RefCell::new(FxIndexSet::default()),
+            seen: RefCell::new(FxHashSet::default()),
             cache: RefCell::new(FxHashMap::default()),
             depth: Cell::new(0),
             fallback,
@@ -99,24 +98,23 @@ impl<Tag, T: Hash + Eq + Clone, R: Clone> CycleDetector<Tag, T, R> {
             return val.clone();
         }
 
+        // Check depth limit to prevent stack overflow from recursive generic types
+        // with growing specializations (e.g., C[set[T]] -> C[set[set[T]]] -> ...)
+        let current_depth = self.depth.get();
+        if current_depth >= MAX_RECURSION_DEPTH {
+            return self.fallback.clone();
+        }
+
         // We hit a cycle
         if !self.seen.borrow_mut().insert(item.clone()) {
             return self.fallback.clone();
         }
 
-        // Check depth limit to prevent stack overflow from recursive generic types
-        // with growing specializations (e.g., C[set[T]] -> C[set[set[T]]] -> ...)
-        let current_depth = self.depth.get();
-        if current_depth >= MAX_RECURSION_DEPTH {
-            self.seen.borrow_mut().pop();
-            return self.fallback.clone();
-        }
         self.depth.set(current_depth + 1);
-
         let ret = func();
 
         self.depth.set(current_depth);
-        self.seen.borrow_mut().pop();
+        self.seen.borrow_mut().remove(&item);
         self.cache.borrow_mut().insert(item, ret.clone());
 
         ret
@@ -127,24 +125,24 @@ impl<Tag, T: Hash + Eq + Clone, R: Clone> CycleDetector<Tag, T, R> {
             return Some(val.clone());
         }
 
+        // Check depth limit to prevent stack overflow from recursive generic protocols
+        // with growing specializations (e.g., C[set[T]] -> C[set[set[T]]] -> ...)
+        let current_depth = self.depth.get();
+        if current_depth >= MAX_RECURSION_DEPTH {
+            return Some(self.fallback.clone());
+        }
+
         // We hit a cycle
         if !self.seen.borrow_mut().insert(item.clone()) {
             return Some(self.fallback.clone());
         }
 
-        // Check depth limit to prevent stack overflow from recursive generic protocols
-        // with growing specializations (e.g., C[set[T]] -> C[set[set[T]]] -> ...)
-        let current_depth = self.depth.get();
-        if current_depth >= MAX_RECURSION_DEPTH {
-            self.seen.borrow_mut().pop();
-            return Some(self.fallback.clone());
-        }
         self.depth.set(current_depth + 1);
 
         let ret = func()?;
 
         self.depth.set(current_depth);
-        self.seen.borrow_mut().pop();
+        self.seen.borrow_mut().remove(&item);
         self.cache.borrow_mut().insert(item, ret.clone());
 
         Some(ret)

crates/ty_python_semantic/src/types/infer/builder/type_expression.rs

@@ -1083,13 +1083,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                     &mut self.inner_expression_inference_state,
                     InnerExpressionInferenceState::Get,
                 );
-                let union = union
-                    .elements(self.db())
-                    .iter()
-                    .fold(UnionBuilder::new(self.db()), |builder, elem| {
-                        builder.add(self.infer_subscript_type_expression(subscript, *elem))
-                    })
-                    .build();
+                let union = union.map(self.db(), |element| {
+                    self.infer_subscript_type_expression(subscript, *element)
+                });
                 self.inner_expression_inference_state = previous_slice_inference_state;
                 union
             }

crates/ty_python_semantic/src/types/narrow.rs

@@ -12,7 +12,7 @@ use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::function::KnownFunction;
 use crate::types::infer::{ExpressionInference, infer_same_file_expression_type};
 use crate::types::typed_dict::{
-    SynthesizedTypedDictType, TypedDictFieldBuilder, TypedDictSchema, TypedDictType,
+    SynthesizedTypedDictType, TypedDictField, TypedDictFieldBuilder, TypedDictSchema, TypedDictType,
 };
 use crate::types::{
     CallableType, ClassLiteral, ClassType, IntersectionBuilder, IntersectionType, KnownClass,
@@ -926,10 +926,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
                 .build();
 
             // Keep order: first literal complement, then broader arms.
-            let result = UnionBuilder::new(self.db)
-                .add(narrowed_single)
-                .add(rest_union)
-                .build();
+            let result = UnionType::from_elements(self.db, [narrowed_single, rest_union]);
             Some(result)
         } else {
             None
@@ -1099,6 +1096,75 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
             }
         }
 
+        // Narrow unions and intersections of `TypedDict` in cases where required keys are
+        // excluded:
+        //
+        // class Foo(TypedDict):
+        //     foo: int
+        // class Bar(TypedDict):
+        //     bar: int
+        //
+        // def _(u: Foo | Bar):
+        //     if "foo" not in u:
+        //         reveal_type(u)  # revealed: Bar
+        if matches!(&**ops, [ast::CmpOp::In | ast::CmpOp::NotIn])
+            && let Type::StringLiteral(key) = inference.expression_type(&**left)
+            && let Some(rhs_place_expr) = place_expr(&comparators[0])
+            && let rhs_type = inference.expression_type(&comparators[0])
+            && is_typeddict_or_union_with_typeddicts(self.db, rhs_type)
+        {
+            let is_negative_check = is_positive == (ops[0] == ast::CmpOp::NotIn);
+            if is_negative_check {
+                let requires_key = |td: TypedDictType<'db>| -> bool {
+                    td.items(self.db)
+                        .get(key.value(self.db))
+                        .is_some_and(TypedDictField::is_required)
+                };
+
+                let narrowed = match rhs_type {
+                    Type::TypedDict(td) => {
+                        if requires_key(td) {
+                            Type::Never
+                        } else {
+                            rhs_type
+                        }
+                    }
+                    Type::Intersection(intersection) => {
+                        if intersection
+                            .positive(self.db)
+                            .iter()
+                            .copied()
+                            .filter_map(Type::as_typed_dict)
+                            .any(requires_key)
+                        {
+                            Type::Never
+                        } else {
+                            rhs_type
+                        }
+                    }
+                    Type::Union(union) => {
+                        // remove all members of the union that would require the key
+                        union.filter(self.db, |ty| match ty {
+                            Type::TypedDict(td) => !requires_key(*td),
+                            Type::Intersection(intersection) => !intersection
+                                .positive(self.db)
+                                .iter()
+                                .copied()
+                                .filter_map(Type::as_typed_dict)
+                                .any(requires_key),
+                            _ => true,
+                        })
+                    }
+                    _ => rhs_type,
+                };
+
+                if narrowed != rhs_type {
+                    let place = self.expect_place(&rhs_place_expr);
+                    constraints.insert(place, NarrowingConstraint::typeguard(narrowed));
+                }
+            }
+        }
+
         let mut last_rhs_ty: Option<Type> = None;
 
         for (op, (left, right)) in std::iter::zip(&**ops, comparator_tuples) {
@@ -1677,18 +1743,13 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
 fn is_typeddict_or_union_with_typeddicts<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool {
     match ty {
         Type::TypedDict(_) => true,
-        Type::Union(union) => {
-            union
-                .elements(db)
-                .iter()
-                .any(|union_member_ty| match union_member_ty {
-                    Type::TypedDict(_) => true,
-                    Type::Intersection(intersection) => {
-                        intersection.positive(db).iter().any(Type::is_typed_dict)
-                    }
-                    _ => false,
-                })
+        Type::Intersection(intersection) => {
+            intersection.positive(db).iter().any(Type::is_typed_dict)
         }
+        Type::Union(union) => union
+            .elements(db)
+            .iter()
+            .any(|union_member_ty| is_typeddict_or_union_with_typeddicts(db, *union_member_ty)),
         _ => false,
     }
 }