[red-knot] Extract a TypeInferenceContext

crates/red_knot_python_semantic/src/types.rs

@@ -1,4 +1,4 @@
-use infer::TypeInferenceBuilder;
+use infer::TypeInferenceContext;
 use ruff_db::files::File;
 use ruff_python_ast as ast;
 
@@ -482,22 +482,22 @@ impl<'db> Type<'db> {
     ///
     /// Returns `None` if `self` is not a callable type.
     #[must_use]
-    pub fn call(&self, db: &'db dyn Db) -> Option<Type<'db>> {
+    fn call(&self, db: &'db dyn Db, _context: &mut TypeInferenceContext<'db>) -> Option<Type<'db>> {
         match self {
             Type::Function(function_type) => Some(function_type.return_type(db)),
 
             // TODO annotated return type on `__new__` or metaclass `__call__`
             Type::Class(class) => Some(Type::Instance(*class)),
 
-            // TODO: handle classes which implement the Callable protocol
+            // TODO: handle classes which implement `__call__`
             Type::Instance(_instance_ty) => Some(Type::Unknown),
 
             // `Any` is callable, and its return type is also `Any`.
             Type::Any => Some(Type::Any),
 
             Type::Unknown => Some(Type::Unknown),
 
-            // TODO: union and intersection types, if they reduce to `Callable`
+            // TODO: union and intersection types
             Type::Union(_) => Some(Type::Unknown),
             Type::Intersection(_) => Some(Type::Unknown),
 
@@ -513,11 +513,14 @@ impl<'db> Type<'db> {
     /// for y in x:
     ///     pass
     /// ```
-    fn iterate(&self, db: &'db dyn Db) -> IterationOutcome<'db> {
+    /// Return None and emit a diagnostic if this type is not iterable.
+    fn iterate(
+        &self,
+        db: &'db dyn Db,
+        context: &mut TypeInferenceContext<'db>,
+    ) -> Option<Type<'db>> {
         if let Type::Tuple(tuple_type) = self {
-            return IterationOutcome::Iterable {
-                element_ty: UnionType::from_elements(db, &**tuple_type.elements(db)),
-            };
+            return Some(UnionType::from_elements(db, &**tuple_type.elements(db)));
         }
 
         // `self` represents the type of the iterable;
@@ -526,19 +529,16 @@ impl<'db> Type<'db> {
 
         let dunder_iter_method = iterable_meta_type.member(db, "__iter__");
         if !dunder_iter_method.is_unbound() {
-            let Some(iterator_ty) = dunder_iter_method.call(db) else {
-                return IterationOutcome::NotIterable {
-                    not_iterable_ty: *self,
-                };
+            let Some(iterator_ty) = dunder_iter_method.call(db, context) else {
+                context.not_iterable_diagnostic(*self);
+                return None;
             };
 
             let dunder_next_method = iterator_ty.to_meta_type(db).member(db, "__next__");
-            return dunder_next_method
-                .call(db)
-                .map(|element_ty| IterationOutcome::Iterable { element_ty })
-                .unwrap_or(IterationOutcome::NotIterable {
-                    not_iterable_ty: *self,
-                });
+            return dunder_next_method.call(db, context).or_else(|| {
+                context.not_iterable_diagnostic(*self);
+                None
+            });
         }
 
         // Although it's not considered great practice,
@@ -549,12 +549,10 @@ impl<'db> Type<'db> {
         // accepting `int` or `SupportsIndex`
         let dunder_get_item_method = iterable_meta_type.member(db, "__getitem__");
 
-        dunder_get_item_method
-            .call(db)
-            .map(|element_ty| IterationOutcome::Iterable { element_ty })
-            .unwrap_or(IterationOutcome::NotIterable {
-                not_iterable_ty: *self,
-            })
+        dunder_get_item_method.call(db, context).or_else(|| {
+            context.not_iterable_diagnostic(*self);
+            None
+        })
     }
 
     #[must_use]
@@ -619,28 +617,6 @@ impl<'db> From<&Type<'db>> for Type<'db> {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum IterationOutcome<'db> {
-    Iterable { element_ty: Type<'db> },
-    NotIterable { not_iterable_ty: Type<'db> },
-}
-
-impl<'db> IterationOutcome<'db> {
-    fn unwrap_with_diagnostic(
-        self,
-        iterable_node: ast::AnyNodeRef,
-        inference_builder: &mut TypeInferenceBuilder<'db>,
-    ) -> Type<'db> {
-        match self {
-            Self::Iterable { element_ty } => element_ty,
-            Self::NotIterable { not_iterable_ty } => {
-                inference_builder.not_iterable_diagnostic(iterable_node, not_iterable_ty);
-                Type::Unknown
-            }
-        }
-    }
-}
-
 #[salsa::interned]
 pub struct FunctionType<'db> {
     /// name of the function at definition