Add API to emit type-checking diagnostics (#12988)

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

crates/red_knot_python_semantic/src/db.rs

@@ -1,15 +1,18 @@
+use ruff_db::files::File;
 use ruff_db::{Db as SourceDb, Upcast};
 
 /// Database giving access to semantic information about a Python program.
 #[salsa::db]
-pub trait Db: SourceDb + Upcast<dyn SourceDb> {}
+pub trait Db: SourceDb + Upcast<dyn SourceDb> {
+    fn is_file_open(&self, file: File) -> bool;
+}
 
 #[cfg(test)]
 pub(crate) mod tests {
     use std::sync::Arc;
 
     use crate::module_resolver::vendored_typeshed_stubs;
-    use ruff_db::files::Files;
+    use ruff_db::files::{File, Files};
     use ruff_db::system::{DbWithTestSystem, System, TestSystem};
     use ruff_db::vendored::VendoredFileSystem;
     use ruff_db::{Db as SourceDb, Upcast};
@@ -91,7 +94,11 @@ pub(crate) mod tests {
     }
 
     #[salsa::db]
-    impl Db for TestDb {}
+    impl Db for TestDb {
+        fn is_file_open(&self, file: File) -> bool {
+            !file.path(self).is_vendored_path()
+        }
+    }
 
     #[salsa::db]
     impl salsa::Database for TestDb {

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -154,6 +154,10 @@ impl<'db> SemanticIndex<'db> {
         &self.scopes[id]
     }
 
+    pub(crate) fn scope_ids(&self) -> impl Iterator<Item = ScopeId> {
+        self.scope_ids_by_scope.iter().copied()
+    }
+
     /// Returns the id of the parent scope.
     pub(crate) fn parent_scope_id(&self, scope_id: FileScopeId) -> Option<FileScopeId> {
         let scope = self.scope(scope_id);

crates/red_knot_python_semantic/src/types.rs

@@ -5,21 +5,37 @@ use crate::builtins::builtins_scope;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::symbol::{ScopeId, ScopedSymbolId};
 use crate::semantic_index::{
-    global_scope, symbol_table, use_def_map, DefinitionWithConstraints,
+    global_scope, semantic_index, symbol_table, use_def_map, DefinitionWithConstraints,
     DefinitionWithConstraintsIterator,
 };
 use crate::types::narrow::narrowing_constraint;
 use crate::{Db, FxOrderSet};
 
+pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
+pub(crate) use self::diagnostic::TypeCheckDiagnostics;
+pub(crate) use self::infer::{
+    infer_definition_types, infer_expression_types, infer_scope_types, TypeInference,
+};
+
 mod builder;
+mod diagnostic;
 mod display;
 mod infer;
 mod narrow;
 
-pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
-pub(crate) use self::infer::{
-    infer_definition_types, infer_expression_types, infer_scope_types, TypeInference,
-};
+pub fn check_types(db: &dyn Db, file: File) -> TypeCheckDiagnostics {
+    let _span = tracing::trace_span!("check_types", file=?file.path(db)).entered();
+
+    let index = semantic_index(db, file);
+    let mut diagnostics = TypeCheckDiagnostics::new();
+
+    for scope_id in index.scope_ids() {
+        let result = infer_scope_types(db, scope_id);
+        diagnostics.extend(result.diagnostics());
+    }
+
+    diagnostics
+}
 
 /// Infer the public type of a symbol (its type as seen from outside its scope).
 pub(crate) fn symbol_ty<'db>(
@@ -333,3 +349,48 @@ pub struct IntersectionType<'db> {
     /// directly in intersections rather than as a separate type.
     negative: FxOrderSet<Type<'db>>,
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Context;
+
+    use ruff_db::files::system_path_to_file;
+    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
+
+    use crate::db::tests::TestDb;
+    use crate::{Program, ProgramSettings, PythonVersion, SearchPathSettings};
+
+    #[test]
+    fn check_types() -> anyhow::Result<()> {
+        let mut db = TestDb::new();
+
+        db.write_file("src/foo.py", "import bar\n")
+            .context("Failed to write foo.py")?;
+
+        Program::from_settings(
+            &db,
+            ProgramSettings {
+                target_version: PythonVersion::default(),
+                search_paths: SearchPathSettings {
+                    extra_paths: Vec::new(),
+                    src_root: SystemPathBuf::from("/src"),
+                    site_packages: vec![],
+                    custom_typeshed: None,
+                },
+            },
+        )
+        .expect("Valid search path settings");
+
+        let foo = system_path_to_file(&db, "src/foo.py").context("Failed to resolve foo.py")?;
+
+        let diagnostics = super::check_types(&db, foo);
+
+        assert_eq!(diagnostics.len(), 1);
+        assert_eq!(
+            diagnostics[0].message(),
+            "Import 'bar' could not be resolved."
+        );
+
+        Ok(())
+    }
+}

crates/red_knot_python_semantic/src/types/diagnostic.rs

@@ -0,0 +1,111 @@
+use ruff_db::files::File;
+use ruff_text_size::{Ranged, TextRange};
+use std::fmt::Formatter;
+use std::ops::Deref;
+use std::sync::Arc;
+
+#[derive(Debug, Eq, PartialEq)]
+pub struct TypeCheckDiagnostic {
+    // TODO: Don't use string keys for rules
+    pub(super) rule: String,
+    pub(super) message: String,
+    pub(super) range: TextRange,
+    pub(super) file: File,
+}
+
+impl TypeCheckDiagnostic {
+    pub fn rule(&self) -> &str {
+        &self.rule
+    }
+
+    pub fn message(&self) -> &str {
+        &self.message
+    }
+
+    pub fn file(&self) -> File {
+        self.file
+    }
+}
+
+impl Ranged for TypeCheckDiagnostic {
+    fn range(&self) -> TextRange {
+        self.range
+    }
+}
+
+/// A collection of type check diagnostics.
+///
+/// The diagnostics are wrapped in an `Arc` because they need to be cloned multiple times
+/// when going from `infer_expression` to `check_file`. We could consider
+/// making [`TypeCheckDiagnostic`] a Salsa struct to have them Arena-allocated (once the Tables refactor is done).
+/// Using Salsa struct does have the downside that it leaks the Salsa dependency into diagnostics and
+/// each Salsa-struct comes with an overhead.
+#[derive(Default, Eq, PartialEq)]
+pub struct TypeCheckDiagnostics {
+    inner: Vec<std::sync::Arc<TypeCheckDiagnostic>>,
+}
+
+impl TypeCheckDiagnostics {
+    pub fn new() -> Self {
+        Self { inner: Vec::new() }
+    }
+
+    pub(super) fn push(&mut self, diagnostic: TypeCheckDiagnostic) {
+        self.inner.push(Arc::new(diagnostic));
+    }
+
+    pub(crate) fn shrink_to_fit(&mut self) {
+        self.inner.shrink_to_fit();
+    }
+}
+
+impl Extend<TypeCheckDiagnostic> for TypeCheckDiagnostics {
+    fn extend<T: IntoIterator<Item = TypeCheckDiagnostic>>(&mut self, iter: T) {
+        self.inner.extend(iter.into_iter().map(std::sync::Arc::new));
+    }
+}
+
+impl Extend<std::sync::Arc<TypeCheckDiagnostic>> for TypeCheckDiagnostics {
+    fn extend<T: IntoIterator<Item = Arc<TypeCheckDiagnostic>>>(&mut self, iter: T) {
+        self.inner.extend(iter);
+    }
+}
+
+impl<'a> Extend<&'a std::sync::Arc<TypeCheckDiagnostic>> for TypeCheckDiagnostics {
+    fn extend<T: IntoIterator<Item = &'a Arc<TypeCheckDiagnostic>>>(&mut self, iter: T) {
+        self.inner
+            .extend(iter.into_iter().map(std::sync::Arc::clone));
+    }
+}
+
+impl std::fmt::Debug for TypeCheckDiagnostics {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.inner.fmt(f)
+    }
+}
+
+impl Deref for TypeCheckDiagnostics {
+    type Target = [std::sync::Arc<TypeCheckDiagnostic>];
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl IntoIterator for TypeCheckDiagnostics {
+    type Item = Arc<TypeCheckDiagnostic>;
+    type IntoIter = std::vec::IntoIter<std::sync::Arc<TypeCheckDiagnostic>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.inner.into_iter()
+    }
+}
+
+impl<'a> IntoIterator for &'a TypeCheckDiagnostics {
+    type Item = &'a Arc<TypeCheckDiagnostic>;
+    type IntoIter = std::slice::Iter<'a, std::sync::Arc<TypeCheckDiagnostic>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.inner.iter()
+    }
+}

crates/red_knot_python_semantic/src/types/infer.rs

@@ -29,7 +29,8 @@ use salsa::plumbing::AsId;
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
 use ruff_python_ast as ast;
-use ruff_python_ast::{Expr, ExprContext};
+use ruff_python_ast::{AnyNodeRef, ExprContext};
+use ruff_text_size::Ranged;
 
 use crate::builtins::builtins_scope;
 use crate::module_name::ModuleName;
@@ -40,6 +41,7 @@ use crate::semantic_index::expression::Expression;
 use crate::semantic_index::semantic_index;
 use crate::semantic_index::symbol::{FileScopeId, NodeWithScopeKind, NodeWithScopeRef, ScopeId};
 use crate::semantic_index::SemanticIndex;
+use crate::types::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
 use crate::types::{
     builtins_symbol_ty_by_name, definitions_ty, global_symbol_ty_by_name, ClassType, FunctionType,
     Name, Type, UnionBuilder,
@@ -123,13 +125,16 @@ pub(crate) enum InferenceRegion<'db> {
 }
 
 /// The inferred types for a single region.
-#[derive(Debug, Eq, PartialEq, Default, Clone)]
+#[derive(Debug, Eq, PartialEq, Default)]
 pub(crate) struct TypeInference<'db> {
     /// The types of every expression in this region.
     expressions: FxHashMap<ScopedExpressionId, Type<'db>>,
 
     /// The types of every definition in this region.
     definitions: FxHashMap<Definition<'db>, Type<'db>>,
+
+    /// The diagnostics for this region.
+    diagnostics: TypeCheckDiagnostics,
 }
 
 impl<'db> TypeInference<'db> {
@@ -142,9 +147,14 @@ impl<'db> TypeInference<'db> {
         self.definitions[&definition]
     }
 
+    pub(crate) fn diagnostics(&self) -> &[std::sync::Arc<TypeCheckDiagnostic>] {
+        &self.diagnostics
+    }
+
     fn shrink_to_fit(&mut self) {
         self.expressions.shrink_to_fit();
         self.definitions.shrink_to_fit();
+        self.diagnostics.shrink_to_fit();
     }
 }
 
@@ -235,6 +245,7 @@ impl<'db> TypeInferenceBuilder<'db> {
     fn extend(&mut self, inference: &TypeInference<'db>) {
         self.types.definitions.extend(inference.definitions.iter());
         self.types.expressions.extend(inference.expressions.iter());
+        self.types.diagnostics.extend(&inference.diagnostics);
     }
 
     /// Infers types in the given [`InferenceRegion`].
@@ -855,7 +866,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             asname: _,
         } = alias;
 
-        let module_ty = self.module_ty_from_name(ModuleName::new(name));
+        let module_ty = self.module_ty_from_name(ModuleName::new(name), alias.into());
         self.types.definitions.insert(definition, module_ty);
     }
 
@@ -953,7 +964,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             ModuleName::new(module_name)
         };
 
-        let module_ty = self.module_ty_from_name(module_name);
+        let module_ty = self.module_ty_from_name(module_name, import_from.into());
 
         let ast::Alias {
             range: _,
@@ -984,10 +995,26 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    fn module_ty_from_name(&self, module_name: Option<ModuleName>) -> Type<'db> {
-        module_name
-            .and_then(|module_name| resolve_module(self.db, module_name))
-            .map_or(Type::Unknown, |module| Type::Module(module.file()))
+    fn module_ty_from_name(
+        &mut self,
+        module_name: Option<ModuleName>,
+        node: AnyNodeRef,
+    ) -> Type<'db> {
+        let Some(module_name) = module_name else {
+            return Type::Unknown;
+        };
+
+        if let Some(module) = resolve_module(self.db, module_name.clone()) {
+            Type::Module(module.file())
+        } else {
+            self.add_diagnostic(
+                node,
+                "unresolved-import",
+                format_args!("Import '{module_name}' could not be resolved."),
+            );
+
+            Type::Unknown
+        }
     }
 
     fn infer_decorator(&mut self, decorator: &ast::Decorator) -> Type<'db> {
@@ -1059,7 +1086,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             ast::Expr::Yield(yield_expression) => self.infer_yield_expression(yield_expression),
             ast::Expr::YieldFrom(yield_from) => self.infer_yield_from_expression(yield_from),
             ast::Expr::Await(await_expression) => self.infer_await_expression(await_expression),
-            Expr::IpyEscapeCommand(_) => todo!("Implement Ipy escape command support"),
+            ast::Expr::IpyEscapeCommand(_) => todo!("Implement Ipy escape command support"),
         };
 
         let expr_id = expression.scoped_ast_id(self.db, self.scope);
@@ -1706,6 +1733,28 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
+    /// Adds a new diagnostic.
+    ///
+    /// The diagnostic does not get added if the rule isn't enabled for this file.
+    fn add_diagnostic(&mut self, node: AnyNodeRef, rule: &str, message: std::fmt::Arguments) {
+        if !self.db.is_file_open(self.file) {
+            return;
+        }
+
+        // TODO: Don't emit the diagnostic if:
+        // * The enclosing node contains any syntax errors
+        // * The rule is disabled for this file. We probably want to introduce a new query that
+        //   returns a rule selector for a given file that respects the package's settings,
+        //   any global pragma comments in the file, and any per-file-ignores.
+
+        self.types.diagnostics.push(TypeCheckDiagnostic {
+            file: self.file,
+            rule: rule.to_string(),
+            message: message.to_string(),
+            range: node.range(),
+        });
+    }
+
     pub(super) fn finish(mut self) -> TypeInference<'db> {
         self.infer_region();
         self.types.shrink_to_fit();