From 6197f0f62df2c9bb91b198c91d3274d71267bd34 Mon Sep 17 00:00:00 2001
From: Artem Baikuzin <artem@asterapbx.ru>
Date: Tue, 21 Feb 2017 21:28:14 +0300
Subject: [PATCH] Add block support for Workbook#initialize. This closes #31

---
 lib/writeexcel/version.rb  |  2 +-
 lib/writeexcel/workbook.rb |  7 +++++++
 test/test_workbook.rb      | 10 ++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/lib/writeexcel/version.rb b/lib/writeexcel/version.rb
index 77f3937b..055d3bf7 100644
--- a/lib/writeexcel/version.rb
+++ b/lib/writeexcel/version.rb
@@ -1,5 +1,5 @@
 require 'writeexcel'
 
 class WriteExcel < Workbook
-  VERSION = "1.0.5"
+  VERSION = "1.0.6"
 end
diff --git a/lib/writeexcel/workbook.rb b/lib/writeexcel/workbook.rb
index f348273f..bb1634c2 100644
--- a/lib/writeexcel/workbook.rb
+++ b/lib/writeexcel/workbook.rb
@@ -130,6 +130,13 @@ def initialize(file, default_formats = {})
 
     # Set colour palette.
     set_palette_xl97
+
+    return unless block_given?
+    begin
+      yield self
+    ensure
+      close
+    end
   end
 
   #
diff --git a/test/test_workbook.rb b/test/test_workbook.rb
index 2711913f..0a6feaea 100644
--- a/test/test_workbook.rb
+++ b/test/test_workbook.rb
@@ -13,6 +13,16 @@ def test_new
     assert_kind_of(Workbook, @workbook)
   end
 
+  def test_new_with_block
+    test_file = StringIO.new
+    workbook = Workbook.new(test_file) do |book|
+      book.add_worksheet('test_block')
+    end
+
+    assert_kind_of(Workbook, workbook)
+    assert_true(test_file.closed?)
+  end
+
   def test_add_worksheet
     sheetnames = ['sheet1', 'sheet2']
     (0 .. sheetnames.size-1).each do |i|