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

Co-authored-by: Micha Reiser <[email protected]>

Author: MatthewMckee4

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;