language: Fix auto-indentation for Python code blocks in Markdown (#43853)

Jeff Brennan and Smit Barmase created

Closes #43722

Release Notes:

- Fixed an issue where auto-indentation didn’t work correctly for Python
code blocks in Markdown.

---------

Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>

Change summary

crates/editor/src/editor_tests.rs | 42 +++++++++++++++++++++++++++++++++
crates/language/src/buffer.rs     | 35 +++++++++++++++++++-------
2 files changed, 67 insertions(+), 10 deletions(-)

Detailed changes

crates/editor/src/editor_tests.rs 🔗

@@ -25972,6 +25972,48 @@ async fn test_indent_on_newline_for_python(cx: &mut TestAppContext) {
     "});
 }
 
+#[gpui::test]
+async fn test_python_indent_in_markdown(cx: &mut TestAppContext) {
+    init_test(cx, |_| {});
+
+    let language_registry = Arc::new(language::LanguageRegistry::test(cx.executor()));
+    let python_lang = languages::language("python", tree_sitter_python::LANGUAGE.into());
+    language_registry.add(markdown_lang());
+    language_registry.add(python_lang);
+
+    let mut cx = EditorTestContext::new(cx).await;
+    cx.update_buffer(|buffer, cx| {
+        buffer.set_language_registry(language_registry);
+        buffer.set_language(Some(markdown_lang()), cx);
+    });
+
+    // Test that `else:` correctly outdents to match `if:` inside the Python code block
+    cx.set_state(indoc! {"
+        # Heading
+
+        ```python
+        def main():
+            if condition:
+                pass
+                ˇ
+        ```
+    "});
+    cx.update_editor(|editor, window, cx| {
+        editor.handle_input("else:", window, cx);
+    });
+    cx.run_until_parked();
+    cx.assert_editor_state(indoc! {"
+        # Heading
+
+        ```python
+        def main():
+            if condition:
+                pass
+            else:ˇ
+        ```
+    "});
+}
+
 #[gpui::test]
 async fn test_tab_in_leading_whitespace_auto_indents_for_bash(cx: &mut TestAppContext) {
     init_test(cx, |_| {});

crates/language/src/buffer.rs 🔗

@@ -3216,6 +3216,7 @@ impl BufferSnapshot {
         struct StartPosition {
             start: Point,
             suffix: SharedString,
+            language: Arc<Language>,
         }
 
         // Find the suggested indentation ranges based on the syntax tree.
@@ -3259,6 +3260,7 @@ impl BufferSnapshot {
                     start_positions.push(StartPosition {
                         start: Point::from_ts_point(capture.node.start_position()),
                         suffix: suffix.clone(),
+                        language: mat.language.clone(),
                     });
                 }
             }
@@ -3319,7 +3321,7 @@ impl BufferSnapshot {
 
         // Find the suggested indentation increases and decreased based on regexes.
         let mut regex_outdent_map = HashMap::default();
-        let mut last_seen_suffix: HashMap<String, Vec<Point>> = HashMap::default();
+        let mut last_seen_suffix: HashMap<String, Vec<StartPosition>> = HashMap::default();
         let mut start_positions_iter = start_positions.iter().peekable();
 
         let mut indent_change_rows = Vec::<(u32, Ordering)>::new();
@@ -3327,14 +3329,21 @@ impl BufferSnapshot {
             Point::new(prev_non_blank_row.unwrap_or(row_range.start), 0)
                 ..Point::new(row_range.end, 0),
             |row, line| {
-                if config
+                let indent_len = self.indent_size_for_line(row).len;
+                let row_language = self.language_at(Point::new(row, indent_len)).cloned();
+                let row_language_config = row_language
+                    .as_ref()
+                    .map(|lang| lang.config())
+                    .unwrap_or(config);
+
+                if row_language_config
                     .decrease_indent_pattern
                     .as_ref()
                     .is_some_and(|regex| regex.is_match(line))
                 {
                     indent_change_rows.push((row, Ordering::Less));
                 }
-                if config
+                if row_language_config
                     .increase_indent_pattern
                     .as_ref()
                     .is_some_and(|regex| regex.is_match(line))
@@ -3343,16 +3352,16 @@ impl BufferSnapshot {
                 }
                 while let Some(pos) = start_positions_iter.peek() {
                     if pos.start.row < row {
-                        let pos = start_positions_iter.next().unwrap();
+                        let pos = start_positions_iter.next().unwrap().clone();
                         last_seen_suffix
                             .entry(pos.suffix.to_string())
                             .or_default()
-                            .push(pos.start);
+                            .push(pos);
                     } else {
                         break;
                     }
                 }
-                for rule in &config.decrease_indent_patterns {
+                for rule in &row_language_config.decrease_indent_patterns {
                     if rule.pattern.as_ref().is_some_and(|r| r.is_match(line)) {
                         let row_start_column = self.indent_size_for_line(row).len;
                         let basis_row = rule
@@ -3360,10 +3369,16 @@ impl BufferSnapshot {
                             .iter()
                             .filter_map(|valid_suffix| last_seen_suffix.get(valid_suffix))
                             .flatten()
-                            .filter(|start_point| start_point.column <= row_start_column)
-                            .max_by_key(|start_point| start_point.row);
-                        if let Some(outdent_to_row) = basis_row {
-                            regex_outdent_map.insert(row, outdent_to_row.row);
+                            .filter(|pos| {
+                                row_language
+                                    .as_ref()
+                                    .or(self.language.as_ref())
+                                    .is_some_and(|lang| Arc::ptr_eq(lang, &pos.language))
+                            })
+                            .filter(|pos| pos.start.column <= row_start_column)
+                            .max_by_key(|pos| pos.start.row);
+                        if let Some(outdent_to) = basis_row {
+                            regex_outdent_map.insert(row, outdent_to.start.row);
                         }
                         break;
                     }