[ty] Implemented support for "selection range" language server feature (#19567)

This PR adds support for the "selection range" language server feature.
This feature was recently requested by a ty user in [this feature
request](https://github.com/astral-sh/ty/issues/882).

This feature allows a client to implement "smart selection expansion"
based on the structure of the parse tree. For example, if you type
"shift-ctrl-right-arrow" in VS Code, the current selection will be
expanded to include the parent AST node. Conversely,
"shift-ctrl-left-arrow" shrinks the selection.

We will probably need to tune the granularity of selection expansion
based on user feedback. The initial implementation includes most AST
nodes, but users may find this to be too fine-grained. We have the
option of skipping some AST nodes that are not as meaningful when
editing code.

Co-authored-by: UnboundVariable <unbound@gmail.com>

crates/ty_ide/src/lib.rs

@@ -12,6 +12,7 @@ mod hover;
 mod inlay_hints;
 mod markup;
 mod references;
+mod selection_range;
 mod semantic_tokens;
 mod signature_help;
 mod stub_mapping;
@@ -28,6 +29,7 @@ pub use hover::hover;
 pub use inlay_hints::inlay_hints;
 pub use markup::MarkupKind;
 pub use references::ReferencesMode;
+pub use selection_range::selection_range;
 pub use semantic_tokens::{
     SemanticToken, SemanticTokenModifier, SemanticTokenType, SemanticTokens, semantic_tokens,
 };

crates/ty_ide/src/selection_range.rs

@@ -0,0 +1,304 @@
+use ruff_db::files::File;
+use ruff_db::parsed::parsed_module;
+use ruff_text_size::{Ranged, TextRange, TextSize};
+
+use crate::Db;
+use crate::find_node::covering_node;
+
+/// Returns a list of nested selection ranges, where each range contains the next one.
+/// The first range in the list is the largest range containing the cursor position.
+pub fn selection_range(db: &dyn Db, file: File, offset: TextSize) -> Vec<TextRange> {
+    let parsed = parsed_module(db, file).load(db);
+    let range = TextRange::new(offset, offset);
+
+    let covering = covering_node(parsed.syntax().into(), range);
+
+    let mut ranges = Vec::new();
+    for node in covering.ancestors() {
+        if should_include_in_selection(node) {
+            let range = node.range();
+            // Eliminate duplicates when parent and child nodes have the same range
+            if ranges.last() != Some(&range) {
+                ranges.push(range);
+            }
+        }
+    }
+
+    ranges
+}
+
+/// Determines if a node should be included in the selection range hierarchy.
+/// This filters out intermediate nodes that don't provide meaningful selections.
+fn should_include_in_selection(node: ruff_python_ast::AnyNodeRef) -> bool {
+    use ruff_python_ast::AnyNodeRef;
+
+    // We will likely need to tune this based on user feedback. Some users may
+    // prefer finer-grained selections while others may prefer coarser-grained.
+    match node {
+        // Exclude nodes that don't represent meaningful semantic units for selection
+        AnyNodeRef::StmtExpr(_) => false,
+
+        _ => true,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::CursorTest;
+    use insta::assert_snapshot;
+    use ruff_db::diagnostic::{Annotation, Diagnostic, DiagnosticId, LintName, Severity, Span};
+    use ruff_db::files::FileRange;
+    use ruff_text_size::Ranged;
+
+    /// Test selection range on a simple expression
+    #[test]
+    fn test_selection_range_simple_expression() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+x = 1 + <CURSOR>2
+",
+            )
+            .build();
+
+        assert_snapshot!(test.selection_range(), @r"
+        info[selection-range]: Selection Range 0
+         --> main.py:2:1
+          |
+        2 | x = 1 + 2
+          | ^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 1
+         --> main.py:2:5
+          |
+        2 | x = 1 + 2
+          |     ^^^^^
+          |
+
+        info[selection-range]: Selection Range 2
+         --> main.py:2:9
+          |
+        2 | x = 1 + 2
+          |         ^
+          |
+        ");
+    }
+
+    /// Test selection range on a function call
+    #[test]
+    fn test_selection_range_function_call() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+print(\"he<CURSOR>llo\")
+",
+            )
+            .build();
+
+        assert_snapshot!(test.selection_range(), @r#"
+        info[selection-range]: Selection Range 0
+         --> main.py:2:1
+          |
+        2 | print("hello")
+          | ^^^^^^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 1
+         --> main.py:2:6
+          |
+        2 | print("hello")
+          |      ^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 2
+         --> main.py:2:7
+          |
+        2 | print("hello")
+          |       ^^^^^^^
+          |
+        "#);
+    }
+
+    /// Test selection range on a function definition
+    #[test]
+    fn test_selection_range_function_definition() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+def my_<CURSOR>function():
+    return 42
+",
+            )
+            .build();
+
+        assert_snapshot!(test.selection_range(), @r"
+        info[selection-range]: Selection Range 0
+         --> main.py:2:1
+          |
+        2 | / def my_function():
+        3 | |     return 42
+          | |_____________^
+          |
+
+        info[selection-range]: Selection Range 1
+         --> main.py:2:5
+          |
+        2 | def my_function():
+          |     ^^^^^^^^^^^
+        3 |     return 42
+          |
+        ");
+    }
+
+    /// Test selection range on a class definition
+    #[test]
+    fn test_selection_range_class_definition() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+class My<CURSOR>Class:
+    def __init__(self):
+        self.value = 1
+",
+            )
+            .build();
+
+        assert_snapshot!(test.selection_range(), @r"
+        info[selection-range]: Selection Range 0
+         --> main.py:2:1
+          |
+        2 | / class MyClass:
+        3 | |     def __init__(self):
+        4 | |         self.value = 1
+          | |______________________^
+          |
+
+        info[selection-range]: Selection Range 1
+         --> main.py:2:7
+          |
+        2 | class MyClass:
+          |       ^^^^^^^
+        3 |     def __init__(self):
+        4 |         self.value = 1
+          |
+        ");
+    }
+
+    /// Test selection range on a deeply nested expression with comprehension, lambda, and subscript
+    #[test]
+    fn test_selection_range_deeply_nested_expression() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+result = [(lambda x: x[key.<CURSOR>attr])(item) for item in data if item is not None]
+",
+            )
+            .build();
+
+        assert_snapshot!(test.selection_range(), @r"
+        info[selection-range]: Selection Range 0
+         --> main.py:2:1
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 1
+         --> main.py:2:10
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          |          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 2
+         --> main.py:2:11
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          |           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 3
+         --> main.py:2:12
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          |            ^^^^^^^^^^^^^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 4
+         --> main.py:2:22
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          |                      ^^^^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 5
+         --> main.py:2:24
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          |                        ^^^^^^^^
+          |
+
+        info[selection-range]: Selection Range 6
+         --> main.py:2:28
+          |
+        2 | result = [(lambda x: x[key.attr])(item) for item in data if item is not None]
+          |                            ^^^^
+          |
+        ");
+    }
+
+    impl CursorTest {
+        fn selection_range(&self) -> String {
+            let ranges = selection_range(&self.db, self.cursor.file, self.cursor.offset);
+
+            if ranges.is_empty() {
+                return "No selection range found".to_string();
+            }
+
+            // Create one diagnostic per range for clearer visualization
+            let diagnostics: Vec<SelectionRangeDiagnostic> = ranges
+                .iter()
+                .enumerate()
+                .map(|(index, &range)| {
+                    SelectionRangeDiagnostic::new(FileRange::new(self.cursor.file, range), index)
+                })
+                .collect();
+
+            self.render_diagnostics(diagnostics)
+        }
+    }
+
+    struct SelectionRangeDiagnostic {
+        range: FileRange,
+        index: usize,
+    }
+
+    impl SelectionRangeDiagnostic {
+        fn new(range: FileRange, index: usize) -> Self {
+            Self { range, index }
+        }
+    }
+
+    impl crate::tests::IntoDiagnostic for SelectionRangeDiagnostic {
+        fn into_diagnostic(self) -> Diagnostic {
+            let mut diagnostic = Diagnostic::new(
+                DiagnosticId::Lint(LintName::of("selection-range")),
+                Severity::Info,
+                format!("Selection Range {}", self.index),
+            );
+
+            diagnostic.annotate(Annotation::primary(
+                Span::from(self.range.file()).with_range(self.range.range()),
+            ));
+
+            diagnostic
+        }
+    }
+}

crates/ty_server/src/server.rs

@@ -7,9 +7,10 @@ use lsp_server::Connection;
 use lsp_types::{
     ClientCapabilities, DeclarationCapability, DiagnosticOptions, DiagnosticServerCapabilities,
     HoverProviderCapability, InitializeParams, InlayHintOptions, InlayHintServerCapabilities,
-    MessageType, SemanticTokensLegend, SemanticTokensOptions, SemanticTokensServerCapabilities,
-    ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
-    TextDocumentSyncOptions, TypeDefinitionProviderCapability, Url, WorkDoneProgressOptions,
+    MessageType, SelectionRangeProviderCapability, SemanticTokensLegend, SemanticTokensOptions,
+    SemanticTokensServerCapabilities, ServerCapabilities, SignatureHelpOptions,
+    TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions,
+    TypeDefinitionProviderCapability, Url, WorkDoneProgressOptions,
 };
 use ruff_db::system::System;
 use std::num::NonZeroUsize;
@@ -241,6 +242,7 @@ impl Server {
                 trigger_characters: Some(vec!['.'.to_string()]),
                 ..Default::default()
             }),
+            selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
             document_symbol_provider: Some(lsp_types::OneOf::Left(true)),
             workspace_symbol_provider: Some(lsp_types::OneOf::Left(true)),
             ..Default::default()

crates/ty_server/src/server/api.rs

@@ -84,6 +84,9 @@ pub(super) fn request(req: server::Request) -> Task {
         >(
             req, BackgroundSchedule::LatencySensitive
         ),
+        requests::SelectionRangeRequestHandler::METHOD => background_document_request_task::<
+            requests::SelectionRangeRequestHandler,
+        >(req, BackgroundSchedule::Worker),
         requests::DocumentSymbolRequestHandler::METHOD => background_document_request_task::<
             requests::DocumentSymbolRequestHandler,
         >(req, BackgroundSchedule::Worker),

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

@@ -8,6 +8,7 @@ mod goto_references;
 mod goto_type_definition;
 mod hover;
 mod inlay_hints;
+mod selection_range;
 mod semantic_tokens;
 mod semantic_tokens_range;
 mod shutdown;
@@ -25,6 +26,7 @@ pub(super) use goto_references::ReferencesRequestHandler;
 pub(super) use goto_type_definition::GotoTypeDefinitionRequestHandler;
 pub(super) use hover::HoverRequestHandler;
 pub(super) use inlay_hints::InlayHintRequestHandler;
+pub(super) use selection_range::SelectionRangeRequestHandler;
 pub(super) use semantic_tokens::SemanticTokensRequestHandler;
 pub(super) use semantic_tokens_range::SemanticTokensRangeRequestHandler;
 pub(super) use shutdown::ShutdownHandler;

crates/ty_server/src/server/api/requests/selection_range.rs

@@ -0,0 +1,73 @@
+use std::borrow::Cow;
+
+use lsp_types::request::SelectionRangeRequest;
+use lsp_types::{SelectionRange as LspSelectionRange, SelectionRangeParams, Url};
+use ruff_db::source::{line_index, source_text};
+use ty_ide::selection_range;
+use ty_project::ProjectDatabase;
+
+use crate::document::{PositionExt, ToRangeExt};
+use crate::server::api::traits::{
+    BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
+};
+use crate::session::DocumentSnapshot;
+use crate::session::client::Client;
+
+pub(crate) struct SelectionRangeRequestHandler;
+
+impl RequestHandler for SelectionRangeRequestHandler {
+    type RequestType = SelectionRangeRequest;
+}
+
+impl BackgroundDocumentRequestHandler for SelectionRangeRequestHandler {
+    fn document_url(params: &SelectionRangeParams) -> Cow<Url> {
+        Cow::Borrowed(&params.text_document.uri)
+    }
+
+    fn run_with_snapshot(
+        db: &ProjectDatabase,
+        snapshot: DocumentSnapshot,
+        _client: &Client,
+        params: SelectionRangeParams,
+    ) -> crate::server::Result<Option<Vec<LspSelectionRange>>> {
+        if snapshot.client_settings().is_language_services_disabled() {
+            return Ok(None);
+        }
+
+        let Some(file) = snapshot.file(db) else {
+            return Ok(None);
+        };
+
+        let source = source_text(db, file);
+        let line_index = line_index(db, file);
+
+        let mut results = Vec::new();
+
+        for position in params.positions {
+            let offset = position.to_text_size(&source, &line_index, snapshot.encoding());
+
+            let ranges = selection_range(db, file, offset);
+            if !ranges.is_empty() {
+                // Convert ranges to nested LSP SelectionRange structure
+                let mut lsp_range = None;
+                for &range in &ranges {
+                    lsp_range = Some(LspSelectionRange {
+                        range: range.to_lsp_range(&source, &line_index, snapshot.encoding()),
+                        parent: lsp_range.map(Box::new),
+                    });
+                }
+                if let Some(range) = lsp_range {
+                    results.push(range);
+                }
+            }
+        }
+
+        if results.is_empty() {
+            Ok(None)
+        } else {
+            Ok(Some(results))
+        }
+    }
+}
+
+impl RetriableRequestHandler for SelectionRangeRequestHandler {}

crates/ty_server/tests/e2e/snapshots/e2e__initialize__initialization.snap

@@ -9,6 +9,7 @@ expression: initialization_result
       "openClose": true,
       "change": 2
     },
+    "selectionRangeProvider": true,
     "hoverProvider": true,
     "completionProvider": {
       "triggerCharacters": [

crates/ty_server/tests/e2e/snapshots/e2e__initialize__initialization_with_workspace.snap

@@ -9,6 +9,7 @@ expression: initialization_result
       "openClose": true,
       "change": 2
     },
+    "selectionRangeProvider": true,
     "hoverProvider": true,
     "completionProvider": {
       "triggerCharacters": [