[red-knot] Add inlay type hints (#17214)

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/red_knot_ide/src/inlay_hints.rs

@@ -0,0 +1,279 @@
+use crate::Db;
+use red_knot_python_semantic::types::Type;
+use red_knot_python_semantic::{HasType, SemanticModel};
+use ruff_db::files::File;
+use ruff_db::parsed::parsed_module;
+use ruff_python_ast::visitor::source_order::{self, SourceOrderVisitor, TraversalSignal};
+use ruff_python_ast::{AnyNodeRef, Expr, Stmt};
+use ruff_text_size::{Ranged, TextRange, TextSize};
+use std::fmt;
+use std::fmt::Formatter;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct InlayHint<'db> {
+    pub position: TextSize,
+    pub content: InlayHintContent<'db>,
+}
+
+impl<'db> InlayHint<'db> {
+    pub const fn display(&self, db: &'db dyn Db) -> DisplayInlayHint<'_, 'db> {
+        self.content.display(db)
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum InlayHintContent<'db> {
+    Type(Type<'db>),
+    ReturnType(Type<'db>),
+}
+
+impl<'db> InlayHintContent<'db> {
+    pub const fn display(&self, db: &'db dyn Db) -> DisplayInlayHint<'_, 'db> {
+        DisplayInlayHint { db, hint: self }
+    }
+}
+
+pub struct DisplayInlayHint<'a, 'db> {
+    db: &'db dyn Db,
+    hint: &'a InlayHintContent<'db>,
+}
+
+impl fmt::Display for DisplayInlayHint<'_, '_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self.hint {
+            InlayHintContent::Type(ty) => {
+                write!(f, ": {}", ty.display(self.db.upcast()))
+            }
+            InlayHintContent::ReturnType(ty) => {
+                write!(f, " -> {}", ty.display(self.db.upcast()))
+            }
+        }
+    }
+}
+
+pub fn inlay_hints(db: &dyn Db, file: File, range: TextRange) -> Vec<InlayHint<'_>> {
+    let mut visitor = InlayHintVisitor::new(db, file, range);
+
+    let ast = parsed_module(db.upcast(), file);
+
+    visitor.visit_body(ast.suite());
+
+    visitor.hints
+}
+
+struct InlayHintVisitor<'db> {
+    model: SemanticModel<'db>,
+    hints: Vec<InlayHint<'db>>,
+    in_assignment: bool,
+    range: TextRange,
+}
+
+impl<'db> InlayHintVisitor<'db> {
+    fn new(db: &'db dyn Db, file: File, range: TextRange) -> Self {
+        Self {
+            model: SemanticModel::new(db.upcast(), file),
+            hints: Vec::new(),
+            in_assignment: false,
+            range,
+        }
+    }
+
+    fn add_type_hint(&mut self, position: TextSize, ty: Type<'db>) {
+        self.hints.push(InlayHint {
+            position,
+            content: InlayHintContent::Type(ty),
+        });
+    }
+}
+
+impl SourceOrderVisitor<'_> for InlayHintVisitor<'_> {
+    fn enter_node(&mut self, node: AnyNodeRef<'_>) -> TraversalSignal {
+        if self.range.intersect(node.range()).is_some() {
+            TraversalSignal::Traverse
+        } else {
+            TraversalSignal::Skip
+        }
+    }
+
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        let node = AnyNodeRef::from(stmt);
+
+        if !self.enter_node(node).is_traverse() {
+            return;
+        }
+
+        match stmt {
+            Stmt::Assign(assign) => {
+                self.in_assignment = true;
+                for target in &assign.targets {
+                    self.visit_expr(target);
+                }
+                self.in_assignment = false;
+
+                return;
+            }
+            // TODO
+            Stmt::FunctionDef(_) => {}
+            Stmt::For(_) => {}
+            Stmt::Expr(_) => {
+                // Don't traverse into expression statements because we don't show any hints.
+                return;
+            }
+            _ => {}
+        }
+
+        source_order::walk_stmt(self, stmt);
+    }
+
+    fn visit_expr(&mut self, expr: &'_ Expr) {
+        if !self.in_assignment {
+            return;
+        }
+
+        match expr {
+            Expr::Name(name) => {
+                if name.ctx.is_store() {
+                    let ty = expr.inferred_type(&self.model);
+                    self.add_type_hint(expr.range().end(), ty);
+                }
+            }
+            _ => {
+                source_order::walk_expr(self, expr);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use insta::assert_snapshot;
+    use ruff_db::{
+        files::{system_path_to_file, File},
+        source::source_text,
+    };
+    use ruff_text_size::TextSize;
+
+    use crate::db::tests::TestDb;
+
+    use red_knot_python_semantic::{
+        Program, ProgramSettings, PythonPath, PythonPlatform, SearchPathSettings,
+    };
+    use ruff_db::system::{DbWithWritableSystem, SystemPathBuf};
+    use ruff_python_ast::PythonVersion;
+
+    pub(super) fn inlay_hint_test(source: &str) -> InlayHintTest {
+        const START: &str = "<START>";
+        const END: &str = "<END>";
+
+        let mut db = TestDb::new();
+
+        let start = source.find(START);
+        let end = source
+            .find(END)
+            .map(|x| if start.is_some() { x - START.len() } else { x })
+            .unwrap_or(source.len());
+
+        let range = TextRange::new(
+            TextSize::try_from(start.unwrap_or_default()).unwrap(),
+            TextSize::try_from(end).unwrap(),
+        );
+
+        let source = source.replace(START, "");
+        let source = source.replace(END, "");
+
+        db.write_file("main.py", source)
+            .expect("write to memory file system to be successful");
+
+        let file = system_path_to_file(&db, "main.py").expect("newly written file to existing");
+
+        Program::from_settings(
+            &db,
+            ProgramSettings {
+                python_version: PythonVersion::latest(),
+                python_platform: PythonPlatform::default(),
+                search_paths: SearchPathSettings {
+                    extra_paths: vec![],
+                    src_roots: vec![SystemPathBuf::from("/")],
+                    custom_typeshed: None,
+                    python_path: PythonPath::KnownSitePackages(vec![]),
+                },
+            },
+        )
+        .expect("Default settings to be valid");
+
+        InlayHintTest { db, file, range }
+    }
+
+    pub(super) struct InlayHintTest {
+        pub(super) db: TestDb,
+        pub(super) file: File,
+        pub(super) range: TextRange,
+    }
+
+    impl InlayHintTest {
+        fn inlay_hints(&self) -> String {
+            let hints = inlay_hints(&self.db, self.file, self.range);
+
+            let mut buf = source_text(&self.db, self.file).as_str().to_string();
+
+            let mut offset = 0;
+
+            for hint in hints {
+                let end_position = (hint.position.to_u32() as usize) + offset;
+                let hint_str = format!("[{}]", hint.display(&self.db));
+                buf.insert_str(end_position, &hint_str);
+                offset += hint_str.len();
+            }
+
+            buf
+        }
+    }
+
+    #[test]
+    fn test_assign_statement() {
+        let test = inlay_hint_test("x = 1");
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        x[: Literal[1]] = 1
+        ");
+    }
+
+    #[test]
+    fn test_tuple_assignment() {
+        let test = inlay_hint_test("x, y = (1, 'abc')");
+
+        assert_snapshot!(test.inlay_hints(), @r#"
+        x[: Literal[1]], y[: Literal["abc"]] = (1, 'abc')
+        "#);
+    }
+
+    #[test]
+    fn test_nested_tuple_assignment() {
+        let test = inlay_hint_test("x, (y, z) = (1, ('abc', 2))");
+
+        assert_snapshot!(test.inlay_hints(), @r#"
+        x[: Literal[1]], (y[: Literal["abc"]], z[: Literal[2]]) = (1, ('abc', 2))
+        "#);
+    }
+
+    #[test]
+    fn test_assign_statement_with_type_annotation() {
+        let test = inlay_hint_test("x: int = 1");
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        x: int = 1
+        ");
+    }
+
+    #[test]
+    fn test_assign_statement_out_of_range() {
+        let test = inlay_hint_test("<START>x = 1<END>\ny = 2");
+
+        assert_snapshot!(test.inlay_hints(), @r"
+        x[: Literal[1]] = 1
+        y = 2
+        ");
+    }
+}

crates/red_knot_ide/src/lib.rs

@@ -2,11 +2,13 @@ mod db;
 mod find_node;
 mod goto;
 mod hover;
+mod inlay_hints;
 mod markup;
 
 pub use db::Db;
 pub use goto::goto_type_definition;
 pub use hover::hover;
+pub use inlay_hints::inlay_hints;
 pub use markup::MarkupKind;
 
 use rustc_hash::FxHashSet;

crates/red_knot_server/src/document.rs

@@ -8,7 +8,7 @@ mod text_document;
 pub(crate) use location::ToLink;
 use lsp_types::{PositionEncodingKind, Url};
 pub use notebook::NotebookDocument;
-pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, ToRangeExt};
+pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, TextSizeExt, ToRangeExt};
 pub(crate) use text_document::DocumentVersion;
 pub use text_document::TextDocument;
 

crates/red_knot_server/src/document/range.rs

@@ -28,6 +28,29 @@ pub(crate) trait PositionExt {
     fn to_text_size(&self, text: &str, index: &LineIndex, encoding: PositionEncoding) -> TextSize;
 }
 
+pub(crate) trait TextSizeExt {
+    fn to_position(
+        self,
+        text: &str,
+        index: &LineIndex,
+        encoding: PositionEncoding,
+    ) -> types::Position
+    where
+        Self: Sized;
+}
+
+impl TextSizeExt for TextSize {
+    fn to_position(
+        self,
+        text: &str,
+        index: &LineIndex,
+        encoding: PositionEncoding,
+    ) -> types::Position {
+        let source_location = offset_to_source_location(self, text, index, encoding);
+        source_location_to_position(&source_location)
+    }
+}
+
 pub(crate) trait ToRangeExt {
     fn to_lsp_range(
         &self,
@@ -107,18 +130,8 @@ impl ToRangeExt for TextRange {
         encoding: PositionEncoding,
     ) -> types::Range {
         types::Range {
-            start: source_location_to_position(&offset_to_source_location(
-                self.start(),
-                text,
-                index,
-                encoding,
-            )),
-            end: source_location_to_position(&offset_to_source_location(
-                self.end(),
-                text,
-                index,
-                encoding,
-            )),
+            start: self.start().to_position(text, index, encoding),
+            end: self.end().to_position(text, index, encoding),
         }
     }
 

crates/red_knot_server/src/server.rs

@@ -8,8 +8,9 @@ use std::panic::PanicInfo;
 use lsp_server::Message;
 use lsp_types::{
     ClientCapabilities, DiagnosticOptions, DiagnosticServerCapabilities, HoverProviderCapability,
-    MessageType, ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
-    TextDocumentSyncOptions, TypeDefinitionProviderCapability, Url,
+    InlayHintOptions, InlayHintServerCapabilities, MessageType, ServerCapabilities,
+    TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
+    TypeDefinitionProviderCapability, Url,
 };
 
 use self::connection::{Connection, ConnectionInitializer};
@@ -222,6 +223,9 @@ impl Server {
             )),
             type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
             hover_provider: Some(HoverProviderCapability::Simple(true)),
+            inlay_hint_provider: Some(lsp_types::OneOf::Right(
+                InlayHintServerCapabilities::Options(InlayHintOptions::default()),
+            )),
             ..Default::default()
         }
     }

crates/red_knot_server/src/server/api.rs

@@ -33,6 +33,9 @@ pub(super) fn request<'a>(req: server::Request) -> Task<'a> {
         request::HoverRequestHandler::METHOD => {
             background_request_task::<request::HoverRequestHandler>(req, BackgroundSchedule::Worker)
         }
+        request::InlayHintRequestHandler::METHOD => background_request_task::<
+            request::InlayHintRequestHandler,
+        >(req, BackgroundSchedule::Worker),
 
         method => {
             tracing::warn!("Received request {method} which does not have a handler");

crates/red_knot_server/src/server/api/requests.rs

@@ -1,7 +1,9 @@
 mod diagnostic;
 mod goto_type_definition;
 mod hover;
+mod inlay_hints;
 
 pub(super) use diagnostic::DocumentDiagnosticRequestHandler;
 pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
 pub(super) use hover::HoverRequestHandler;
+pub(super) use inlay_hints::InlayHintRequestHandler;

crates/red_knot_server/src/server/api/requests/inlay_hints.rs

@@ -0,0 +1,62 @@
+use std::borrow::Cow;
+
+use crate::document::{RangeExt, TextSizeExt};
+use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::server::client::Notifier;
+use crate::DocumentSnapshot;
+use lsp_types::request::InlayHintRequest;
+use lsp_types::{InlayHintParams, Url};
+use red_knot_ide::inlay_hints;
+use red_knot_project::ProjectDatabase;
+use ruff_db::source::{line_index, source_text};
+
+pub(crate) struct InlayHintRequestHandler;
+
+impl RequestHandler for InlayHintRequestHandler {
+    type RequestType = InlayHintRequest;
+}
+
+impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
+    fn document_url(params: &InlayHintParams) -> Cow<Url> {
+        Cow::Borrowed(&params.text_document.uri)
+    }
+
+    fn run_with_snapshot(
+        snapshot: DocumentSnapshot,
+        db: ProjectDatabase,
+        _notifier: Notifier,
+        params: InlayHintParams,
+    ) -> crate::server::Result<Option<Vec<lsp_types::InlayHint>>> {
+        let Some(file) = snapshot.file(&db) else {
+            tracing::debug!("Failed to resolve file for {:?}", params);
+            return Ok(None);
+        };
+
+        let index = line_index(&db, file);
+        let source = source_text(&db, file);
+
+        let range = params
+            .range
+            .to_text_range(&source, &index, snapshot.encoding());
+
+        let inlay_hints = inlay_hints(&db, file, range);
+
+        let inlay_hints = inlay_hints
+            .into_iter()
+            .map(|hint| lsp_types::InlayHint {
+                position: hint
+                    .position
+                    .to_position(&source, &index, snapshot.encoding()),
+                label: lsp_types::InlayHintLabel::String(hint.display(&db).to_string()),
+                kind: Some(lsp_types::InlayHintKind::TYPE),
+                tooltip: None,
+                padding_left: None,
+                padding_right: None,
+                data: None,
+                text_edits: None,
+            })
+            .collect();
+
+        Ok(Some(inlay_hints))
+    }
+}

crates/red_knot_wasm/src/lib.rs

@@ -1,7 +1,7 @@
 use std::any::Any;
 
 use js_sys::{Error, JsString};
-use red_knot_ide::{goto_type_definition, hover, MarkupKind};
+use red_knot_ide::{goto_type_definition, hover, inlay_hints, MarkupKind};
 use red_knot_project::metadata::options::Options;
 use red_knot_project::metadata::value::ValueSource;
 use red_knot_project::watch::{ChangeEvent, ChangedKind, CreatedKind, DeletedKind};
@@ -20,7 +20,7 @@ use ruff_db::Upcast;
 use ruff_notebook::Notebook;
 use ruff_python_formatter::formatted_file;
 use ruff_source_file::{LineIndex, OneIndexed, SourceLocation};
-use ruff_text_size::Ranged;
+use ruff_text_size::{Ranged, TextSize};
 use wasm_bindgen::prelude::*;
 
 #[wasm_bindgen(start)]
@@ -216,17 +216,7 @@ impl Workspace {
         let source = source_text(&self.db, file_id.file);
         let index = line_index(&self.db, file_id.file);
 
-        let offset = index.offset(
-            OneIndexed::new(position.line).ok_or_else(|| {
-                Error::new("Invalid value `0` for `position.line`. The line index is 1-indexed.")
-            })?,
-            OneIndexed::new(position.column).ok_or_else(|| {
-                Error::new(
-                    "Invalid value `0` for `position.column`. The column index is 1-indexed.",
-                )
-            })?,
-            &source,
-        );
+        let offset = position.to_text_size(&source, &index)?;
 
         let Some(targets) = goto_type_definition(&self.db, file_id.file, offset) else {
             return Ok(Vec::new());
@@ -258,17 +248,7 @@ impl Workspace {
         let source = source_text(&self.db, file_id.file);
         let index = line_index(&self.db, file_id.file);
 
-        let offset = index.offset(
-            OneIndexed::new(position.line).ok_or_else(|| {
-                Error::new("Invalid value `0` for `position.line`. The line index is 1-indexed.")
-            })?,
-            OneIndexed::new(position.column).ok_or_else(|| {
-                Error::new(
-                    "Invalid value `0` for `position.column`. The column index is 1-indexed.",
-                )
-            })?,
-            &source,
-        );
+        let offset = position.to_text_size(&source, &index)?;
 
         let Some(range_info) = hover(&self.db, file_id.file, offset) else {
             return Ok(None);
@@ -283,6 +263,26 @@ impl Workspace {
             range: source_range,
         }))
     }
+
+    #[wasm_bindgen(js_name = "inlayHints")]
+    pub fn inlay_hints(&self, file_id: &FileHandle, range: Range) -> Result<Vec<InlayHint>, Error> {
+        let index = line_index(&self.db, file_id.file);
+        let source = source_text(&self.db, file_id.file);
+
+        let result = inlay_hints(
+            &self.db,
+            file_id.file,
+            range.to_text_range(&index, &source)?,
+        );
+
+        Ok(result
+            .into_iter()
+            .map(|hint| InlayHint {
+                markdown: hint.display(&self.db).to_string(),
+                position: Position::from_text_size(hint.position, &index, &source),
+            })
+            .collect())
+    }
 }
 
 pub(crate) fn into_error<E: std::fmt::Display>(err: E) -> Error {
@@ -369,6 +369,14 @@ pub struct Range {
     pub end: Position,
 }
 
+#[wasm_bindgen]
+impl Range {
+    #[wasm_bindgen(constructor)]
+    pub fn new(start: Position, end: Position) -> Self {
+        Self { start, end }
+    }
+}
+
 impl Range {
     fn from_file_range(db: &dyn Db, file_range: FileRange) -> Self {
         let index = line_index(db.upcast(), file_range.file());
@@ -382,9 +390,21 @@ impl Range {
         line_index: &LineIndex,
         source: &str,
     ) -> Self {
-        let start = line_index.source_location(text_range.start(), source);
-        let end = line_index.source_location(text_range.end(), source);
-        Self::from((start, end))
+        Self {
+            start: Position::from_text_size(text_range.start(), line_index, source),
+            end: Position::from_text_size(text_range.end(), line_index, source),
+        }
+    }
+
+    fn to_text_range(
+        self,
+        line_index: &LineIndex,
+        source: &str,
+    ) -> Result<ruff_text_size::TextRange, Error> {
+        let start = self.start.to_text_size(source, line_index)?;
+        let end = self.end.to_text_size(source, line_index)?;
+
+        Ok(ruff_text_size::TextRange::new(start, end))
     }
 }
 
@@ -415,6 +435,28 @@ impl Position {
     }
 }
 
+impl Position {
+    fn to_text_size(self, text: &str, index: &LineIndex) -> Result<TextSize, Error> {
+        let text_size = index.offset(
+            OneIndexed::new(self.line).ok_or_else(|| {
+                Error::new("Invalid value `0` for `position.line`. The line index is 1-indexed.")
+            })?,
+            OneIndexed::new(self.column).ok_or_else(|| {
+                Error::new(
+                    "Invalid value `0` for `position.column`. The column index is 1-indexed.",
+                )
+            })?,
+            text,
+        );
+
+        Ok(text_size)
+    }
+
+    fn from_text_size(offset: TextSize, line_index: &LineIndex, source: &str) -> Self {
+        line_index.source_location(offset, source).into()
+    }
+}
+
 impl From<SourceLocation> for Position {
     fn from(location: SourceLocation) -> Self {
         Self {
@@ -433,13 +475,13 @@ pub enum Severity {
     Fatal,
 }
 
-impl From<ruff_db::diagnostic::Severity> for Severity {
-    fn from(value: ruff_db::diagnostic::Severity) -> Self {
+impl From<diagnostic::Severity> for Severity {
+    fn from(value: diagnostic::Severity) -> Self {
         match value {
-            ruff_db::diagnostic::Severity::Info => Self::Info,
-            ruff_db::diagnostic::Severity::Warning => Self::Warning,
-            ruff_db::diagnostic::Severity::Error => Self::Error,
-            ruff_db::diagnostic::Severity::Fatal => Self::Fatal,
+            diagnostic::Severity::Info => Self::Info,
+            diagnostic::Severity::Warning => Self::Warning,
+            diagnostic::Severity::Error => Self::Error,
+            diagnostic::Severity::Fatal => Self::Fatal,
         }
     }
 }
@@ -481,6 +523,14 @@ pub struct Hover {
     pub range: Range,
 }
 
+#[wasm_bindgen]
+pub struct InlayHint {
+    #[wasm_bindgen(getter_with_clone)]
+    pub markdown: String,
+
+    pub position: Position,
+}
+
 #[derive(Debug, Clone)]
 struct WasmSystem {
     fs: MemoryFileSystem,