update: stash save/pop uncommitted changes.

Also:
- return to your previous branch after `brew update`.

Closes Homebrew/homebrew#38568.

Signed-off-by: Mike McQuaid <mike@mikemcquaid.com>

2 files changed, 42 insertions, 14 deletions

diff --git a/Library/Homebrew/cmd/update.rb b/Library/Homebrew/cmd/update.rb
index c01fb827a..37610956a 100644
--- a/Library/Homebrew/cmd/update.rb
+++ b/Library/Homebrew/cmd/update.rb
@@ -122,25 +122,53 @@ class Updater
 
   def initialize(repository)
     @repository = repository
+    @stashed = false
   end
 
-  def pull!
-    safe_system "git", "checkout", "-q", "master"
+  def pull!(options={})
+    quiet = []
+    quiet << "--quiet" unless ARGV.verbose?
+
+    unless system "git", "diff", "--quiet"
+      unless options[:silent]
+        puts "Stashing your changes:"
+        system "git", "status", "--short", "--untracked-files"
+      end
+      safe_system "git", "stash", "save", "--include-untracked", *quiet
+      @stashed = true
+    end
+
+    @initial_branch = `git symbolic-ref --short HEAD`.chomp
+    if @initial_branch != "master" && !@initial_branch.empty?
+      safe_system "git", "checkout", "master", *quiet
+    end
 
     @initial_revision = read_current_revision
 
     # ensure we don't munge line endings on checkout
     safe_system "git", "config", "core.autocrlf", "false"
 
-    args = ["pull"]
+    args = ["pull", "origin"]
     args << "--rebase" if ARGV.include? "--rebase"
-    args << "-q" unless ARGV.verbose?
-    args << "origin"
     # the refspec ensures that 'origin/master' gets updated
     args << "refs/heads/master:refs/remotes/origin/master"
+    args += quiet
 
     reset_on_interrupt { safe_system "git", *args }
 
+    if @initial_branch != "master" && !@initial_branch.empty?
+      safe_system "git", "checkout", @initial_branch, *quiet
+    end
+
+    if @stashed
+      safe_system "git", "stash", "pop", *quiet
+      unless options[:silent]
+        puts "Restored your changes:"
+        system "git", "status", "--short", "--untracked-files"
+      end
+      @stashed = false
+    end
+
     @current_revision = read_current_revision
   end
 
@@ -148,7 +176,9 @@ class Updater
     ignore_interrupts { yield }
   ensure
     if $?.signaled? && $?.termsig == 2 # SIGINT
+      safe_system "git", "checkout", @initial_branch
       safe_system "git", "reset", "--hard", @initial_revision
+      safe_system "git", "stash", "pop" if @stashed
     end
   end
 
diff --git a/Library/Homebrew/test/test_updater.rb b/Library/Homebrew/test/test_updater.rb
index c9f7ea4b8..849095833 100644
--- a/Library/Homebrew/test/test_updater.rb
+++ b/Library/Homebrew/test/test_updater.rb
@@ -5,7 +5,7 @@ require 'yaml'
 
 class UpdaterTests < Homebrew::TestCase
   class UpdaterMock < ::Updater
-    attr_accessor :diff
+    attr_accessor :diff, :expected, :called
 
     def initialize(*)
       super
@@ -29,10 +29,7 @@ class UpdaterTests < Homebrew::TestCase
       end
     end
     alias_method :safe_system, :`
-
-    def expectations_met?
-      @expected == @called
-    end
+    alias_method :system, :`
 
     def inspect
       "#<#{self.class.name}>"
@@ -60,14 +57,15 @@ class UpdaterTests < Homebrew::TestCase
     Formulary.stubs(:factory).returns(stub(:pkg_version => "1.0"))
     FormulaVersions.stubs(:new).returns(stub(:formula_at_revision => "2.0"))
     @updater.diff = fixture(fixture_name)
-    @updater.in_repo_expect("git checkout -q master")
+    @updater.in_repo_expect("git diff --quiet", true)
+    @updater.in_repo_expect("git symbolic-ref --short HEAD", "master")
     @updater.in_repo_expect("git rev-parse -q --verify HEAD", "1234abcd")
     @updater.in_repo_expect("git config core.autocrlf false")
-    @updater.in_repo_expect("git pull -q origin refs/heads/master:refs/remotes/origin/master")
+    @updater.in_repo_expect("git pull origin refs/heads/master:refs/remotes/origin/master --quiet")
     @updater.in_repo_expect("git rev-parse -q --verify HEAD", "3456cdef")
-    @updater.pull!
+    @updater.pull!(:silent => true)
     @report.update(@updater.report)
-    assert_predicate @updater, :expectations_met?
+    assert_equal @updater.expected, @updater.called
   end
 
   def test_update_homebrew_without_any_changes