add remove local branch option when removing a worktree

Author: ozankasikci

src/commands/interactive/command.rs

@@ -15,7 +15,10 @@ use ratatui::{
 
 use super::{
     Action, EventSource, Focus, Selection, StatusMessage, WorktreeEntry,
-    dialog::{CreateDialog, CreateDialogFocus, Dialog, MergeDialog, MergeDialogFocus},
+    dialog::{
+        CreateDialog, CreateDialogFocus, Dialog, MergeDialog, MergeDialogFocus, RemoveDialog,
+        RemoveDialogFocus,
+    },
     view::{DetailData, DialogView, Snapshot},
 };
 
@@ -74,7 +77,7 @@ where
 
     pub fn run<F, G>(mut self, mut on_remove: F, mut on_create: G) -> Result<Option<Selection>>
     where
-        F: FnMut(&str) -> Result<()>,
+        F: FnMut(&str, bool) -> Result<()>,
         G: FnMut(&str, Option<&str>) -> Result<()>,
     {
         self.terminal
@@ -96,7 +99,7 @@ where
         on_create: &mut G,
     ) -> Result<Option<Selection>>
     where
-        F: FnMut(&str) -> Result<()>,
+        F: FnMut(&str, bool) -> Result<()>,
         G: FnMut(&str, Option<&str>) -> Result<()>,
     {
         let mut state = ListState::default();
@@ -123,15 +126,15 @@ where
         on_create: &mut G,
     ) -> Result<LoopControl>
     where
-        F: FnMut(&str) -> Result<()>,
+        F: FnMut(&str, bool) -> Result<()>,
         G: FnMut(&str, Option<&str>) -> Result<()>,
     {
         if let Some(dialog) = self.dialog.clone() {
             match dialog {
-                Dialog::ConfirmRemove { index } => {
+                Dialog::Remove(_) => {
                     if let Event::Key(key) = event {
                         if key.kind == KeyEventKind::Press {
-                            self.handle_confirm(index, key.code, state, on_remove)?;
+                            self.handle_remove_dialog_key(key, state, on_remove)?;
                         }
                     }
                     return Ok(LoopControl::Continue);
@@ -280,7 +283,7 @@ where
                     }
                     Action::Remove => {
                         if let Some(index) = self.selected {
-                            self.dialog = Some(Dialog::ConfirmRemove { index });
+                            self.dialog = Some(Dialog::Remove(RemoveDialog::new(index)));
                         } else {
                             self.status =
                                 Some(StatusMessage::info("No worktree selected to remove."));
@@ -319,57 +322,123 @@ where
         Ok(LoopControl::Continue)
     }
 
-    fn handle_confirm<F>(
+    fn handle_remove_dialog_key<F>(
         &mut self,
-        index: usize,
-        code: KeyCode,
+        key: KeyEvent,
         state: &mut ListState,
         on_remove: &mut F,
     ) -> Result<()>
     where
-        F: FnMut(&str) -> Result<()>,
+        F: FnMut(&str, bool) -> Result<()>,
     {
-        match code {
-            KeyCode::Char('y') | KeyCode::Char('Y') | KeyCode::Enter => {
-                if let Some(entry) = self.worktrees.get(index).cloned() {
-                    match on_remove(&entry.name) {
-                        Ok(()) => {
-                            self.worktrees.remove(index);
-                            let removal_dir = entry
-                                .path
-                                .parent()
-                                .map(|parent| parent.display().to_string())
-                                .unwrap_or_else(|| entry.path.display().to_string());
-                            let message = format!(
-                                "Removed worktree `{}` from `{}`.",
-                                entry.name, removal_dir
-                            );
-                            self.selected = None;
-                            self.focus = Focus::Worktrees;
-                            self.sync_selection(state);
-                            self.status = None;
-                            self.dialog = Some(Dialog::Info { message });
-                            return Ok(());
-                        }
-                        Err(err) => {
-                            self.status = Some(StatusMessage::error(format!(
-                                "Failed to remove `{}`: {err}",
-                                entry.name
-                            )));
-                            self.dialog = None;
-                            return Ok(());
-                        }
-                    }
-                }
-                self.dialog = None;
-            }
+        let dialog_state = self.dialog.take();
+        let Some(Dialog::Remove(mut dialog)) = dialog_state else {
+            self.dialog = dialog_state;
+            return Ok(());
+        };
+
+        let mut reinstate = true;
+
+        match key.code {
             KeyCode::Esc | KeyCode::Char('n') | KeyCode::Char('N') => {
+                reinstate = false;
                 self.status = Some(StatusMessage::info("Removal cancelled."));
-                self.dialog = None;
             }
+            KeyCode::Char('y') | KeyCode::Char('Y') => {
+                reinstate = false;
+                self.perform_remove(dialog.index, dialog.remove_local_branch(), state, on_remove)?;
+            }
+            KeyCode::Tab => dialog.focus_next(),
+            KeyCode::BackTab => dialog.focus_prev(),
+            KeyCode::Up | KeyCode::Char('k') => match dialog.focus {
+                RemoveDialogFocus::Options => dialog.move_option(-1),
+                RemoveDialogFocus::Buttons => dialog.focus = RemoveDialogFocus::Options,
+            },
+            KeyCode::Down | KeyCode::Char('j') => match dialog.focus {
+                RemoveDialogFocus::Options => dialog.move_option(1),
+                RemoveDialogFocus::Buttons => {}
+            },
+            KeyCode::Left => {
+                if dialog.focus == RemoveDialogFocus::Buttons {
+                    dialog.move_button(-1);
+                }
+            }
+            KeyCode::Right => {
+                if dialog.focus == RemoveDialogFocus::Buttons {
+                    dialog.move_button(1);
+                }
+            }
+            KeyCode::Char(' ') => {
+                if dialog.focus == RemoveDialogFocus::Options {
+                    dialog.toggle_selected_option();
+                }
+            }
+            KeyCode::Enter => match dialog.focus {
+                RemoveDialogFocus::Options => dialog.toggle_selected_option(),
+                RemoveDialogFocus::Buttons => {
+                    if dialog.buttons_selected == 0 {
+                        reinstate = false;
+                        self.status = Some(StatusMessage::info("Removal cancelled."));
+                    } else {
+                        reinstate = false;
+                        self.perform_remove(
+                            dialog.index,
+                            dialog.remove_local_branch(),
+                            state,
+                            on_remove,
+                        )?;
+                    }
+                }
+            },
             _ => {}
         }
 
+        if reinstate {
+            self.dialog = Some(Dialog::Remove(dialog));
+        }
+
+        Ok(())
+    }
+
+    fn perform_remove<F>(
+        &mut self,
+        index: usize,
+        remove_local_branch: bool,
+        state: &mut ListState,
+        on_remove: &mut F,
+    ) -> Result<()>
+    where
+        F: FnMut(&str, bool) -> Result<()>,
+    {
+        if let Some(entry) = self.worktrees.get(index).cloned() {
+            match on_remove(&entry.name, remove_local_branch) {
+                Ok(()) => {
+                    self.worktrees.remove(index);
+                    let removal_dir = entry
+                        .path
+                        .parent()
+                        .map(|parent| parent.display().to_string())
+                        .unwrap_or_else(|| entry.path.display().to_string());
+                    let message =
+                        format!("Removed worktree `{}` from `{}`.", entry.name, removal_dir);
+                    self.selected = None;
+                    self.focus = Focus::Worktrees;
+                    self.sync_selection(state);
+                    self.status = None;
+                    self.dialog = Some(Dialog::Info { message });
+                }
+                Err(err) => {
+                    self.status = Some(StatusMessage::error(format!(
+                        "Failed to remove `{}`: {err}",
+                        entry.name
+                    )));
+                    self.dialog = None;
+                }
+            }
+        } else {
+            self.dialog = None;
+        }
+
         Ok(())
     }
 
@@ -769,11 +838,12 @@ where
         let detail = self.current_entry().map(build_detail_data);
 
         let dialog = match self.dialog.clone() {
-            Some(Dialog::ConfirmRemove { index }) => {
+            Some(Dialog::Remove(dialog)) => {
                 self.worktrees
-                    .get(index)
-                    .map(|entry| DialogView::ConfirmRemove {
+                    .get(dialog.index)
+                    .map(|entry| DialogView::Remove {
                         name: entry.name.clone(),
+                        dialog: dialog.into(),
                     })
             }
             Some(Dialog::Info { message }) => Some(DialogView::Info { message }),

src/commands/interactive/dialog.rs

@@ -189,6 +189,96 @@ impl CreateDialogView {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum RemoveDialogFocus {
+    Options,
+    Buttons,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct RemoveDialog {
+    pub(crate) index: usize,
+    pub(crate) focus: RemoveDialogFocus,
+    pub(crate) options_selected: usize,
+    pub(crate) buttons_selected: usize,
+    pub(crate) remove_local_branch: bool,
+}
+
+impl RemoveDialog {
+    const OPTION_COUNT: usize = 1;
+    const BUTTON_COUNT: usize = 2;
+
+    pub(crate) fn new(index: usize) -> Self {
+        Self {
+            index,
+            focus: RemoveDialogFocus::Options,
+            options_selected: 0,
+            buttons_selected: 1,
+            remove_local_branch: true,
+        }
+    }
+
+    pub(crate) fn focus_next(&mut self) {
+        self.focus = match self.focus {
+            RemoveDialogFocus::Options => RemoveDialogFocus::Buttons,
+            RemoveDialogFocus::Buttons => RemoveDialogFocus::Options,
+        };
+    }
+
+    pub(crate) fn focus_prev(&mut self) {
+        self.focus_next();
+    }
+
+    pub(crate) fn move_option(&mut self, delta: isize) {
+        let len = Self::OPTION_COUNT as isize;
+        let current = self.options_selected as isize;
+        let next = (current + delta).rem_euclid(len);
+        self.options_selected = next as usize;
+    }
+
+    pub(crate) fn move_button(&mut self, delta: isize) {
+        let len = Self::BUTTON_COUNT as isize;
+        let current = self.buttons_selected as isize;
+        let next = (current + delta).rem_euclid(len);
+        self.buttons_selected = next as usize;
+    }
+
+    pub(crate) fn toggle_selected_option(&mut self) {
+        if self.options_selected == 0 {
+            self.remove_local_branch = !self.remove_local_branch;
+        }
+    }
+
+    pub(crate) fn remove_local_branch(&self) -> bool {
+        self.remove_local_branch
+    }
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct RemoveDialogView {
+    pub(crate) focus: RemoveDialogFocus,
+    pub(crate) options_selected: usize,
+    pub(crate) buttons_selected: usize,
+    pub(crate) remove_local_branch: bool,
+}
+
+impl From<&RemoveDialog> for RemoveDialogView {
+    fn from(dialog: &RemoveDialog) -> Self {
+        Self {
+            focus: dialog.focus,
+            options_selected: dialog.options_selected,
+            buttons_selected: dialog.buttons_selected,
+            remove_local_branch: dialog.remove_local_branch,
+        }
+    }
+}
+
+impl From<RemoveDialog> for RemoveDialogView {
+    fn from(dialog: RemoveDialog) -> Self {
+        Self::from(&dialog)
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub(crate) enum MergeDialogFocus {
     Options,
@@ -300,7 +390,7 @@ impl From<MergeDialog> for MergeDialogView {
 
 #[derive(Clone, Debug)]
 pub(crate) enum Dialog {
-    ConfirmRemove { index: usize },
+    Remove(RemoveDialog),
     Info { message: String },
     Create(CreateDialog),
     Merge(MergeDialog),

src/commands/interactive/runtime.rs

@@ -65,8 +65,10 @@ pub fn run(repo: &Repo) -> Result<()> {
         default_branch,
     );
     let result = command.run(
-        |name| {
-            let command = RemoveCommand::new(name.to_owned(), false).with_quiet(true);
+        |name, remove_local_branch| {
+            let command = RemoveCommand::new(name.to_owned(), false)
+                .with_quiet(true)
+                .with_remove_local_branch(remove_local_branch);
             command.execute(repo)
         },
         |name, base| {