repl_editor.rs

  1//! REPL operations on an [`Editor`].
  2
  3use std::ops::Range;
  4use std::sync::Arc;
  5
  6use anyhow::{Context as _, Result};
  7use editor::Editor;
  8use gpui::{App, Entity, WeakEntity, Window, prelude::*};
  9use language::{BufferSnapshot, Language, LanguageName, Point};
 10use project::{ProjectItem as _, WorktreeId};
 11
 12use crate::repl_store::ReplStore;
 13use crate::session::SessionEvent;
 14use crate::{
 15    ClearOutputs, Interrupt, JupyterSettings, KernelSpecification, Restart, Session, Shutdown,
 16};
 17
 18pub fn assign_kernelspec(
 19    kernel_specification: KernelSpecification,
 20    weak_editor: WeakEntity<Editor>,
 21    window: &mut Window,
 22    cx: &mut App,
 23) -> Result<()> {
 24    let store = ReplStore::global(cx);
 25    if !store.read(cx).is_enabled() {
 26        return Ok(());
 27    }
 28
 29    let worktree_id = crate::repl_editor::worktree_id_for_editor(weak_editor.clone(), cx)
 30        .context("editor is not in a worktree")?;
 31
 32    store.update(cx, |store, cx| {
 33        store.set_active_kernelspec(worktree_id, kernel_specification.clone(), cx);
 34    });
 35
 36    let fs = store.read(cx).fs().clone();
 37
 38    if let Some(session) = store.read(cx).get_session(weak_editor.entity_id()).cloned() {
 39        // Drop previous session, start new one
 40        session.update(cx, |session, cx| {
 41            session.clear_outputs(cx);
 42            session.shutdown(window, cx);
 43            cx.notify();
 44        });
 45    }
 46
 47    let session =
 48        cx.new(|cx| Session::new(weak_editor.clone(), fs, kernel_specification, window, cx));
 49
 50    weak_editor
 51        .update(cx, |_editor, cx| {
 52            cx.notify();
 53
 54            cx.subscribe(&session, {
 55                let store = store.clone();
 56                move |_this, _session, event, cx| match event {
 57                    SessionEvent::Shutdown(shutdown_event) => {
 58                        store.update(cx, |store, _cx| {
 59                            store.remove_session(shutdown_event.entity_id());
 60                        });
 61                    }
 62                }
 63            })
 64            .detach();
 65        })
 66        .ok();
 67
 68    store.update(cx, |store, _cx| {
 69        store.insert_session(weak_editor.entity_id(), session.clone());
 70    });
 71
 72    Ok(())
 73}
 74
 75pub fn run(
 76    editor: WeakEntity<Editor>,
 77    move_down: bool,
 78    window: &mut Window,
 79    cx: &mut App,
 80) -> Result<()> {
 81    let store = ReplStore::global(cx);
 82    if !store.read(cx).is_enabled() {
 83        return Ok(());
 84    }
 85
 86    let editor = editor.upgrade().context("editor was dropped")?;
 87    let selected_range = editor
 88        .update(cx, |editor, cx| editor.selections.newest_adjusted(cx))
 89        .range();
 90    let multibuffer = editor.read(cx).buffer().clone();
 91    let Some(buffer) = multibuffer.read(cx).as_singleton() else {
 92        return Ok(());
 93    };
 94
 95    let Some(project_path) = buffer.read(cx).project_path(cx) else {
 96        return Ok(());
 97    };
 98
 99    let (runnable_ranges, next_cell_point) =
100        runnable_ranges(&buffer.read(cx).snapshot(), selected_range, cx);
101
102    for runnable_range in runnable_ranges {
103        let Some(language) = multibuffer.read(cx).language_at(runnable_range.start, cx) else {
104            continue;
105        };
106
107        let kernel_specification = store
108            .read(cx)
109            .active_kernelspec(project_path.worktree_id, Some(language.clone()), cx)
110            .with_context(|| format!("No kernel found for language: {}", language.name()))?;
111
112        let fs = store.read(cx).fs().clone();
113
114        let session = if let Some(session) = store.read(cx).get_session(editor.entity_id()).cloned()
115        {
116            session
117        } else {
118            let weak_editor = editor.downgrade();
119            let session =
120                cx.new(|cx| Session::new(weak_editor, fs, kernel_specification, window, cx));
121
122            editor.update(cx, |_editor, cx| {
123                cx.notify();
124
125                cx.subscribe(&session, {
126                    let store = store.clone();
127                    move |_this, _session, event, cx| match event {
128                        SessionEvent::Shutdown(shutdown_event) => {
129                            store.update(cx, |store, _cx| {
130                                store.remove_session(shutdown_event.entity_id());
131                            });
132                        }
133                    }
134                })
135                .detach();
136            });
137
138            store.update(cx, |store, _cx| {
139                store.insert_session(editor.entity_id(), session.clone());
140            });
141
142            session
143        };
144
145        let selected_text;
146        let anchor_range;
147        let next_cursor;
148        {
149            let snapshot = multibuffer.read(cx).read(cx);
150            selected_text = snapshot
151                .text_for_range(runnable_range.clone())
152                .collect::<String>();
153            anchor_range = snapshot.anchor_before(runnable_range.start)
154                ..snapshot.anchor_after(runnable_range.end);
155            next_cursor = next_cell_point.map(|point| snapshot.anchor_after(point));
156        }
157
158        session.update(cx, |session, cx| {
159            session.execute(
160                selected_text,
161                anchor_range,
162                next_cursor,
163                move_down,
164                window,
165                cx,
166            );
167        });
168    }
169
170    anyhow::Ok(())
171}
172
173pub enum SessionSupport {
174    ActiveSession(Entity<Session>),
175    Inactive(KernelSpecification),
176    RequiresSetup(LanguageName),
177    Unsupported,
178}
179
180pub fn worktree_id_for_editor(editor: WeakEntity<Editor>, cx: &mut App) -> Option<WorktreeId> {
181    editor.upgrade().and_then(|editor| {
182        editor
183            .read(cx)
184            .buffer()
185            .read(cx)
186            .as_singleton()?
187            .read(cx)
188            .project_path(cx)
189            .map(|path| path.worktree_id)
190    })
191}
192
193pub fn session(editor: WeakEntity<Editor>, cx: &mut App) -> SessionSupport {
194    let store = ReplStore::global(cx);
195    let entity_id = editor.entity_id();
196
197    if let Some(session) = store.read(cx).get_session(entity_id).cloned() {
198        return SessionSupport::ActiveSession(session);
199    };
200
201    let Some(language) = get_language(editor.clone(), cx) else {
202        return SessionSupport::Unsupported;
203    };
204
205    let worktree_id = worktree_id_for_editor(editor, cx);
206
207    let Some(worktree_id) = worktree_id else {
208        return SessionSupport::Unsupported;
209    };
210
211    let kernelspec = store
212        .read(cx)
213        .active_kernelspec(worktree_id, Some(language.clone()), cx);
214
215    match kernelspec {
216        Some(kernelspec) => SessionSupport::Inactive(kernelspec),
217        None => {
218            // For language_supported, need to check available kernels for language
219            if language_supported(&language, cx) {
220                SessionSupport::RequiresSetup(language.name())
221            } else {
222                SessionSupport::Unsupported
223            }
224        }
225    }
226}
227
228pub fn clear_outputs(editor: WeakEntity<Editor>, cx: &mut App) {
229    let store = ReplStore::global(cx);
230    let entity_id = editor.entity_id();
231    let Some(session) = store.read(cx).get_session(entity_id).cloned() else {
232        return;
233    };
234    session.update(cx, |session, cx| {
235        session.clear_outputs(cx);
236        cx.notify();
237    });
238}
239
240pub fn interrupt(editor: WeakEntity<Editor>, cx: &mut App) {
241    let store = ReplStore::global(cx);
242    let entity_id = editor.entity_id();
243    let Some(session) = store.read(cx).get_session(entity_id).cloned() else {
244        return;
245    };
246
247    session.update(cx, |session, cx| {
248        session.interrupt(cx);
249        cx.notify();
250    });
251}
252
253pub fn shutdown(editor: WeakEntity<Editor>, window: &mut Window, cx: &mut App) {
254    let store = ReplStore::global(cx);
255    let entity_id = editor.entity_id();
256    let Some(session) = store.read(cx).get_session(entity_id).cloned() else {
257        return;
258    };
259
260    session.update(cx, |session, cx| {
261        session.shutdown(window, cx);
262        cx.notify();
263    });
264}
265
266pub fn restart(editor: WeakEntity<Editor>, window: &mut Window, cx: &mut App) {
267    let Some(editor) = editor.upgrade() else {
268        return;
269    };
270
271    let entity_id = editor.entity_id();
272
273    let Some(session) = ReplStore::global(cx)
274        .read(cx)
275        .get_session(entity_id)
276        .cloned()
277    else {
278        return;
279    };
280
281    session.update(cx, |session, cx| {
282        session.restart(window, cx);
283        cx.notify();
284    });
285}
286
287pub fn setup_editor_session_actions(editor: &mut Editor, editor_handle: WeakEntity<Editor>) {
288    editor
289        .register_action({
290            let editor_handle = editor_handle.clone();
291            move |_: &ClearOutputs, _, cx| {
292                if !JupyterSettings::enabled(cx) {
293                    return;
294                }
295
296                crate::clear_outputs(editor_handle.clone(), cx);
297            }
298        })
299        .detach();
300
301    editor
302        .register_action({
303            let editor_handle = editor_handle.clone();
304            move |_: &Interrupt, _, cx| {
305                if !JupyterSettings::enabled(cx) {
306                    return;
307                }
308
309                crate::interrupt(editor_handle.clone(), cx);
310            }
311        })
312        .detach();
313
314    editor
315        .register_action({
316            let editor_handle = editor_handle.clone();
317            move |_: &Shutdown, window, cx| {
318                if !JupyterSettings::enabled(cx) {
319                    return;
320                }
321
322                crate::shutdown(editor_handle.clone(), window, cx);
323            }
324        })
325        .detach();
326
327    editor
328        .register_action({
329            let editor_handle = editor_handle;
330            move |_: &Restart, window, cx| {
331                if !JupyterSettings::enabled(cx) {
332                    return;
333                }
334
335                crate::restart(editor_handle.clone(), window, cx);
336            }
337        })
338        .detach();
339}
340
341fn cell_range(buffer: &BufferSnapshot, start_row: u32, end_row: u32) -> Range<Point> {
342    let mut snippet_end_row = end_row;
343    while buffer.is_line_blank(snippet_end_row) && snippet_end_row > start_row {
344        snippet_end_row -= 1;
345    }
346    Point::new(start_row, 0)..Point::new(snippet_end_row, buffer.line_len(snippet_end_row))
347}
348
349// Returns the ranges of the snippets in the buffer and the next point for moving the cursor to
350fn jupytext_cells(
351    buffer: &BufferSnapshot,
352    range: Range<Point>,
353) -> (Vec<Range<Point>>, Option<Point>) {
354    let mut current_row = range.start.row;
355
356    let Some(language) = buffer.language() else {
357        return (Vec::new(), None);
358    };
359
360    let default_scope = language.default_scope();
361    let comment_prefixes = default_scope.line_comment_prefixes();
362    if comment_prefixes.is_empty() {
363        return (Vec::new(), None);
364    }
365
366    let jupytext_prefixes = comment_prefixes
367        .iter()
368        .map(|comment_prefix| format!("{comment_prefix}%%"))
369        .collect::<Vec<_>>();
370
371    let mut snippet_start_row = None;
372    loop {
373        if jupytext_prefixes
374            .iter()
375            .any(|prefix| buffer.contains_str_at(Point::new(current_row, 0), prefix))
376        {
377            snippet_start_row = Some(current_row);
378            break;
379        } else if current_row > 0 {
380            current_row -= 1;
381        } else {
382            break;
383        }
384    }
385
386    let mut snippets = Vec::new();
387    if let Some(mut snippet_start_row) = snippet_start_row {
388        for current_row in range.start.row + 1..=buffer.max_point().row {
389            if jupytext_prefixes
390                .iter()
391                .any(|prefix| buffer.contains_str_at(Point::new(current_row, 0), prefix))
392            {
393                snippets.push(cell_range(buffer, snippet_start_row, current_row - 1));
394
395                if current_row <= range.end.row {
396                    snippet_start_row = current_row;
397                } else {
398                    // Return our snippets as well as the next point for moving the cursor to
399                    return (snippets, Some(Point::new(current_row, 0)));
400                }
401            }
402        }
403
404        // Go to the end of the buffer (no more jupytext cells found)
405        snippets.push(cell_range(
406            buffer,
407            snippet_start_row,
408            buffer.max_point().row,
409        ));
410    }
411
412    (snippets, None)
413}
414
415fn runnable_ranges(
416    buffer: &BufferSnapshot,
417    range: Range<Point>,
418    cx: &mut App,
419) -> (Vec<Range<Point>>, Option<Point>) {
420    if let Some(language) = buffer.language()
421        && language.name() == "Markdown".into()
422    {
423        return (markdown_code_blocks(buffer, range, cx), None);
424    }
425
426    let (jupytext_snippets, next_cursor) = jupytext_cells(buffer, range.clone());
427    if !jupytext_snippets.is_empty() {
428        return (jupytext_snippets, next_cursor);
429    }
430
431    let snippet_range = cell_range(buffer, range.start.row, range.end.row);
432    let start_language = buffer.language_at(snippet_range.start);
433    let end_language = buffer.language_at(snippet_range.end);
434
435    if start_language
436        .zip(end_language)
437        .is_some_and(|(start, end)| start == end)
438    {
439        (vec![snippet_range], None)
440    } else {
441        (Vec::new(), None)
442    }
443}
444
445// We allow markdown code blocks to end in a trailing newline in order to render the output
446// below the final code fence. This is different than our behavior for selections and Jupytext cells.
447fn markdown_code_blocks(
448    buffer: &BufferSnapshot,
449    range: Range<Point>,
450    cx: &mut App,
451) -> Vec<Range<Point>> {
452    buffer
453        .injections_intersecting_range(range)
454        .filter(|(_, language)| language_supported(language, cx))
455        .map(|(content_range, _)| {
456            buffer.offset_to_point(content_range.start)..buffer.offset_to_point(content_range.end)
457        })
458        .collect()
459}
460
461fn language_supported(language: &Arc<Language>, cx: &mut App) -> bool {
462    let store = ReplStore::global(cx);
463    let store_read = store.read(cx);
464
465    // Since we're just checking for general language support, we only need to look at
466    // the pure Jupyter kernels - these are all the globally available ones
467    store_read.pure_jupyter_kernel_specifications().any(|spec| {
468        // Convert to lowercase for case-insensitive comparison since kernels might report "python" while our language is "Python"
469        spec.language().as_ref().to_lowercase() == language.name().as_ref().to_lowercase()
470    })
471}
472
473fn get_language(editor: WeakEntity<Editor>, cx: &mut App) -> Option<Arc<Language>> {
474    editor
475        .update(cx, |editor, cx| {
476            let selection = editor.selections.newest::<usize>(cx);
477            let buffer = editor.buffer().read(cx).snapshot(cx);
478            buffer.language_at(selection.head()).cloned()
479        })
480        .ok()
481        .flatten()
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487    use gpui::App;
488    use indoc::indoc;
489    use language::{Buffer, Language, LanguageConfig, LanguageRegistry};
490
491    #[gpui::test]
492    fn test_snippet_ranges(cx: &mut App) {
493        // Create a test language
494        let test_language = Arc::new(Language::new(
495            LanguageConfig {
496                name: "TestLang".into(),
497                line_comments: vec!["# ".into()],
498                ..Default::default()
499            },
500            None,
501        ));
502
503        let buffer = cx.new(|cx| {
504            Buffer::local(
505                indoc! { r#"
506                    print(1 + 1)
507                    print(2 + 2)
508
509                    print(4 + 4)
510
511
512                "# },
513                cx,
514            )
515            .with_language(test_language, cx)
516        });
517        let snapshot = buffer.read(cx).snapshot();
518
519        // Single-point selection
520        let (snippets, _) = runnable_ranges(&snapshot, Point::new(0, 4)..Point::new(0, 4), cx);
521        let snippets = snippets
522            .into_iter()
523            .map(|range| snapshot.text_for_range(range).collect::<String>())
524            .collect::<Vec<_>>();
525        assert_eq!(snippets, vec!["print(1 + 1)"]);
526
527        // Multi-line selection
528        let (snippets, _) = runnable_ranges(&snapshot, Point::new(0, 5)..Point::new(2, 0), cx);
529        let snippets = snippets
530            .into_iter()
531            .map(|range| snapshot.text_for_range(range).collect::<String>())
532            .collect::<Vec<_>>();
533        assert_eq!(
534            snippets,
535            vec![indoc! { r#"
536                print(1 + 1)
537                print(2 + 2)"# }]
538        );
539
540        // Trimming multiple trailing blank lines
541        let (snippets, _) = runnable_ranges(&snapshot, Point::new(0, 5)..Point::new(5, 0), cx);
542
543        let snippets = snippets
544            .into_iter()
545            .map(|range| snapshot.text_for_range(range).collect::<String>())
546            .collect::<Vec<_>>();
547        assert_eq!(
548            snippets,
549            vec![indoc! { r#"
550                print(1 + 1)
551                print(2 + 2)
552
553                print(4 + 4)"# }]
554        );
555    }
556
557    #[gpui::test]
558    fn test_jupytext_snippet_ranges(cx: &mut App) {
559        // Create a test language
560        let test_language = Arc::new(Language::new(
561            LanguageConfig {
562                name: "TestLang".into(),
563                line_comments: vec!["# ".into()],
564                ..Default::default()
565            },
566            None,
567        ));
568
569        let buffer = cx.new(|cx| {
570            Buffer::local(
571                indoc! { r#"
572                    # Hello!
573                    # %% [markdown]
574                    # This is some arithmetic
575                    print(1 + 1)
576                    print(2 + 2)
577
578                    # %%
579                    print(3 + 3)
580                    print(4 + 4)
581
582                    print(5 + 5)
583
584
585
586                "# },
587                cx,
588            )
589            .with_language(test_language, cx)
590        });
591        let snapshot = buffer.read(cx).snapshot();
592
593        // Jupytext snippet surrounding an empty selection
594        let (snippets, _) = runnable_ranges(&snapshot, Point::new(2, 5)..Point::new(2, 5), cx);
595
596        let snippets = snippets
597            .into_iter()
598            .map(|range| snapshot.text_for_range(range).collect::<String>())
599            .collect::<Vec<_>>();
600        assert_eq!(
601            snippets,
602            vec![indoc! { r#"
603                # %% [markdown]
604                # This is some arithmetic
605                print(1 + 1)
606                print(2 + 2)"# }]
607        );
608
609        // Jupytext snippets intersecting a non-empty selection
610        let (snippets, _) = runnable_ranges(&snapshot, Point::new(2, 5)..Point::new(6, 2), cx);
611        let snippets = snippets
612            .into_iter()
613            .map(|range| snapshot.text_for_range(range).collect::<String>())
614            .collect::<Vec<_>>();
615        assert_eq!(
616            snippets,
617            vec![
618                indoc! { r#"
619                    # %% [markdown]
620                    # This is some arithmetic
621                    print(1 + 1)
622                    print(2 + 2)"#
623                },
624                indoc! { r#"
625                    # %%
626                    print(3 + 3)
627                    print(4 + 4)
628
629                    print(5 + 5)"#
630                }
631            ]
632        );
633    }
634
635    #[gpui::test]
636    fn test_markdown_code_blocks(cx: &mut App) {
637        use crate::kernels::LocalKernelSpecification;
638        use jupyter_protocol::JupyterKernelspec;
639
640        // Initialize settings
641        settings::init(cx);
642        editor::init(cx);
643
644        // Initialize the ReplStore with a fake filesystem
645        let fs = Arc::new(project::RealFs::new(None, cx.background_executor().clone()));
646        ReplStore::init(fs, cx);
647
648        // Add mock kernel specifications for TypeScript and Python
649        let store = ReplStore::global(cx);
650        store.update(cx, |store, cx| {
651            let typescript_spec = KernelSpecification::Jupyter(LocalKernelSpecification {
652                name: "typescript".into(),
653                kernelspec: JupyterKernelspec {
654                    argv: vec![],
655                    display_name: "TypeScript".into(),
656                    language: "typescript".into(),
657                    interrupt_mode: None,
658                    metadata: None,
659                    env: None,
660                },
661                path: std::path::PathBuf::new(),
662            });
663
664            let python_spec = KernelSpecification::Jupyter(LocalKernelSpecification {
665                name: "python".into(),
666                kernelspec: JupyterKernelspec {
667                    argv: vec![],
668                    display_name: "Python".into(),
669                    language: "python".into(),
670                    interrupt_mode: None,
671                    metadata: None,
672                    env: None,
673                },
674                path: std::path::PathBuf::new(),
675            });
676
677            store.set_kernel_specs_for_testing(vec![typescript_spec, python_spec], cx);
678        });
679
680        let markdown = languages::language("markdown", tree_sitter_md::LANGUAGE.into());
681        let typescript = languages::language(
682            "typescript",
683            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
684        );
685        let python = languages::language("python", tree_sitter_python::LANGUAGE.into());
686        let language_registry = Arc::new(LanguageRegistry::new(cx.background_executor().clone()));
687        language_registry.add(markdown.clone());
688        language_registry.add(typescript);
689        language_registry.add(python);
690
691        // Two code blocks intersecting with selection
692        let buffer = cx.new(|cx| {
693            let mut buffer = Buffer::local(
694                indoc! { r#"
695                    Hey this is Markdown!
696
697                    ```typescript
698                    let foo = 999;
699                    console.log(foo + 1999);
700                    ```
701
702                    ```typescript
703                    console.log("foo")
704                    ```
705                    "#
706                },
707                cx,
708            );
709            buffer.set_language_registry(language_registry.clone());
710            buffer.set_language(Some(markdown.clone()), cx);
711            buffer
712        });
713        let snapshot = buffer.read(cx).snapshot();
714
715        let (snippets, _) = runnable_ranges(&snapshot, Point::new(3, 5)..Point::new(8, 5), cx);
716        let snippets = snippets
717            .into_iter()
718            .map(|range| snapshot.text_for_range(range).collect::<String>())
719            .collect::<Vec<_>>();
720
721        assert_eq!(
722            snippets,
723            vec![
724                indoc! { r#"
725                    let foo = 999;
726                    console.log(foo + 1999);
727                    "#
728                },
729                "console.log(\"foo\")\n"
730            ]
731        );
732
733        // Three code blocks intersecting with selection
734        let buffer = cx.new(|cx| {
735            let mut buffer = Buffer::local(
736                indoc! { r#"
737                    Hey this is Markdown!
738
739                    ```typescript
740                    let foo = 999;
741                    console.log(foo + 1999);
742                    ```
743
744                    ```ts
745                    console.log("foo")
746                    ```
747
748                    ```typescript
749                    console.log("another code block")
750                    ```
751                "# },
752                cx,
753            );
754            buffer.set_language_registry(language_registry.clone());
755            buffer.set_language(Some(markdown.clone()), cx);
756            buffer
757        });
758        let snapshot = buffer.read(cx).snapshot();
759
760        let (snippets, _) = runnable_ranges(&snapshot, Point::new(3, 5)..Point::new(12, 5), cx);
761        let snippets = snippets
762            .into_iter()
763            .map(|range| snapshot.text_for_range(range).collect::<String>())
764            .collect::<Vec<_>>();
765
766        assert_eq!(
767            snippets,
768            vec![
769                indoc! { r#"
770                    let foo = 999;
771                    console.log(foo + 1999);
772                    "#
773                },
774                "console.log(\"foo\")\n",
775                "console.log(\"another code block\")\n",
776            ]
777        );
778
779        // Python code block
780        let buffer = cx.new(|cx| {
781            let mut buffer = Buffer::local(
782                indoc! { r#"
783                    Hey this is Markdown!
784
785                    ```python
786                    print("hello there")
787                    print("hello there")
788                    print("hello there")
789                    ```
790                "# },
791                cx,
792            );
793            buffer.set_language_registry(language_registry.clone());
794            buffer.set_language(Some(markdown.clone()), cx);
795            buffer
796        });
797        let snapshot = buffer.read(cx).snapshot();
798
799        let (snippets, _) = runnable_ranges(&snapshot, Point::new(4, 5)..Point::new(5, 5), cx);
800        let snippets = snippets
801            .into_iter()
802            .map(|range| snapshot.text_for_range(range).collect::<String>())
803            .collect::<Vec<_>>();
804
805        assert_eq!(
806            snippets,
807            vec![indoc! { r#"
808                print("hello there")
809                print("hello there")
810                print("hello there")
811                "#
812            },]
813        );
814    }
815}