Auto merge of #5681 - bundler:seg-pristine-specific-gems, r=colby-swandale

[Pristine] Allow passing a list of gems to pristine

In the process of testing this, I discovered and fixed a bug that had pristine git gems not get added files removed.

3 files changed, 65 insertions, 11 deletions

diff --git a/lib/bundler/cli.rb b/lib/bundler/cli.rb
index 10287d3ac7..c9ddec93f4 100644
--- a/lib/bundler/cli.rb
+++ b/lib/bundler/cli.rb
@@ -553,10 +553,10 @@ module Bundler
       Issue.new.run
     end
 
-    desc "pristine", "Restores installed gems to pristine condition from files located in the gem cache. Gem installed from a git repository will be issued `git checkout --force`."
-    def pristine
+    desc "pristine [GEMS...]", "Restores installed gems to pristine condition from files located in the gem cache. Gem installed from a git repository will be issued `git checkout --force`."
+    def pristine(*gems)
       require "bundler/cli/pristine"
-      Pristine.new.run
+      Pristine.new(gems).run
     end
 
     if Bundler.feature_flag.plugins?
diff --git a/lib/bundler/cli/pristine.rb b/lib/bundler/cli/pristine.rb
index 30542b583e..cfd90da34b 100644
--- a/lib/bundler/cli/pristine.rb
+++ b/lib/bundler/cli/pristine.rb
@@ -3,14 +3,21 @@ require "bundler/cli/common"
 
 module Bundler
   class CLI::Pristine
+    def initialize(gems)
+      @gems = gems
+    end
+
     def run
+      CLI::Common.ensure_all_gems_in_lockfile!(@gems)
+
       Bundler.load.specs.each do |spec|
         next if spec.name == "bundler" # Source::Rubygems doesn't install bundler
+        next if !@gems.empty? && !@gems.include?(spec.name)
 
         gem_name = "#{spec.name} (#{spec.version}#{spec.git_version})"
         gem_name += " (#{spec.platform})" if !spec.platform.nil? && spec.platform != Gem::Platform::RUBY
 
-        case spec.source
+        case source = spec.source
         when Source::Rubygems
           cached_gem = spec.cache_file
           unless File.exist?(cached_gem)
@@ -19,11 +26,11 @@ module Bundler
           end
 
           FileUtils.rm_rf spec.full_gem_path
-          spec.source.install(spec, :force => true)
+          source.install(spec, :force => true)
         when Source::Git
-          git_source = spec.source
-          git_source.remote!
-          git_source.install(spec, :force => true)
+          source.remote!
+          FileUtils.rm_rf spec.full_gem_path
+          source.install(spec, :force => true)
         else
           Bundler.ui.warn("Cannot pristine #{gem_name}. Gem is sourced from local path.")
         end
diff --git a/spec/commands/pristine_spec.rb b/spec/commands/pristine_spec.rb
index cc01aabbf7..9c3918f51e 100644
--- a/spec/commands/pristine_spec.rb
+++ b/spec/commands/pristine_spec.rb
@@ -52,12 +52,23 @@ RSpec.describe "bundle pristine" do
       changed_file = Pathname.new(spec.full_gem_path).join("lib/foo.rb")
       diff = "#Pristine spec changes"
 
-      File.open(changed_file, "a") {|f| f.puts "#Pristine spec changes" }
+      File.open(changed_file, "a") {|f| f.puts diff }
       expect(File.read(changed_file)).to include(diff)
 
-      bundle "pristine"
+      bundle! "pristine"
       expect(File.read(changed_file)).to_not include(diff)
     end
+
+    it "removes added files" do
+      spec = Bundler.definition.specs["foo"].first
+      changes_txt = Pathname.new(spec.full_gem_path).join("lib/changes.txt")
+
+      FileUtils.touch(changes_txt)
+      expect(changes_txt).to be_file
+
+      bundle! "pristine"
+      expect(changes_txt).not_to be_file
+    end
   end
 
   context "when sourced from gemspec" do
@@ -66,7 +77,7 @@ RSpec.describe "bundle pristine" do
       changed_file = Pathname.new(spec.full_gem_path).join("lib/baz.rb")
       diff = "#Pristine spec changes"
 
-      File.open(changed_file, "a") {|f| f.puts "#Pristine spec changes" }
+      File.open(changed_file, "a") {|f| f.puts diff }
       expect(File.read(changed_file)).to include(diff)
 
       bundle "pristine"
@@ -98,4 +109,40 @@ RSpec.describe "bundle pristine" do
       expect(changes_txt).to be_file
     end
   end
+
+  context "when passing a list of gems to pristine" do
+    it "resets them" do
+      foo = Bundler.definition.specs["foo"].first
+      foo_changes_txt = Pathname.new(foo.full_gem_path).join("lib/changes.txt")
+      FileUtils.touch(foo_changes_txt)
+      expect(foo_changes_txt).to be_file
+      foo_ref = Spec::Builders::GitReader.new(lib_path("foo")).ref_for("HEAD", 6)
+
+      bar = Bundler.definition.specs["bar"].first
+      bar_changes_txt = Pathname.new(bar.full_gem_path).join("lib/changes.txt")
+      FileUtils.touch(bar_changes_txt)
+      expect(bar_changes_txt).to be_file
+
+      weakling = Bundler.definition.specs["weakling"].first
+      weakling_changes_txt = Pathname.new(weakling.full_gem_path).join("lib/changes.txt")
+      FileUtils.touch(weakling_changes_txt)
+      expect(weakling_changes_txt).to be_file
+
+      bundle! "pristine foo bar weakling"
+
+      expect(out).to eq(strip_whitespace(<<-EOS).strip)
+        Cannot pristine bar (1.0). Gem is sourced from local path.
+        Using foo 1.0 from #{lib_path("foo")} (at master@#{foo_ref})
+        Installing weakling 1.0
+      EOS
+      expect(weakling_changes_txt).not_to be_file
+      expect(foo_changes_txt).not_to be_file
+      expect(bar_changes_txt).to be_file
+    end
+
+    it "raises when one of them is not in the lockfile" do
+      bundle "pristine abcabcabc"
+      expect(out).to include("Could not find gem 'abcabcabc'.")
+    end
+  end
 end