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);
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.clone(), 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            if language_supported(&language.clone()) {
219                SessionSupport::RequiresSetup(language.name())
220            } else {
221                SessionSupport::Unsupported
222            }
223        }
224    }
225}
226
227pub fn clear_outputs(editor: WeakEntity<Editor>, cx: &mut App) {
228    let store = ReplStore::global(cx);
229    let entity_id = editor.entity_id();
230    let Some(session) = store.read(cx).get_session(entity_id).cloned() else {
231        return;
232    };
233    session.update(cx, |session, cx| {
234        session.clear_outputs(cx);
235        cx.notify();
236    });
237}
238
239pub fn interrupt(editor: WeakEntity<Editor>, cx: &mut App) {
240    let store = ReplStore::global(cx);
241    let entity_id = editor.entity_id();
242    let Some(session) = store.read(cx).get_session(entity_id).cloned() else {
243        return;
244    };
245
246    session.update(cx, |session, cx| {
247        session.interrupt(cx);
248        cx.notify();
249    });
250}
251
252pub fn shutdown(editor: WeakEntity<Editor>, window: &mut Window, cx: &mut App) {
253    let store = ReplStore::global(cx);
254    let entity_id = editor.entity_id();
255    let Some(session) = store.read(cx).get_session(entity_id).cloned() else {
256        return;
257    };
258
259    session.update(cx, |session, cx| {
260        session.shutdown(window, cx);
261        cx.notify();
262    });
263}
264
265pub fn restart(editor: WeakEntity<Editor>, window: &mut Window, cx: &mut App) {
266    let Some(editor) = editor.upgrade() else {
267        return;
268    };
269
270    let entity_id = editor.entity_id();
271
272    let Some(session) = ReplStore::global(cx)
273        .read(cx)
274        .get_session(entity_id)
275        .cloned()
276    else {
277        return;
278    };
279
280    session.update(cx, |session, cx| {
281        session.restart(window, cx);
282        cx.notify();
283    });
284}
285
286pub fn setup_editor_session_actions(editor: &mut Editor, editor_handle: WeakEntity<Editor>) {
287    editor
288        .register_action({
289            let editor_handle = editor_handle.clone();
290            move |_: &ClearOutputs, _, cx| {
291                if !JupyterSettings::enabled(cx) {
292                    return;
293                }
294
295                crate::clear_outputs(editor_handle.clone(), cx);
296            }
297        })
298        .detach();
299
300    editor
301        .register_action({
302            let editor_handle = editor_handle.clone();
303            move |_: &Interrupt, _, cx| {
304                if !JupyterSettings::enabled(cx) {
305                    return;
306                }
307
308                crate::interrupt(editor_handle.clone(), cx);
309            }
310        })
311        .detach();
312
313    editor
314        .register_action({
315            let editor_handle = editor_handle.clone();
316            move |_: &Shutdown, window, cx| {
317                if !JupyterSettings::enabled(cx) {
318                    return;
319                }
320
321                crate::shutdown(editor_handle.clone(), window, cx);
322            }
323        })
324        .detach();
325
326    editor
327        .register_action({
328            let editor_handle = editor_handle.clone();
329            move |_: &Restart, window, cx| {
330                if !JupyterSettings::enabled(cx) {
331                    return;
332                }
333
334                crate::restart(editor_handle.clone(), window, cx);
335            }
336        })
337        .detach();
338}
339
340fn cell_range(buffer: &BufferSnapshot, start_row: u32, end_row: u32) -> Range<Point> {
341    let mut snippet_end_row = end_row;
342    while buffer.is_line_blank(snippet_end_row) && snippet_end_row > start_row {
343        snippet_end_row -= 1;
344    }
345    Point::new(start_row, 0)..Point::new(snippet_end_row, buffer.line_len(snippet_end_row))
346}
347
348// Returns the ranges of the snippets in the buffer and the next point for moving the cursor to
349fn jupytext_cells(
350    buffer: &BufferSnapshot,
351    range: Range<Point>,
352) -> (Vec<Range<Point>>, Option<Point>) {
353    let mut current_row = range.start.row;
354
355    let Some(language) = buffer.language() else {
356        return (Vec::new(), None);
357    };
358
359    let default_scope = language.default_scope();
360    let comment_prefixes = default_scope.line_comment_prefixes();
361    if comment_prefixes.is_empty() {
362        return (Vec::new(), None);
363    }
364
365    let jupytext_prefixes = comment_prefixes
366        .iter()
367        .map(|comment_prefix| format!("{comment_prefix}%%"))
368        .collect::<Vec<_>>();
369
370    let mut snippet_start_row = None;
371    loop {
372        if jupytext_prefixes
373            .iter()
374            .any(|prefix| buffer.contains_str_at(Point::new(current_row, 0), prefix))
375        {
376            snippet_start_row = Some(current_row);
377            break;
378        } else if current_row > 0 {
379            current_row -= 1;
380        } else {
381            break;
382        }
383    }
384
385    let mut snippets = Vec::new();
386    if let Some(mut snippet_start_row) = snippet_start_row {
387        for current_row in range.start.row + 1..=buffer.max_point().row {
388            if jupytext_prefixes
389                .iter()
390                .any(|prefix| buffer.contains_str_at(Point::new(current_row, 0), prefix))
391            {
392                snippets.push(cell_range(buffer, snippet_start_row, current_row - 1));
393
394                if current_row <= range.end.row {
395                    snippet_start_row = current_row;
396                } else {
397                    // Return our snippets as well as the next point for moving the cursor to
398                    return (snippets, Some(Point::new(current_row, 0)));
399                }
400            }
401        }
402
403        // Go to the end of the buffer (no more jupytext cells found)
404        snippets.push(cell_range(
405            buffer,
406            snippet_start_row,
407            buffer.max_point().row,
408        ));
409    }
410
411    (snippets, None)
412}
413
414fn runnable_ranges(
415    buffer: &BufferSnapshot,
416    range: Range<Point>,
417) -> (Vec<Range<Point>>, Option<Point>) {
418    if let Some(language) = buffer.language() {
419        if language.name() == "Markdown".into() {
420            return (markdown_code_blocks(buffer, range.clone()), None);
421        }
422    }
423
424    let (jupytext_snippets, next_cursor) = jupytext_cells(buffer, range.clone());
425    if !jupytext_snippets.is_empty() {
426        return (jupytext_snippets, next_cursor);
427    }
428
429    let snippet_range = cell_range(buffer, range.start.row, range.end.row);
430    let start_language = buffer.language_at(snippet_range.start);
431    let end_language = buffer.language_at(snippet_range.end);
432
433    if start_language
434        .zip(end_language)
435        .map_or(false, |(start, end)| start == end)
436    {
437        (vec![snippet_range], None)
438    } else {
439        (Vec::new(), None)
440    }
441}
442
443// We allow markdown code blocks to end in a trailing newline in order to render the output
444// below the final code fence. This is different than our behavior for selections and Jupytext cells.
445fn markdown_code_blocks(buffer: &BufferSnapshot, range: Range<Point>) -> Vec<Range<Point>> {
446    buffer
447        .injections_intersecting_range(range)
448        .filter(|(_, language)| language_supported(language))
449        .map(|(content_range, _)| {
450            buffer.offset_to_point(content_range.start)..buffer.offset_to_point(content_range.end)
451        })
452        .collect()
453}
454
455fn language_supported(language: &Arc<Language>) -> bool {
456    match language.name().as_ref() {
457        "TypeScript" | "Python" => true,
458        _ => false,
459    }
460}
461
462fn get_language(editor: WeakEntity<Editor>, cx: &mut App) -> Option<Arc<Language>> {
463    editor
464        .update(cx, |editor, cx| {
465            let selection = editor.selections.newest::<usize>(cx);
466            let buffer = editor.buffer().read(cx).snapshot(cx);
467            buffer.language_at(selection.head()).cloned()
468        })
469        .ok()
470        .flatten()
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476    use gpui::App;
477    use indoc::indoc;
478    use language::{Buffer, Language, LanguageConfig, LanguageRegistry};
479
480    #[gpui::test]
481    fn test_snippet_ranges(cx: &mut App) {
482        // Create a test language
483        let test_language = Arc::new(Language::new(
484            LanguageConfig {
485                name: "TestLang".into(),
486                line_comments: vec!["# ".into()],
487                ..Default::default()
488            },
489            None,
490        ));
491
492        let buffer = cx.new(|cx| {
493            Buffer::local(
494                indoc! { r#"
495                    print(1 + 1)
496                    print(2 + 2)
497
498                    print(4 + 4)
499
500
501                "# },
502                cx,
503            )
504            .with_language(test_language, cx)
505        });
506        let snapshot = buffer.read(cx).snapshot();
507
508        // Single-point selection
509        let (snippets, _) = runnable_ranges(&snapshot, Point::new(0, 4)..Point::new(0, 4));
510        let snippets = snippets
511            .into_iter()
512            .map(|range| snapshot.text_for_range(range).collect::<String>())
513            .collect::<Vec<_>>();
514        assert_eq!(snippets, vec!["print(1 + 1)"]);
515
516        // Multi-line selection
517        let (snippets, _) = runnable_ranges(&snapshot, Point::new(0, 5)..Point::new(2, 0));
518        let snippets = snippets
519            .into_iter()
520            .map(|range| snapshot.text_for_range(range).collect::<String>())
521            .collect::<Vec<_>>();
522        assert_eq!(
523            snippets,
524            vec![indoc! { r#"
525                print(1 + 1)
526                print(2 + 2)"# }]
527        );
528
529        // Trimming multiple trailing blank lines
530        let (snippets, _) = runnable_ranges(&snapshot, Point::new(0, 5)..Point::new(5, 0));
531
532        let snippets = snippets
533            .into_iter()
534            .map(|range| snapshot.text_for_range(range).collect::<String>())
535            .collect::<Vec<_>>();
536        assert_eq!(
537            snippets,
538            vec![indoc! { r#"
539                print(1 + 1)
540                print(2 + 2)
541
542                print(4 + 4)"# }]
543        );
544    }
545
546    #[gpui::test]
547    fn test_jupytext_snippet_ranges(cx: &mut App) {
548        // Create a test language
549        let test_language = Arc::new(Language::new(
550            LanguageConfig {
551                name: "TestLang".into(),
552                line_comments: vec!["# ".into()],
553                ..Default::default()
554            },
555            None,
556        ));
557
558        let buffer = cx.new(|cx| {
559            Buffer::local(
560                indoc! { r#"
561                    # Hello!
562                    # %% [markdown]
563                    # This is some arithmetic
564                    print(1 + 1)
565                    print(2 + 2)
566
567                    # %%
568                    print(3 + 3)
569                    print(4 + 4)
570
571                    print(5 + 5)
572
573
574
575                "# },
576                cx,
577            )
578            .with_language(test_language, cx)
579        });
580        let snapshot = buffer.read(cx).snapshot();
581
582        // Jupytext snippet surrounding an empty selection
583        let (snippets, _) = runnable_ranges(&snapshot, Point::new(2, 5)..Point::new(2, 5));
584
585        let snippets = snippets
586            .into_iter()
587            .map(|range| snapshot.text_for_range(range).collect::<String>())
588            .collect::<Vec<_>>();
589        assert_eq!(
590            snippets,
591            vec![indoc! { r#"
592                # %% [markdown]
593                # This is some arithmetic
594                print(1 + 1)
595                print(2 + 2)"# }]
596        );
597
598        // Jupytext snippets intersecting a non-empty selection
599        let (snippets, _) = runnable_ranges(&snapshot, Point::new(2, 5)..Point::new(6, 2));
600        let snippets = snippets
601            .into_iter()
602            .map(|range| snapshot.text_for_range(range).collect::<String>())
603            .collect::<Vec<_>>();
604        assert_eq!(
605            snippets,
606            vec![
607                indoc! { r#"
608                    # %% [markdown]
609                    # This is some arithmetic
610                    print(1 + 1)
611                    print(2 + 2)"#
612                },
613                indoc! { r#"
614                    # %%
615                    print(3 + 3)
616                    print(4 + 4)
617
618                    print(5 + 5)"#
619                }
620            ]
621        );
622    }
623
624    #[gpui::test]
625    fn test_markdown_code_blocks(cx: &mut App) {
626        let markdown = languages::language("markdown", tree_sitter_md::LANGUAGE.into());
627        let typescript = languages::language(
628            "typescript",
629            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
630        );
631        let python = languages::language("python", tree_sitter_python::LANGUAGE.into());
632        let language_registry = Arc::new(LanguageRegistry::new(cx.background_executor().clone()));
633        language_registry.add(markdown.clone());
634        language_registry.add(typescript.clone());
635        language_registry.add(python.clone());
636
637        // Two code blocks intersecting with selection
638        let buffer = cx.new(|cx| {
639            let mut buffer = Buffer::local(
640                indoc! { r#"
641                    Hey this is Markdown!
642
643                    ```typescript
644                    let foo = 999;
645                    console.log(foo + 1999);
646                    ```
647
648                    ```typescript
649                    console.log("foo")
650                    ```
651                    "#
652                },
653                cx,
654            );
655            buffer.set_language_registry(language_registry.clone());
656            buffer.set_language(Some(markdown.clone()), cx);
657            buffer
658        });
659        let snapshot = buffer.read(cx).snapshot();
660
661        let (snippets, _) = runnable_ranges(&snapshot, Point::new(3, 5)..Point::new(8, 5));
662        let snippets = snippets
663            .into_iter()
664            .map(|range| snapshot.text_for_range(range).collect::<String>())
665            .collect::<Vec<_>>();
666
667        assert_eq!(
668            snippets,
669            vec![
670                indoc! { r#"
671                    let foo = 999;
672                    console.log(foo + 1999);
673                    "#
674                },
675                "console.log(\"foo\")\n"
676            ]
677        );
678
679        // Three code blocks intersecting with selection
680        let buffer = cx.new(|cx| {
681            let mut buffer = Buffer::local(
682                indoc! { r#"
683                    Hey this is Markdown!
684
685                    ```typescript
686                    let foo = 999;
687                    console.log(foo + 1999);
688                    ```
689
690                    ```ts
691                    console.log("foo")
692                    ```
693
694                    ```typescript
695                    console.log("another code block")
696                    ```
697                "# },
698                cx,
699            );
700            buffer.set_language_registry(language_registry.clone());
701            buffer.set_language(Some(markdown.clone()), cx);
702            buffer
703        });
704        let snapshot = buffer.read(cx).snapshot();
705
706        let (snippets, _) = runnable_ranges(&snapshot, Point::new(3, 5)..Point::new(12, 5));
707        let snippets = snippets
708            .into_iter()
709            .map(|range| snapshot.text_for_range(range).collect::<String>())
710            .collect::<Vec<_>>();
711
712        assert_eq!(
713            snippets,
714            vec![
715                indoc! { r#"
716                    let foo = 999;
717                    console.log(foo + 1999);
718                    "#
719                },
720                "console.log(\"foo\")\n",
721                "console.log(\"another code block\")\n",
722            ]
723        );
724
725        // Python code block
726        let buffer = cx.new(|cx| {
727            let mut buffer = Buffer::local(
728                indoc! { r#"
729                    Hey this is Markdown!
730
731                    ```python
732                    print("hello there")
733                    print("hello there")
734                    print("hello there")
735                    ```
736                "# },
737                cx,
738            );
739            buffer.set_language_registry(language_registry.clone());
740            buffer.set_language(Some(markdown.clone()), cx);
741            buffer
742        });
743        let snapshot = buffer.read(cx).snapshot();
744
745        let (snippets, _) = runnable_ranges(&snapshot, Point::new(4, 5)..Point::new(5, 5));
746        let snippets = snippets
747            .into_iter()
748            .map(|range| snapshot.text_for_range(range).collect::<String>())
749            .collect::<Vec<_>>();
750
751        assert_eq!(
752            snippets,
753            vec![indoc! { r#"
754                print("hello there")
755                print("hello there")
756                print("hello there")
757                "#
758            },]
759        );
760    }
761}