handle alternate default branches

3 files changed, 32 insertions, 1 deletions

diff --git a/src/prompt.rs b/src/prompt.rs
index 6c95c0e..f319a96 100644
--- a/src/prompt.rs
+++ b/src/prompt.rs
@@ -299,7 +299,7 @@ fn format_vcs(vcs_info: Option<&dyn vcs::VcsInfo>) -> Option<String> {
         let branch = vcs_info
             .branch()
             .map(|branch| {
-                if branch == "master" {
+                if Some(branch.clone()) == vcs_info.default_branch() {
                     String::new()
                 } else {
                     branch
@@ -456,6 +456,7 @@ mod test {
         active_operation: vcs::ActiveOperation,
         branch: Option<String>,
         remote_branch_diff: Option<(usize, usize)>,
+        default_branch: Option<String>,
     }
 
     impl vcs::VcsInfo for TestVcs {
@@ -483,6 +484,9 @@ mod test {
         fn remote_branch_diff(&self) -> Option<(usize, usize)> {
             self.remote_branch_diff
         }
+        fn default_branch(&self) -> Option<String> {
+            self.default_branch
+        }
     }
 
     #[test]
@@ -605,6 +609,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("master")),
                 remote_branch_diff: Some((0, 0)),
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(format_vcs(Some(&test_vcs)), Some(String::from("g")));
@@ -620,6 +625,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("dev")),
                 remote_branch_diff: Some((0, 0)),
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(
@@ -638,6 +644,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("master")),
                 remote_branch_diff: None,
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(
@@ -656,6 +663,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("dev")),
                 remote_branch_diff: None,
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(
@@ -674,6 +682,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("master")),
                 remote_branch_diff: None,
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(
@@ -692,6 +701,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("dev")),
                 remote_branch_diff: None,
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(
@@ -710,6 +720,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: None,
                 remote_branch_diff: None,
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(format_vcs(Some(&test_vcs)), Some(String::from("g!")));
@@ -725,6 +736,7 @@ mod test {
                 active_operation: vcs::ActiveOperation::None,
                 branch: Some(String::from("master")),
                 remote_branch_diff: Some((2, 3)),
+                default_branch: Some(String::from("master")),
             };
 
             assert_eq!(
diff --git a/src/vcs/git.rs b/src/vcs/git.rs
index f4666d1..8442857 100644
--- a/src/vcs/git.rs
+++ b/src/vcs/git.rs
@@ -9,6 +9,7 @@ pub struct GitInfo {
     active_operation: super::ActiveOperation,
     branch: Option<String>,
     remote_branch_diff: Option<(usize, usize)>,
+    default_branch: Option<String>,
 }
 
 impl GitInfo {
@@ -127,6 +128,18 @@ impl GitInfo {
                 })
             });
         talk_about_time!("remote branch diff");
+
+        let default_branch = git
+            .find_reference("refs/remotes/origin/HEAD")
+            .ok()
+            .and_then(|r| r.resolve().ok())
+            .and_then(|r| {
+                r.shorthand()
+                    .map(|s| s.strip_prefix("origin/").unwrap_or(s))
+                    .map(String::from)
+            });
+        eprintln!("{:?}", default_branch);
+        talk_about_time!("default branch");
         stop_talking_about_time!();
 
         GitInfo {
@@ -137,6 +150,7 @@ impl GitInfo {
             active_operation,
             branch,
             remote_branch_diff,
+            default_branch,
         }
     }
 }
@@ -173,6 +187,10 @@ impl super::VcsInfo for GitInfo {
     fn remote_branch_diff(&self) -> Option<(usize, usize)> {
         self.remote_branch_diff
     }
+
+    fn default_branch(&self) -> Option<String> {
+        self.default_branch.clone()
+    }
 }
 
 pub fn detect() -> Option<Box<dyn super::VcsInfo>> {
diff --git a/src/vcs/mod.rs b/src/vcs/mod.rs
index 96340cb..f4edb56 100644
--- a/src/vcs/mod.rs
+++ b/src/vcs/mod.rs
@@ -24,6 +24,7 @@ pub trait VcsInfo {
     fn active_operation(&self) -> ActiveOperation;
     fn branch(&self) -> Option<String>;
     fn remote_branch_diff(&self) -> Option<(usize, usize)>;
+    fn default_branch(&self) -> Option<String>;
 
     fn is_dirty(&self) -> bool {
         let diff = self.remote_branch_diff();