LSP: Add support for custom extensions (#24463)

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

Author: silverstein

crates/ruff_server/Cargo.toml

@@ -48,6 +48,7 @@ libc = { workspace = true }
 
 [dev-dependencies]
 insta = { workspace = true, features = ["filters", "json"] }
+ruff_linter = { workspace = true, features = ["test-rules"] }
 dunce = { workspace = true }
 regex = { workspace = true }
 smallvec = { workspace = true }

crates/ruff_server/src/fix.rs

@@ -27,13 +27,13 @@ pub(crate) fn fix_all(
     linter_settings: &LinterSettings,
     encoding: PositionEncoding,
 ) -> crate::Result<Fixes> {
-    let source_kind = query.make_source_kind();
     let settings = query.settings();
     let document_path = query.virtual_file_path();
 
-    let SourceType::Python(source_type) = query.source_type() else {
+    let SourceType::Python(source_type) = query.source_type_for_lint() else {
         return Ok(Fixes::default());
     };
+    let source_kind = query.make_python_source_kind(source_type);
 
     // If the document is excluded, return an empty list of fixes.
     if is_document_excluded_for_linting(

crates/ruff_server/src/lint.rs

@@ -70,13 +70,13 @@ pub(crate) fn check(
     encoding: PositionEncoding,
     show_syntax_errors: bool,
 ) -> DiagnosticsMap {
-    let source_kind = query.make_source_kind();
     let settings = query.settings();
     let document_path = query.virtual_file_path();
 
-    let SourceType::Python(source_type) = query.source_type() else {
+    let SourceType::Python(source_type) = query.source_type_for_lint() else {
         return DiagnosticsMap::default();
     };
+    let source_kind = query.make_python_source_kind(source_type);
 
     // If the document is excluded, return an empty list of diagnostics.
     if is_document_excluded_for_linting(

crates/ruff_server/src/server/api/requests/code_action.rs

@@ -33,7 +33,7 @@ impl super::BackgroundDocumentRequestHandler for CodeActions {
         let query = snapshot.query();
 
         // Don't provide code actions for non-Python documents (e.g., markdown files).
-        let SourceType::Python(_) = query.source_type() else {
+        let SourceType::Python(_) = query.source_type_for_lint() else {
             return Ok(Some(response));
         };
 

crates/ruff_server/src/server/api/requests/format.rs

@@ -105,6 +105,7 @@ fn format_text_document(
 ) -> Result<super::FormatResponse> {
     let settings = query.settings();
     let file_path = query.virtual_file_path();
+    let source_type = query.source_type_for_format();
 
     // If the document is excluded, return early.
     if is_document_excluded_for_formatting(
@@ -119,7 +120,7 @@ fn format_text_document(
     let source = text_document.contents();
     let formatted = crate::format::format(
         text_document,
-        query.source_type(),
+        source_type,
         &settings.formatter,
         &file_path,
         backend,

crates/ruff_server/src/server/api/requests/format_range.rs

@@ -53,6 +53,7 @@ fn format_text_document_range(
 ) -> Result<super::FormatResponse> {
     let settings = query.settings();
     let file_path = query.virtual_file_path();
+    let source_type = query.source_type_for_format();
 
     // If the document is excluded, return early.
     if is_document_excluded_for_formatting(
@@ -69,7 +70,7 @@ fn format_text_document_range(
     let range = range.to_text_range(text, index, encoding);
     let formatted_range = crate::format::format_range(
         text_document,
-        query.source_type(),
+        source_type,
         &settings.formatter,
         range,
         &file_path,

crates/ruff_server/src/server/api/requests/hover.rs

@@ -33,7 +33,7 @@ pub(crate) fn hover(
     position: &types::TextDocumentPositionParams,
 ) -> Option<types::Hover> {
     // Don't show noqa hover for non-Python documents (e.g., markdown files).
-    let SourceType::Python(_) = snapshot.query().source_type() else {
+    let SourceType::Python(_) = snapshot.query().source_type_for_lint() else {
         return None;
     };
 

crates/ruff_server/src/session/index.rs

@@ -561,12 +561,15 @@ impl DocumentQuery {
         }
     }
 
-    /// Generate a source kind used by the linter.
-    pub(crate) fn make_source_kind(&self) -> ruff_linter::source_kind::SourceKind {
+    /// Generate a Python source kind used by the linter.
+    pub(crate) fn make_python_source_kind(
+        &self,
+        source_type: ruff_python_ast::PySourceType,
+    ) -> ruff_linter::source_kind::SourceKind {
         match self {
             Self::Text { document, .. } => ruff_linter::source_kind::SourceKind::Python {
                 code: document.contents().to_string(),
-                is_stub: ruff_python_ast::PySourceType::from(self.virtual_file_path()).is_stub(),
+                is_stub: source_type.is_stub(),
             },
             Self::Notebook { notebook, .. } => {
                 ruff_linter::source_kind::SourceKind::ipy_notebook(notebook.make_ruff_notebook())
@@ -582,10 +585,26 @@ impl DocumentQuery {
         }
     }
 
-    /// Get the source type of the document associated with this query.
-    pub(crate) fn source_type(&self) -> ruff_python_ast::SourceType {
+    /// Get the source type for linter-oriented operations.
+    pub(crate) fn source_type_for_lint(&self) -> ruff_python_ast::SourceType {
+        match self {
+            Self::Text { settings, .. } => settings
+                .linter
+                .extension
+                .get_source_type(&self.virtual_file_path()),
+            Self::Notebook { .. } => {
+                ruff_python_ast::SourceType::Python(ruff_python_ast::PySourceType::Ipynb)
+            }
+        }
+    }
+
+    /// Get the source type for formatter-oriented operations.
+    pub(crate) fn source_type_for_format(&self) -> ruff_python_ast::SourceType {
         match self {
-            Self::Text { .. } => ruff_python_ast::SourceType::from(self.virtual_file_path()),
+            Self::Text { settings, .. } => settings
+                .formatter
+                .extension
+                .get_source_type(&self.virtual_file_path()),
             Self::Notebook { .. } => {
                 ruff_python_ast::SourceType::Python(ruff_python_ast::PySourceType::Ipynb)
             }

crates/ruff_server/tests/e2e/code_action.rs

@@ -1,6 +1,5 @@
 use anyhow::Result;
 use insta::assert_json_snapshot;
-use lsp_types::notification::PublishDiagnostics;
 
 use crate::TestServerBuilder;
 
@@ -25,8 +24,6 @@ fn code_actions_for_python() -> Result<()> {
 
     server.open_text_document("test.py", "import os\n", 1);
 
-    server.await_notification::<PublishDiagnostics>();
-
     let actions = server
         .code_action_request("test.py", vec![])
         .expect("Expected Some response");

crates/ruff_server/tests/e2e/custom_extension.rs

@@ -0,0 +1,102 @@
+use anyhow::Result;
+use insta::assert_json_snapshot;
+use lsp_types::{Position, Range};
+
+use crate::TestServerBuilder;
+
+const CUSTOM_EXTENSION_CONFIG: &str = r#"[tool.ruff]
+preview = true
+extension = { thing = "markdown" }
+
+[tool.ruff.format]
+preview = true
+"#;
+
+const CUSTOM_EXTENSION_MARKDOWN: &str = "# title\n\n```python\nx='hi'\n```\n";
+
+#[test]
+fn format_custom_extension_mapped_to_markdown() -> Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(".")?
+        .with_file("pyproject.toml", CUSTOM_EXTENSION_CONFIG)?
+        .build();
+
+    server.open_text_document("test.thing", CUSTOM_EXTENSION_MARKDOWN, 1);
+
+    let edits = server.format_request("test.thing");
+
+    assert_json_snapshot!(
+        edits,
+        @r#"
+    [
+      {
+        "range": {
+          "start": {
+            "line": 3,
+            "character": 0
+          },
+          "end": {
+            "line": 4,
+            "character": 0
+          }
+        },
+        "newText": "x = \"hi\"\n"
+      }
+    ]
+    "#
+    );
+
+    Ok(())
+}
+
+#[test]
+fn range_format_custom_extension_mapped_to_markdown_is_unsupported() -> Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(".")?
+        .with_file("pyproject.toml", CUSTOM_EXTENSION_CONFIG)?
+        .build();
+
+    server.open_text_document("test.thing", CUSTOM_EXTENSION_MARKDOWN, 1);
+
+    let edits = server.format_range_request(
+        "test.thing",
+        Range {
+            start: Position {
+                line: 2,
+                character: 0,
+            },
+            end: Position {
+                line: 3,
+                character: 6,
+            },
+        },
+    );
+
+    assert_json_snapshot!(edits, @"null");
+
+    Ok(())
+}
+
+#[test]
+fn lint_custom_extension_mapped_to_markdown_emits_no_diagnostics() -> Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(".")?
+        .with_file("pyproject.toml", CUSTOM_EXTENSION_CONFIG)?
+        .build();
+
+    server.open_text_document("test.thing", CUSTOM_EXTENSION_MARKDOWN, 1);
+
+    let diagnostics = server.document_diagnostic_request("test.thing", None);
+
+    assert_json_snapshot!(
+        diagnostics,
+        @r#"
+    {
+      "kind": "full",
+      "items": []
+    }
+    "#
+    );
+
+    Ok(())
+}