Use configured extension mapping to select code block language

Author: amyreese

Cargo.lock

@@ -500,6 +500,11 @@ impl ExtensionMapping {
         let ext = path.extension()?.to_str()?;
         self.0.get(ext).copied()
     }
+
+    /// Return the [`Language`] for a given file extension.
+    pub fn get_extension(&self, ext: &str) -> Option<Language> {
+        self.0.get(ext).copied()
+    }
 }
 
 impl From<FxHashMap<String, Language>> for ExtensionMapping {

crates/ruff_linter/src/settings/types.rs

@@ -17,8 +17,12 @@ ruff_source_file = { workspace = true }
 ruff_text_size = { workspace = true }
 ruff_workspace = { workspace = true }
 
-insta = { workspace = true }
 regex = { workspace = true }
 
+[dev-dependencies]
+ruff_linter = { workspace = true }
+
+insta = { workspace = true }
+
 [lints]
 workspace = true

crates/ruff_markdown/Cargo.toml

@@ -88,7 +88,11 @@ pub fn format_code_blocks(
                             let end = code_line.start();
                             let unformatted_code = dedent(&source[TextRange::new(start, end)]);
 
-                            let py_source_type = PySourceType::from_extension(&language);
+                            // map code block to source type, accounting for configured extension mappings
+                            let py_source_type = match settings.extension.get_extension(&language) {
+                                None => PySourceType::from_extension(&language),
+                                Some(language) => PySourceType::from(language),
+                            };
                             let options =
                                 settings.to_format_options(py_source_type, &unformatted_code, path);
 
@@ -128,6 +132,7 @@ pub fn format_code_blocks(
 #[cfg(test)]
 mod tests {
     use insta::assert_snapshot;
+    use ruff_linter::settings::types::{ExtensionMapping, ExtensionPair, Language};
     use ruff_workspace::FormatterSettings;
 
     use crate::{MarkdownResult, format_code_blocks};
@@ -380,4 +385,27 @@ print( 'hello' )
             &FormatterSettings::default()
         ), @"Unchanged");
     }
+
+    #[test]
+    fn format_code_blocks_extension_mapping() {
+        // format "py" mapped as "pyi" instead
+        let code = r#"
+```py
+def foo(): ...
+def bar(): ...
+```
+        "#;
+        let mapping = ExtensionMapping::from_iter([ExtensionPair {
+            extension: "py".to_string(),
+            language: Language::Pyi,
+        }]);
+        assert_snapshot!(format_code_blocks(
+            code,
+            None,
+            &FormatterSettings {
+                extension: mapping,
+                ..Default::default()
+            }
+        ), @"Unchanged");
+    }
 }