diff --git a/src/cmd.rs b/src/cmd.rs index a79ec91..75a7140 100644 --- a/src/cmd.rs +++ b/src/cmd.rs @@ -66,16 +66,46 @@ pub struct Worktree { #[derive(Parser)] pub enum WorktreeAction { #[clap(about = "Add a new worktree")] - Add(WorktreeActionArgs), + Add(WorktreeAddArgs), #[clap(about = "Add an existing worktree")] - Delete(WorktreeActionArgs), + Delete(WorktreeDeleteArgs), + #[clap(about = "Show state of existing worktrees")] + Status(WorktreeStatusArgs), + #[clap(about = "Clean all worktrees that do not contain uncommited/unpushed changes")] + Clean(WorktreeCleanArgs), } #[derive(Parser)] -pub struct WorktreeActionArgs { +pub struct WorktreeAddArgs { #[clap(about = "Name of the worktree")] pub name: String, + + #[clap( + short = 'p', + long = "branch-prefix", + about = "Prefix for the branch name" + )] + pub branch_prefix: Option, + #[clap(short = 't', long = "track", about = "Remote branch to track")] + pub track: Option, } +#[derive(Parser)] +pub struct WorktreeDeleteArgs { + #[clap(about = "Name of the worktree")] + pub name: String, + + #[clap( + long = "force", + about = "Force deletion, even when there are uncommitted/unpushed changes" + )] + pub force: bool, +} + +#[derive(Parser)] +pub struct WorktreeStatusArgs {} + +#[derive(Parser)] +pub struct WorktreeCleanArgs {} pub fn parse() -> Opts { Opts::parse() diff --git a/src/lib.rs b/src/lib.rs index d9b7890..393b527 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -99,6 +99,16 @@ fn expand_path(path: &Path) -> PathBuf { Path::new(&expanded_path).to_path_buf() } +fn get_default_branch(repo: &git2::Repository) -> Result { + match repo.find_branch("main", git2::BranchType::Local) { + Ok(branch) => Ok(branch), + Err(_) => match repo.find_branch("master", git2::BranchType::Local) { + Ok(branch) => Ok(branch), + Err(_) => Err(String::from("Could not determine default branch")), + }, + } +} + fn sync_trees(config: Config) { for tree in config.trees { let repos = tree.repos.unwrap_or_default(); @@ -448,6 +458,7 @@ fn add_table_header(table: &mut Table) { .apply_modifier(comfy_table::modifiers::UTF8_ROUND_CORNERS) .set_header(vec![ Cell::new("Repo"), + Cell::new("Worktree"), Cell::new("Status"), Cell::new("Branches"), Cell::new("HEAD"), @@ -455,26 +466,50 @@ fn add_table_header(table: &mut Table) { ]); } -fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repository) { - let repo_status = get_repo_status(repo_handle); +fn add_worktree_table_header(table: &mut Table) { + table + .load_preset(comfy_table::presets::UTF8_FULL) + .apply_modifier(comfy_table::modifiers::UTF8_ROUND_CORNERS) + .set_header(vec![ + Cell::new("Worktree"), + Cell::new("Status"), + Cell::new("Branch"), + Cell::new("Remote branch"), + ]); +} + +fn add_repo_status( + table: &mut Table, + repo_name: &str, + repo_handle: &git2::Repository, + is_worktree: bool, +) { + let repo_status = get_repo_status(repo_handle, is_worktree); table.add_row(vec![ repo_name, + match is_worktree { + true => "\u{2714}", + false => "", + }, &match repo_status.changes { - Some(changes) => { - let mut out = Vec::new(); - if changes.files_new > 0 { - out.push(format!("New: {}\n", changes.files_new)) + None => String::from("-"), + Some(changes) => match changes { + Some(changes) => { + let mut out = Vec::new(); + if changes.files_new > 0 { + out.push(format!("New: {}\n", changes.files_new)) + } + if changes.files_modified > 0 { + out.push(format!("Modified: {}\n", changes.files_modified)) + } + if changes.files_deleted > 0 { + out.push(format!("Deleted: {}\n", changes.files_deleted)) + } + out.into_iter().collect::().trim().to_string() } - if changes.files_modified > 0 { - out.push(format!("Modified: {}\n", changes.files_modified)) - } - if changes.files_deleted > 0 { - out.push(format!("Deleted: {}\n", changes.files_deleted)) - } - out.into_iter().collect::().trim().to_string() - } - None => String::from("\u{2714}"), + None => String::from("\u{2714}"), + }, }, &repo_status .branches @@ -504,9 +539,12 @@ fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repos .collect::() .trim() .to_string(), - &match repo_status.head { - Some(head) => head, - None => String::from("Empty"), + &match is_worktree { + true => String::from(""), + false => match repo_status.head { + Some(head) => head, + None => String::from("Empty"), + }, }, &repo_status .remotes @@ -518,6 +556,72 @@ fn add_repo_status(table: &mut Table, repo_name: &str, repo_handle: &git2::Repos ]); } +fn add_worktree_status(table: &mut Table, worktree_name: &str, repo: &git2::Repository) { + let repo_status = get_repo_status(repo, false); + + let head = repo.head().unwrap(); + + if !head.is_branch() { + print_error("No branch checked out in worktree"); + process::exit(1); + } + + let local_branch_name = head.shorthand().unwrap(); + let local_branch = repo + .find_branch(local_branch_name, git2::BranchType::Local) + .unwrap(); + + let upstream_output = match local_branch.upstream() { + Ok(remote_branch) => { + let remote_branch_name = remote_branch.name().unwrap().unwrap().to_string(); + + let (ahead, behind) = repo + .graph_ahead_behind( + local_branch.get().peel_to_commit().unwrap().id(), + remote_branch.get().peel_to_commit().unwrap().id(), + ) + .unwrap(); + + format!( + "{}{}\n", + &remote_branch_name, + &match (ahead, behind) { + (0, 0) => String::from(""), + (d, 0) => format!(" [+{}]", &d), + (0, d) => format!(" [-{}]", &d), + (d1, d2) => format!(" [+{}/-{}]", &d1, &d2), + }, + ) + } + Err(_) => String::from(""), + }; + + table.add_row(vec![ + worktree_name, + &match repo_status.changes { + None => String::from(""), + Some(changes) => match changes { + Some(changes) => { + let mut out = Vec::new(); + if changes.files_new > 0 { + out.push(format!("New: {}\n", changes.files_new)) + } + if changes.files_modified > 0 { + out.push(format!("Modified: {}\n", changes.files_modified)) + } + if changes.files_deleted > 0 { + out.push(format!("Deleted: {}\n", changes.files_deleted)) + } + out.into_iter().collect::().trim().to_string() + } + None => String::from("\u{2714}"), + }, + }, + local_branch_name, + &upstream_output, + ]); +} + fn show_single_repo_status(path: &Path, is_worktree: bool) { let mut table = Table::new(); add_table_header(&mut table); @@ -547,7 +651,7 @@ fn show_single_repo_status(path: &Path, is_worktree: bool) { }, }; - add_repo_status(&mut table, &repo_name, &repo_handle.unwrap()); + add_repo_status(&mut table, &repo_name, &repo_handle.unwrap(), is_worktree); println!("{}", table); } @@ -588,12 +692,104 @@ fn show_status(config: Config) { let repo_handle = repo_handle.unwrap(); - add_repo_status(&mut table, &repo.name, &repo_handle); + add_repo_status(&mut table, &repo.name, &repo_handle, repo.worktree_setup); } println!("{}", table); } } +enum WorktreeRemoveFailureReason { + Changes(String), + Error(String), +} + +fn remove_worktree( + name: &str, + worktree_dir: &Path, + force: bool, + main_repo: &git2::Repository, +) -> Result<(), WorktreeRemoveFailureReason> { + if !worktree_dir.exists() { + return Err(WorktreeRemoveFailureReason::Error(format!( + "{} does not exist", + name + ))); + } + let worktree_repo = match open_repo(worktree_dir, false) { + Ok(r) => r, + Err(e) => { + return Err(WorktreeRemoveFailureReason::Error(format!( + "Error opening repo: {}", + e + ))); + } + }; + + let head = worktree_repo.head().unwrap(); + if !head.is_branch() { + return Err(WorktreeRemoveFailureReason::Error(String::from( + "No branch checked out in worktree", + ))); + } + + let branch_name = head.shorthand().unwrap(); + if !branch_name.ends_with(name) { + return Err(WorktreeRemoveFailureReason::Error(format!( + "Branch {} is checked out in worktree, this does not look correct", + &branch_name + ))); + } + + let mut branch = worktree_repo + .find_branch(branch_name, git2::BranchType::Local) + .unwrap(); + + if !force { + let status = get_repo_status(&worktree_repo, false); + if status.changes.is_some() { + return Err(WorktreeRemoveFailureReason::Changes(String::from( + "Changes found in worktree", + ))); + } + + match branch.upstream() { + Ok(remote_branch) => { + let (ahead, behind) = worktree_repo + .graph_ahead_behind( + branch.get().peel_to_commit().unwrap().id(), + remote_branch.get().peel_to_commit().unwrap().id(), + ) + .unwrap(); + + if (ahead, behind) != (0, 0) { + return Err(WorktreeRemoveFailureReason::Changes(format!( + "Branch {} is not in line with remote branch", + name + ))); + } + } + Err(_) => { + return Err(WorktreeRemoveFailureReason::Changes(format!( + "No remote tracking branch for branch {} found", + name + ))); + } + } + } + + if let Err(e) = std::fs::remove_dir_all(&worktree_dir) { + return Err(WorktreeRemoveFailureReason::Error(format!( + "Error deleting {}: {}", + &worktree_dir.display(), + e + ))); + } + main_repo.find_worktree(name).unwrap().prune(None).unwrap(); + branch.delete().unwrap(); + + Ok(()) +} + pub fn run() { let opts = cmd::parse(); @@ -667,33 +863,80 @@ pub fn run() { } }; + let repo = match open_repo(&dir, true) { + Ok(r) => r, + Err(e) => { + match e.kind { + RepoErrorKind::NotFound => { + print_error("Current directory does not contain a worktree setup") + } + _ => print_error(&format!("Error opening repo: {}", e)), + } + process::exit(1); + } + }; + + let worktrees = repo + .worktrees() + .unwrap() + .iter() + .map(|e| e.unwrap().to_string()) + .collect::>(); + match args.action { cmd::WorktreeAction::Add(action_args) => { - let repo = match open_repo(&dir, true) { - Ok(r) => r, - Err(e) => { - match e.kind { - RepoErrorKind::NotFound => print_error( - "Current directory does not contain a worktree setup", - ), - _ => print_error(&format!("Error opening repo: {}", e)), - } - process::exit(1); - } - }; - - let worktrees = repo - .worktrees() - .unwrap() - .iter() - .map(|e| e.unwrap()) - .collect::(); if worktrees.contains(&action_args.name) { - print_error("Worktree directory already exists"); + print_error("Worktree already exists"); process::exit(1); } - match repo.worktree(&action_args.name, &dir.join(&action_args.name), None) { + let branch_name = match action_args.branch_prefix { + Some(prefix) => format!("{}{}", &prefix, &action_args.name), + None => action_args.name.clone(), + }; + + let checkout_commit = match &action_args.track { + Some(upstream_branch_name) => { + match repo.find_branch(upstream_branch_name, git2::BranchType::Remote) { + Ok(branch) => branch.into_reference().peel_to_commit().unwrap(), + Err(_) => { + print_error(&format!( + "Remote branch {} not found", + &upstream_branch_name + )); + process::exit(1); + } + } + } + None => get_default_branch(&repo) + .unwrap() + .into_reference() + .peel_to_commit() + .unwrap(), + }; + + let mut target_branch = + match repo.find_branch(&branch_name, git2::BranchType::Local) { + Ok(branchref) => branchref, + Err(_) => repo.branch(&branch_name, &checkout_commit, false).unwrap(), + }; + + if let Some(upstream_branch_name) = action_args.track { + target_branch + .set_upstream(Some(&upstream_branch_name)) + .unwrap(); + } + + let worktree = repo.worktree( + &action_args.name, + &dir.join(&action_args.name), + Some( + git2::WorktreeAddOptions::new() + .reference(Some(&target_branch.into_reference())), + ), + ); + + match worktree { Ok(_) => print_success(&format!("Worktree {} created", &action_args.name)), Err(e) => { print_error(&format!("Error creating worktree: {}", e)); @@ -701,64 +944,122 @@ pub fn run() { } }; } + cmd::WorktreeAction::Delete(action_args) => { let worktree_dir = dir.join(&action_args.name); - if !worktree_dir.exists() { - print_error(&format!("{} does not exist", &action_args.name)); - process::exit(1); - } - let repo = match open_repo(&worktree_dir, false) { - Ok(r) => r, - Err(e) => { - print_error(&format!("Error opening repo: {}", e)); - process::exit(1); - } - }; - let status = get_repo_status(&repo); - if status.changes.is_some() { - print_error("Changes found in worktree, refusing to delete!"); - process::exit(1); - } - let mut branch = repo - .find_branch(&action_args.name, git2::BranchType::Local) - .unwrap(); - match branch.upstream() { - Ok(remote_branch) => { - let (ahead, behind) = repo - .graph_ahead_behind( - branch.get().peel_to_commit().unwrap().id(), - remote_branch.get().peel_to_commit().unwrap().id(), - ) - .unwrap(); - - if (ahead, behind) != (0, 0) { - print_error(&format!("Branch {} is not in line with remote branch, refusing to delete worktree!", &action_args.name)); - process::exit(1); - } - } - Err(_) => { - print_error(&format!("No remote tracking branch for branch {} found, refusing to delete worktree!", &action_args.name)); - process::exit(1); - } - } - - match std::fs::remove_dir_all(&worktree_dir) { + match remove_worktree( + &action_args.name, + &worktree_dir, + action_args.force, + &repo, + ) { Ok(_) => print_success(&format!("Worktree {} deleted", &action_args.name)), - Err(e) => { - print_error(&format!( - "Error deleting {}: {}", - &worktree_dir.display(), - e - )); + Err(error) => { + match error { + WorktreeRemoveFailureReason::Error(msg) => { + print_error(&msg); + process::exit(1); + } + WorktreeRemoveFailureReason::Changes(changes) => { + print_warning(&format!( + "Changes in worktree: {}. Refusing to delete", + changes + )); + } + } process::exit(1); } } - repo.find_worktree(&action_args.name) - .unwrap() - .prune(None) - .unwrap(); - branch.delete().unwrap(); + } + cmd::WorktreeAction::Status(_args) => { + let mut table = Table::new(); + add_worktree_table_header(&mut table); + for worktree in &worktrees { + let repo_dir = &dir.join(&worktree); + if repo_dir.exists() { + let repo = match open_repo(repo_dir, false) { + Ok(r) => r, + Err(e) => { + print_error(&format!("Error opening repo: {}", e)); + process::exit(1); + } + }; + add_worktree_status(&mut table, worktree, &repo); + } else { + print_warning(&format!( + "Worktree {} does not have a directory", + &worktree + )); + } + } + for entry in std::fs::read_dir(&dir).unwrap() { + let dirname = path_as_string( + &entry + .unwrap() + .path() + .strip_prefix(&dir) + .unwrap() + .to_path_buf(), + ); + if dirname == GIT_MAIN_WORKTREE_DIRECTORY { + continue; + } + if !&worktrees.contains(&dirname) { + print_warning(&format!( + "Found {}, which is not a valid worktree directory!", + &dirname + )); + } + } + println!("{}", table); + } + cmd::WorktreeAction::Clean(_args) => { + for worktree in &worktrees { + let repo_dir = &dir.join(&worktree); + if repo_dir.exists() { + match remove_worktree(worktree, repo_dir, false, &repo) { + Ok(_) => print_success(&format!("Worktree {} deleted", &worktree)), + Err(error) => match error { + WorktreeRemoveFailureReason::Changes(changes) => { + print_warning(&format!( + "Changes found in {}: {}, skipping", + &worktree, &changes + )); + continue; + } + WorktreeRemoveFailureReason::Error(e) => { + print_error(&e); + process::exit(1); + } + }, + } + } else { + print_warning(&format!( + "Worktree {} does not have a directory", + &worktree + )); + } + } + for entry in std::fs::read_dir(&dir).unwrap() { + let dirname = path_as_string( + &entry + .unwrap() + .path() + .strip_prefix(&dir) + .unwrap() + .to_path_buf(), + ); + if dirname == GIT_MAIN_WORKTREE_DIRECTORY { + continue; + } + if !&worktrees.contains(&dirname) { + print_warning(&format!( + "Found {}, which is not a valid worktree directory!", + &dirname + )); + } + } } } } diff --git a/src/repo.rs b/src/repo.rs index 4933fe1..f2535c7 100644 --- a/src/repo.rs +++ b/src/repo.rs @@ -90,11 +90,14 @@ pub struct RepoStatus { pub head: Option, - pub changes: Option, + // None(_) => Could not get changes (e.g. because it's a worktree setup + // Some(None) => No changes + // Some(Some(_)) => Changes + pub changes: Option>, pub worktrees: usize, - pub submodules: Vec<(String, SubmoduleStatus)>, + pub submodules: Option>, pub branches: Vec<(String, Option<(String, RemoteTrackingStatus)>)>, } @@ -253,7 +256,7 @@ pub fn clone_repo( } } -pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus { +pub fn get_repo_status(repo: &git2::Repository, is_worktree: bool) -> RepoStatus { let operation = match repo.state() { git2::RepositoryState::Clean => None, state => Some(state), @@ -268,84 +271,100 @@ pub fn get_repo_status(repo: &git2::Repository) -> RepoStatus { .map(|repo_name| repo_name.unwrap().to_string()) .collect::>(); - let head = match empty { + let head = match is_worktree { true => None, - false => Some(repo.head().unwrap().shorthand().unwrap().to_string()), + false => match empty { + true => None, + false => Some(repo.head().unwrap().shorthand().unwrap().to_string()), + }, }; - let statuses = repo - .statuses(Some( - git2::StatusOptions::new() - .include_ignored(false) - .include_untracked(true), - )) - .unwrap(); - - let changes = match statuses.is_empty() { + let changes = match is_worktree { true => None, false => { - let mut files_new = 0; - let mut files_modified = 0; - let mut files_deleted = 0; - for status in statuses.iter() { - let status_bits = status.status(); - if status_bits.intersects( - git2::Status::INDEX_MODIFIED - | git2::Status::INDEX_RENAMED - | git2::Status::INDEX_TYPECHANGE - | git2::Status::WT_MODIFIED - | git2::Status::WT_RENAMED - | git2::Status::WT_TYPECHANGE, - ) { - files_modified += 1; - } else if status_bits.intersects(git2::Status::INDEX_NEW | git2::Status::WT_NEW) { - files_new += 1; - } else if status_bits - .intersects(git2::Status::INDEX_DELETED | git2::Status::WT_DELETED) - { - files_deleted += 1; + let statuses = repo + .statuses(Some( + git2::StatusOptions::new() + .include_ignored(false) + .include_untracked(true), + )) + .unwrap(); + + match statuses.is_empty() { + true => Some(None), + false => { + let mut files_new = 0; + let mut files_modified = 0; + let mut files_deleted = 0; + for status in statuses.iter() { + let status_bits = status.status(); + if status_bits.intersects( + git2::Status::INDEX_MODIFIED + | git2::Status::INDEX_RENAMED + | git2::Status::INDEX_TYPECHANGE + | git2::Status::WT_MODIFIED + | git2::Status::WT_RENAMED + | git2::Status::WT_TYPECHANGE, + ) { + files_modified += 1; + } else if status_bits + .intersects(git2::Status::INDEX_NEW | git2::Status::WT_NEW) + { + files_new += 1; + } else if status_bits + .intersects(git2::Status::INDEX_DELETED | git2::Status::WT_DELETED) + { + files_deleted += 1; + } + } + if (files_new, files_modified, files_deleted) == (0, 0, 0) { + panic!( + "is_empty() returned true, but no file changes were detected. This is a bug!" + ); + } + Some(Some(RepoChanges { + files_new, + files_modified, + files_deleted, + })) } } - if (files_new, files_modified, files_deleted) == (0, 0, 0) { - panic!( - "is_empty() returned true, but no file changes were detected. This is a bug!" - ); - } - Some(RepoChanges { - files_new, - files_modified, - files_deleted, - }) } }; let worktrees = repo.worktrees().unwrap().len(); - let mut submodules = Vec::new(); - for submodule in repo.submodules().unwrap() { - let submodule_name = submodule.name().unwrap().to_string(); + let submodules = match is_worktree { + true => None, + false => { + let mut submodules = Vec::new(); + for submodule in repo.submodules().unwrap() { + let submodule_name = submodule.name().unwrap().to_string(); - let submodule_status; - let status = repo - .submodule_status(submodule.name().unwrap(), git2::SubmoduleIgnore::None) - .unwrap(); + let submodule_status; + let status = repo + .submodule_status(submodule.name().unwrap(), git2::SubmoduleIgnore::None) + .unwrap(); - if status.intersects( - git2::SubmoduleStatus::WD_INDEX_MODIFIED - | git2::SubmoduleStatus::WD_WD_MODIFIED - | git2::SubmoduleStatus::WD_UNTRACKED, - ) { - submodule_status = SubmoduleStatus::Changed; - } else if status.is_wd_uninitialized() { - submodule_status = SubmoduleStatus::Uninitialized; - } else if status.is_wd_modified() { - submodule_status = SubmoduleStatus::OutOfDate; - } else { - submodule_status = SubmoduleStatus::Clean; + if status.intersects( + git2::SubmoduleStatus::WD_INDEX_MODIFIED + | git2::SubmoduleStatus::WD_WD_MODIFIED + | git2::SubmoduleStatus::WD_UNTRACKED, + ) { + submodule_status = SubmoduleStatus::Changed; + } else if status.is_wd_uninitialized() { + submodule_status = SubmoduleStatus::Uninitialized; + } else if status.is_wd_modified() { + submodule_status = SubmoduleStatus::OutOfDate; + } else { + submodule_status = SubmoduleStatus::Clean; + } + + submodules.push((submodule_name, submodule_status)); + } + Some(submodules) } - - submodules.push((submodule_name, submodule_status)); - } + }; let mut branches = Vec::new(); for (local_branch, _) in repo