From 098bcda6910aa2f734fb0ad59ac089157a239c87 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc=20Cornell=C3=A0?= <hello@mcornella.com>
Date: Thu, 9 Sep 2021 12:57:59 +0200
Subject: [PATCH] feat(update): allow updating from branch set up on install

Closes #8788

Co-authored-by: Nikolas Garofil <nikolas@garofil.be>
---
 tools/install.sh | 42 ++++++++++++++++++++++--------------------
 tools/upgrade.sh | 17 +++++++++++++++--
 2 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/tools/install.sh b/tools/install.sh
index cfc2808fe..510f83ce8 100755
--- a/tools/install.sh
+++ b/tools/install.sh
@@ -53,7 +53,7 @@ KEEP_ZSHRC=${KEEP_ZSHRC:-no}
 
 
 command_exists() {
-	command -v "$@" >/dev/null 2>&1
+  command -v "$@" >/dev/null 2>&1
 }
 
 fmt_error() {
@@ -65,27 +65,27 @@ fmt_underline() {
 }
 
 fmt_code() {
-  # shellcheck disable=SC2016 # backtic in single-quote
+  # shellcheck disable=SC2016 # backtick in single-quote
   printf '`\033[38;5;247m%s%s`\n' "$*" "$RESET"
 }
 
 setup_color() {
-	# Only use colors if connected to a terminal
-	if [ -t 1 ]; then
-		RED=$(printf '\033[31m')
-		GREEN=$(printf '\033[32m')
-		YELLOW=$(printf '\033[33m')
-		BLUE=$(printf '\033[34m')
-		BOLD=$(printf '\033[1m')
-		RESET=$(printf '\033[m')
-	else
-		RED=""
-		GREEN=""
-		YELLOW=""
-		BLUE=""
-		BOLD=""
-		RESET=""
-	fi
+  # Only use colors if connected to a terminal
+  if [ -t 1 ]; then
+    RED=$(printf '\033[31m')
+    GREEN=$(printf '\033[32m')
+    YELLOW=$(printf '\033[33m')
+    BLUE=$(printf '\033[34m')
+    BOLD=$(printf '\033[1m')
+    RESET=$(printf '\033[m')
+  else
+    RED=""
+    GREEN=""
+    YELLOW=""
+    BLUE=""
+    BOLD=""
+    RESET=""
+  fi
 }
 
 setup_ohmyzsh() {
@@ -114,6 +114,8 @@ setup_ohmyzsh() {
     -c fsck.zeroPaddedFilemode=ignore \
     -c fetch.fsck.zeroPaddedFilemode=ignore \
     -c receive.fsck.zeroPaddedFilemode=ignore \
+    -c oh-my-zsh.remote=origin \
+    -c oh-my-zsh.branch="$BRANCH" \
     --depth=1 --branch "$BRANCH" "$REMOTE" "$ZSH" || {
     fmt_error "git clone of oh-my-zsh repo failed"
     exit 1
@@ -157,9 +159,9 @@ setup_zshrc() {
   sed "/^export ZSH=/ c\\
 export ZSH=\"$ZSH\"
 " "$ZSH/templates/zshrc.zsh-template" > ~/.zshrc-omztemp
-	mv -f ~/.zshrc-omztemp ~/.zshrc
+  mv -f ~/.zshrc-omztemp ~/.zshrc
 
-	echo
+  echo
 }
 
 setup_shell() {
diff --git a/tools/upgrade.sh b/tools/upgrade.sh
index 7dec398fd..1424d1433 100755
--- a/tools/upgrade.sh
+++ b/tools/upgrade.sh
@@ -66,10 +66,20 @@ git config rebase.autoStash true
 
 local ret=0
 
+# repository settings
+remote=${"$(git config --local oh-my-zsh.remote)":-origin}
+branch=${"$(git config --local oh-my-zsh.branch)":-master}
+
+# repository state
+last_head=$(git symbolic-ref --quiet --short HEAD || git rev-parse HEAD)
+# checkout update branch
+git checkout -q "$branch" -- || exit 1
+# branch commit before update (used in changelog)
+last_commit=$(git rev-parse "$branch")
+
 # Update Oh My Zsh
 printf "${BLUE}%s${RESET}\n" "Updating Oh My Zsh"
-last_commit=$(git rev-parse HEAD)
-if git pull --rebase --stat origin master; then
+if git pull --rebase --stat $remote $branch; then
   # Check if it was really updated or not
   if [[ "$(git rev-parse HEAD)" = "$last_commit" ]]; then
     message="Oh My Zsh is already at the latest version."
@@ -103,6 +113,9 @@ else
   printf "${RED}%s${RESET}\n" 'There was an error updating. Try again later?'
 fi
 
+# go back to HEAD previous to update
+git checkout -q "$last_head" --
+
 # Unset git-config values set just for the upgrade
 case "$resetAutoStash" in
   "") git config --unset rebase.autoStash ;;