rubyzip · Jan 7, 2011 · Dec 6, 2010
diff --git a/.gitignore b/.gitignore
@@ -10,3 +10,4 @@ test/okToDeleteMoved.txt
 test/output.zip
 test/test_putOnClosedStream.zip
 test/zipWithDirs_copy.zip
+nbproject/*
diff --git a/lib/zip/zip.rb b/lib/zip/zip.rb
@@ -944,10 +944,14 @@ class ZipOutputStream
 
     # Opens the indicated zip file. If a file with that name already
     # exists it will be overwritten.
-    def initialize(fileName)
+    def initialize(fileName, stream=false)
       super()
       @fileName = fileName
-      @outputStream = File.new(@fileName, "wb")
+      if stream
+        @outputStream = StringIO.new
+      else
+        @outputStream = File.new(@fileName, "wb")
+      end
       @entrySet = ZipEntrySet.new
       @compressor = NullCompressor.instance
       @closed = false
@@ -966,6 +970,13 @@ def ZipOutputStream.open(fileName)
       zos.close if zos
     end
 
+    # Same as #open but writes to a filestream instead
+    def ZipOutputStream.write_buffer
+      zos = new('', true)
+      yield zos
+      return zos.close_buffer
+    end
+
     # Closes the stream and writes the central directory to the zip file
     def close
       return if @closed
@@ -976,6 +987,16 @@ def close
       @closed = true
     end
 
+    # Closes the stream and writes the central directory to the zip file
+    def close_buffer
+      return @outputStream if @closed
+      finalize_current_entry
+      update_local_headers
+      write_central_directory
+      @closed = true
+      return @outputStream
+    end
+
     # Closes the current entry and opens a new for writing.
     # +entry+ can be a ZipEntry object or a string.
     def put_next_entry(entryname, comment = nil, extra = nil, compression_method = ZipEntry::DEFLATED,  level = Zlib::DEFAULT_COMPRESSION)
@@ -1385,11 +1406,11 @@ class ZipFile < ZipCentralDirectory
 
     # Opens a zip archive. Pass true as the second parameter to create
     # a new archive if it doesn't exist already.
-    def initialize(fileName, create = nil)
+    def initialize(fileName, create = nil, buffer = false)
       super()
       @name = fileName
       @comment = ""
-      if (File.exists?(fileName))
+      if (File.exists?(fileName)) and !buffer
 	File.open(name, "rb") { |f| read_from_stream(f) }
       elsif (create)
 	@entrySet = ZipEntrySet.new
@@ -1420,6 +1441,17 @@ def ZipFile.open(fileName, create = nil)
       end
     end
 
+    # Same as #open. But outputs data to a buffer instead of a file
+    def ZipFile.add_buffer
+      zf = ZipFile.new('', true, true)
+      begin
+        yield zf
+      ensure
+        buffer = zf.write_buffer
+        return buffer
+      end
+    end
+
     # Returns the zip files comment, if it has one
     attr_accessor :comment
 
@@ -1521,6 +1553,16 @@ def commit
       initialize(name)
     end
 
+    # Write buffer write changes to buffer and return
+    def write_buffer
+      buffer = ZipOutputStream.write_buffer do |zos|
+        @entrySet.each { |e| e.write_to_zip_output_stream(zos) }
+        zos.comment = comment
+      end
+      return buffer
+    end
+
+
     # Closes the zip file committing any changes that has been made.
     def close
       commit

diff --git a/test/ziptest.rb b/test/ziptest.rb
@@ -614,6 +614,16 @@ def test_open
     assert_test_zip_contents(TEST_ZIP)
   end
 
+  def test_write_buffer
+    buffer = ZipOutputStream.write_buffer {
+      |zos|
+      zos.comment = TEST_ZIP.comment
+      write_test_zip(zos)
+    }
+    File.open(TEST_ZIP.zip_name, 'w') { |f| f.write buffer.string }
+    assert_test_zip_contents(TEST_ZIP)
+  end
+
   def test_writingToClosedStream
     assert_i_o_error_in_closed_stream { |zos| zos << "hello world" }
     assert_i_o_error_in_closed_stream { |zos| zos.puts "hello world" }
@@ -1041,6 +1051,23 @@ def setup
 class ZipFileTest < Test::Unit::TestCase
   include CommonZipFileFixture
 
+  def test_createFromScratchToBuffer
+    comment  = "a short comment"
+
+    buffer = ZipFile.add_buffer do |zf|
+      zf.get_output_stream("myFile") { |os| os.write "myFile contains just this" }
+      zf.mkdir("dir1")
+      zf.comment = comment
+    end
+
+    File.open(EMPTY_FILENAME, 'w') { |file| file.write buffer.string }
+    `cp #{EMPTY_FILENAME} ~/test.zip`
+
+    zfRead = ZipFile.new(EMPTY_FILENAME)
+    assert_equal(comment, zfRead.comment)
+    assert_equal(2, zfRead.entries.length)
+  end
+
   def test_createFromScratch
     comment  = "a short comment"
 
@@ -1288,6 +1315,21 @@ def test_commit
     zf.close
   end
 
+  def test_write_buffer
+    newName = "renamedFirst"
+    zf = ZipFile.new(TEST_ZIP.zip_name)
+    oldName = zf.entries.first
+    zf.rename(oldName, newName)
+    buffer = zf.write_buffer
+    File.open(TEST_ZIP.zip_name, 'w') { |f| f.write buffer.string }
+    zfRead = ZipFile.new(TEST_ZIP.zip_name)
+    assert(zfRead.entries.detect { |e| e.name == newName } != nil)
+    assert(zfRead.entries.detect { |e| e.name == oldName } == nil)
+    zfRead.close
+
+    zf.close
+  end
+
   # This test tests that after commit, you
   # can delete the file you used to add the entry to the zip file
   # with