diff --git a/crates/red_knot_ide/src/inlay_hints.rs b/crates/red_knot_ide/src/inlay_hints.rs new file mode 100644 index 0000000000..b282b0edfc --- /dev/null +++ b/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> { + 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>, + 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 = ""; + const END: &str = ""; + + 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("x = 1\ny = 2"); + + assert_snapshot!(test.inlay_hints(), @r" + x[: Literal[1]] = 1 + y = 2 + "); + } +} diff --git a/crates/red_knot_ide/src/lib.rs b/crates/red_knot_ide/src/lib.rs index 69f331f106..48f1145894 100644 --- a/crates/red_knot_ide/src/lib.rs +++ b/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; diff --git a/crates/red_knot_server/src/document.rs b/crates/red_knot_server/src/document.rs index 7d1c896044..7dca10720f 100644 --- a/crates/red_knot_server/src/document.rs +++ b/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; diff --git a/crates/red_knot_server/src/document/range.rs b/crates/red_knot_server/src/document/range.rs index 5fb42ab93d..0e268dcfb1 100644 --- a/crates/red_knot_server/src/document/range.rs +++ b/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), } } diff --git a/crates/red_knot_server/src/server.rs b/crates/red_knot_server/src/server.rs index ba0fa12c56..52e1adcb32 100644 --- a/crates/red_knot_server/src/server.rs +++ b/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() } } diff --git a/crates/red_knot_server/src/server/api.rs b/crates/red_knot_server/src/server/api.rs index f62d155a4e..478666ea61 100644 --- a/crates/red_knot_server/src/server/api.rs +++ b/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::(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"); diff --git a/crates/red_knot_server/src/server/api/requests.rs b/crates/red_knot_server/src/server/api/requests.rs index 6f9ace2b30..b6e907aa0c 100644 --- a/crates/red_knot_server/src/server/api/requests.rs +++ b/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; diff --git a/crates/red_knot_server/src/server/api/requests/inlay_hints.rs b/crates/red_knot_server/src/server/api/requests/inlay_hints.rs new file mode 100644 index 0000000000..f299fdf975 --- /dev/null +++ b/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 { + Cow::Borrowed(¶ms.text_document.uri) + } + + fn run_with_snapshot( + snapshot: DocumentSnapshot, + db: ProjectDatabase, + _notifier: Notifier, + params: InlayHintParams, + ) -> crate::server::Result>> { + 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)) + } +} diff --git a/crates/red_knot_wasm/src/lib.rs b/crates/red_knot_wasm/src/lib.rs index 0e475ecdc9..3faa4f79b1 100644 --- a/crates/red_knot_wasm/src/lib.rs +++ b/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, 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(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 { + 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 { + 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 for Position { fn from(location: SourceLocation) -> Self { Self { @@ -433,13 +475,13 @@ pub enum Severity { Fatal, } -impl From for Severity { - fn from(value: ruff_db::diagnostic::Severity) -> Self { +impl From 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, diff --git a/playground/knot/src/Editor/Editor.tsx b/playground/knot/src/Editor/Editor.tsx index 363704354e..3791f13e7d 100644 --- a/playground/knot/src/Editor/Editor.tsx +++ b/playground/knot/src/Editor/Editor.tsx @@ -12,15 +12,16 @@ import { languages, MarkerSeverity, Position, + Range, Uri, } from "monaco-editor"; import { useCallback, useEffect, useRef } from "react"; import { Theme } from "shared"; import { + Range as KnotRange, Severity, type Workspace, Position as KnotPosition, - type Range as KnotRange, } from "red_knot_wasm"; import IStandaloneCodeEditor = editor.IStandaloneCodeEditor; @@ -142,11 +143,13 @@ class PlaygroundServer languages.TypeDefinitionProvider, editor.ICodeEditorOpener, languages.HoverProvider, + languages.InlayHintsProvider, languages.DocumentFormattingEditProvider { private typeDefinitionProviderDisposable: IDisposable; private editorOpenerDisposable: IDisposable; private hoverDisposable: IDisposable; + private inlayHintsDisposable: IDisposable; private formatDisposable: IDisposable; constructor( @@ -159,11 +162,64 @@ class PlaygroundServer "python", this, ); + this.inlayHintsDisposable = monaco.languages.registerInlayHintsProvider( + "python", + this, + ); this.editorOpenerDisposable = monaco.editor.registerEditorOpener(this); this.formatDisposable = monaco.languages.registerDocumentFormattingEditProvider("python", this); } + provideInlayHints( + _model: editor.ITextModel, + range: Range, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + _token: CancellationToken, + ): languages.ProviderResult { + const workspace = this.props.workspace; + const selectedFile = this.props.files.selected; + + if (selectedFile == null) { + return; + } + + const selectedHandle = this.props.files.handles[selectedFile]; + + if (selectedHandle == null) { + return; + } + + const inlayHints = workspace.inlayHints( + selectedHandle, + iRangeToKnotRange(range), + ); + + if (inlayHints.length === 0) { + return undefined; + } + + return { + dispose: () => {}, + hints: inlayHints.map((hint) => ({ + label: hint.markdown, + position: { + lineNumber: hint.position.line, + column: hint.position.column, + }, + })), + }; + } + + resolveInlayHint( + // eslint-disable-next-line @typescript-eslint/no-unused-vars + _hint: languages.InlayHint, + // eslint-disable-next-line @typescript-eslint/no-unused-vars + _token: CancellationToken, + ): languages.ProviderResult { + return undefined; + } + update(props: PlaygroundServerProps) { this.props = props; } @@ -387,6 +443,7 @@ class PlaygroundServer this.hoverDisposable.dispose(); this.editorOpenerDisposable.dispose(); this.typeDefinitionProviderDisposable.dispose(); + this.inlayHintsDisposable.dispose(); this.formatDisposable.dispose(); } } @@ -399,3 +456,10 @@ function knotRangeToIRange(range: KnotRange): IRange { endColumn: range.end.column, }; } + +function iRangeToKnotRange(range: IRange): KnotRange { + return new KnotRange( + new KnotPosition(range.startLineNumber, range.startColumn), + new KnotPosition(range.endLineNumber, range.endColumn), + ); +} diff --git a/playground/package-lock.json b/playground/package-lock.json index 15a69e37df..b97f5647de 100644 --- a/playground/package-lock.json +++ b/playground/package-lock.json @@ -6094,7 +6094,7 @@ } }, "ruff/ruff_wasm": { - "version": "0.11.2", + "version": "0.11.3", "license": "MIT" }, "shared": {