lib/drg/tasks/gemfile.rb in drg-0.4.1 vs lib/drg/tasks/gemfile.rb in drg-0.4.2

@@ -5,35 +5,46 @@
 
       def initialize
         @file = ::Bundler.default_gemfile
       end
 
+      # Saves a copy of @lines before changing it (note that #dup and #clone weren't working)
+      #
       # @param [GemfileLine] gem
       # @param [String] version to update the gem line with
       def update(gem, version)
+        @saved_lines = Marshal.load Marshal.dump(lines)
         lines[gem] = gem.update version
       end
 
       def find_by_name(name)
         lines.each_with_index.each do |line, index|
           next if line =~ /:?path:?\s*(=>)?\s*/
           next if line =~ /:?git(hub)?:?\s*(=>)?\s*/
-          return GemfileLine.new line, index if line =~ /gem\s*['"]#{name}["']/
+          return GemfileLine.new line, index, name if line =~ /gem\s*['"]#{name}["']/
         end
         nil
       end
 
       def lines
         @lines ||= File.readlines file
       end
 
+      def saved_lines
+        @saved_lines ||= lines
+      end
+
       def write
         File.open file, 'wb' do |f|
           lines.each do |line|
             f << line
           end
         end
       end
 
+      def rollback
+        @lines = saved_lines
+        write
+      end
     end
   end
 end