[red-knot] Add inlay type hints

crates/red_knot_ide/src/inlay_hints.rs

@@ -0,0 +1,193 @@
+use crate::{Db, RangedValue};
+use red_knot_python_semantic::types::Type;
+use red_knot_python_semantic::{HasType, SemanticModel};
+use ruff_db::files::{File, FileRange};
+use ruff_db::parsed::parsed_module;
+use ruff_python_ast::visitor::source_order::{self, SourceOrderVisitor};
+use ruff_python_ast::{Expr, Stmt};
+use ruff_text_size::{Ranged, TextRange};
+use std::fmt;
+use std::fmt::Formatter;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum InlayHintContent<'db> {
+    AssignStatement(Type<'db>),
+    FunctionReturnType(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::AssignStatement(ty) => {
+                write!(f, ": {}", ty.display(self.db.upcast()))
+            }
+            InlayHintContent::FunctionReturnType(ty) => {
+                write!(f, " -> {}", ty.display(self.db.upcast()))
+            }
+        }
+    }
+}
+
+pub fn inlay_hints(db: &dyn Db, file: File) -> Vec<RangedValue<InlayHintContent<'_>>> {
+    let mut visitor = InlayHintVisitor::new(db, file);
+
+    let ast = parsed_module(db.upcast(), file);
+
+    visitor.visit_body(ast.suite());
+
+    let hints = visitor.hints().clone();
+
+    hints
+}
+
+struct InlayHintVisitor<'db> {
+    model: SemanticModel<'db>,
+    file: File,
+    hints: Vec<RangedValue<InlayHintContent<'db>>>,
+}
+
+impl<'db> InlayHintVisitor<'db> {
+    fn new(db: &'db dyn Db, file: File) -> Self {
+        Self {
+            model: SemanticModel::new(db.upcast(), file),
+            file,
+            hints: Vec::new(),
+        }
+    }
+
+    fn hints(&self) -> &Vec<RangedValue<InlayHintContent<'db>>> {
+        &self.hints
+    }
+
+    fn add_hint(&mut self, range: TextRange, ty: Type<'db>) {
+        self.hints.push(RangedValue {
+            range: FileRange::new(self.file, range),
+            value: InlayHintContent::AssignStatement(ty),
+        });
+    }
+}
+
+impl SourceOrderVisitor<'_> for InlayHintVisitor<'_> {
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        match stmt {
+            Stmt::Assign(assign) => {
+                let ty = assign.value.inferred_type(&self.model);
+                for target in &assign.targets {
+                    match target {
+                        Expr::Tuple(tuple) => {
+                            for element in &tuple.elts {
+                                let element_ty = element.inferred_type(&self.model);
+                                self.add_hint(element.range(), element_ty);
+                            }
+                        }
+                        _ => {
+                            self.add_hint(target.range(), ty);
+                        }
+                    }
+                }
+                return;
+            }
+            // TODO
+            Stmt::FunctionDef(_) => {}
+            Stmt::For(_) => {}
+            Stmt::Expr(_) => {}
+            _ => {}
+        }
+
+        source_order::walk_stmt(self, stmt);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use red_knot_python_semantic::types::StringLiteralType;
+    use ruff_db::files::{system_path_to_file, File};
+    use ruff_db::system::DbWithWritableSystem as _;
+    use ruff_text_size::TextSize;
+
+    use crate::db::tests::TestDb;
+
+    struct TestCase {
+        db: TestDb,
+        file: File,
+    }
+
+    fn test_case(content: impl AsRef<str>) -> TestCase {
+        let mut db = TestDb::new();
+        db.write_file("test.py", content).unwrap();
+
+        let file = system_path_to_file(&db, "test.py").unwrap();
+
+        TestCase { db, file }
+    }
+
+    #[test]
+    fn test_assign_statement() {
+        let test_case = test_case("x = 1");
+        let hints = get_inlay_hints(&test_case.db, test_case.file);
+        assert_eq!(hints.len(), 1);
+        assert_eq!(
+            hints[0].value,
+            InlayHintContent::AssignStatement(Type::IntLiteral(1))
+        );
+        assert_eq!(
+            hints[0].range,
+            FileRange::new(
+                test_case.file,
+                TextRange::new(TextSize::from(0), TextSize::from(1))
+            )
+        );
+    }
+
+    #[test]
+    fn test_tuple_assignment() {
+        let test_case = test_case("x, y = (1, 'abc')");
+        let hints = get_inlay_hints(&test_case.db, test_case.file);
+        assert_eq!(hints.len(), 2);
+        assert_eq!(
+            hints[0].value,
+            InlayHintContent::AssignStatement(Type::IntLiteral(1))
+        );
+        assert_eq!(
+            hints[1].value,
+            InlayHintContent::AssignStatement(Type::StringLiteral(StringLiteralType::new(
+                &test_case.db,
+                "abc"
+            )))
+        );
+        assert_eq!(
+            hints[0].range,
+            FileRange::new(
+                test_case.file,
+                TextRange::new(TextSize::from(0), TextSize::from(1))
+            )
+        );
+        assert_eq!(
+            hints[1].range,
+            FileRange::new(
+                test_case.file,
+                TextRange::new(TextSize::from(3), TextSize::from(4))
+            )
+        );
+    }
+
+    #[test]
+    fn test_assign_statement_with_type_annotation() {
+        let test_case = test_case("x: int = 1");
+        let hints = get_inlay_hints(&test_case.db, test_case.file);
+        assert_eq!(hints.len(), 0);
+    }
+}

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::get_inlay_hints;
 pub use markup::MarkupKind;
 
 use rustc_hash::FxHashSet;

crates/red_knot_project/src/metadata/options.rs

@@ -32,6 +32,9 @@ pub struct Options {
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub terminal: Option<TerminalOptions>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub editor: Option<EditorOptions>,
 }
 
 impl Options {
@@ -407,3 +410,14 @@ impl OptionDiagnostic {
         }
     }
 }
+
+#[derive(Debug, Default, Clone, Eq, PartialEq, Combine, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case", deny_unknown_fields)]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct EditorOptions {
+    /// Whether to show inlay hints.
+    ///
+    /// Defaults to `false`.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub inlay_hints: Option<bool>,
+}

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

@@ -35,6 +35,11 @@ pub(super) fn request<'a>(req: server::Request) -> Task<'a> {
         request::HoverRequestHandler::METHOD => background_request_task::<
             request::HoverRequestHandler,
         >(req, BackgroundSchedule::LatencySensitive),
+        request::InlayHintRequestHandler::METHOD => background_request_task::<
+            request::InlayHintRequestHandler,
+        >(
+            req, BackgroundSchedule::LatencySensitive
+        ),
 
         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,78 @@
+use std::borrow::Cow;
+
+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::get_inlay_hints;
+use red_knot_project::Db;
+use red_knot_project::ProjectDatabase;
+use ruff_db::source::{line_index, source_text};
+use ruff_text_size::Ranged;
+
+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 editor_options = db
+            .project()
+            .metadata(&db)
+            .options()
+            .editor
+            .clone()
+            .unwrap_or_default();
+
+        if !editor_options.inlay_hints.unwrap_or(false) {
+            return Ok(None);
+        }
+
+        let inlay_hints = get_inlay_hints(&db, file);
+
+        let index = line_index(&db, file);
+        let source = source_text(&db, file);
+
+        let inlay_hints = inlay_hints
+            .into_iter()
+            .map(|hint| {
+                let end = index.source_location(hint.range.range().end(), &source);
+
+                lsp_types::InlayHint {
+                    position: lsp_types::Position {
+                        line: u32::try_from(end.row.to_zero_indexed())
+                            .expect("row usize fits in u32"),
+                        character: u32::try_from(end.column.to_zero_indexed())
+                            .expect("character usize fits in u32"),
+                    },
+                    label: lsp_types::InlayHintLabel::String(hint.display(&db).to_string()),
+                    kind: None,
+                    tooltip: None,
+                    padding_left: None,
+                    padding_right: None,
+                    data: None,
+                    text_edits: None,
+                }
+            })
+            .collect();
+
+        Ok(Some(inlay_hints))
+    }
+}