Expand the worktree functionality

This commit is contained in:
2021-11-22 21:19:12 +01:00
parent 8ba214d6cf
commit 711d9131da
3 changed files with 511 additions and 161 deletions

View File

@@ -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<String>,
#[clap(short = 't', long = "track", about = "Remote branch to track")]
pub track: Option<String>,
}
#[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()

View File

@@ -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<git2::Branch, String> {
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::<String>().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::<String>().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::<String>()
.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::<String>().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::<Vec<String>>();
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::<String>();
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
));
}
}
}
}
}

View File

@@ -90,11 +90,14 @@ pub struct RepoStatus {
pub head: Option<String>,
pub changes: Option<RepoChanges>,
// None(_) => Could not get changes (e.g. because it's a worktree setup
// Some(None) => No changes
// Some(Some(_)) => Changes
pub changes: Option<Option<RepoChanges>>,
pub worktrees: usize,
pub submodules: Vec<(String, SubmoduleStatus)>,
pub submodules: Option<Vec<(String, SubmoduleStatus)>>,
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::<Vec<String>>();
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