From 8eb1abd459a25df03fece6994a0c2c8ae4e12f68 Mon Sep 17 00:00:00 2001
From: Konstantin Haase <konstantin.mailinglists@googlemail.com>
Date: Fri, 13 May 2011 10:35:17 +0200
Subject: [PATCH] fix handling of response.body

---
 lib/sinatra/base.rb   | 19 +++++++++++++++++--
 test/response_test.rb |  4 ++--
 test/routing_test.rb  |  2 +-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/lib/sinatra/base.rb b/lib/sinatra/base.rb
index ad69ce08..2ff2d206 100644
--- a/lib/sinatra/base.rb
+++ b/lib/sinatra/base.rb
@@ -59,6 +59,20 @@ module Sinatra
   # http://rack.rubyforge.org/doc/classes/Rack/Response.html
   # http://rack.rubyforge.org/doc/classes/Rack/Response/Helpers.html
   class Response < Rack::Response
+    def body=(value)
+      @body = value.respond_to?(:each) ? value : [value.to_str]
+    end
+
+    def each
+      block_given? ? super : enum_for(:each)
+    end
+
+    def finish
+      if body.respond_to? :to_ary and not [204, 304].include?(status.to_i)
+        headers["Content-Length"] = body.inject(0) { |l, p| l + Rack::Utils.bytesize(p) }.to_s
+      end
+      super
+    end
   end
 
   class NotFound < NameError #:nodoc:
@@ -627,9 +641,10 @@ module Sinatra
       @response['Content-Type'] = nil
       invoke { dispatch! }
       invoke { error_block!(response.status) }
+
       unless @response['Content-Type']
-        if body.respond_to?(:to_ary) and body.first.respond_to? :content_type
-          content_type body.first.content_type
+        if body.respond_to? :to_ary and body[0].respond_to? :content_type
+          content_type body[0].content_type
         else
           content_type :html
         end
diff --git a/test/response_test.rb b/test/response_test.rb
index b2d78e47..826defdc 100644
--- a/test/response_test.rb
+++ b/test/response_test.rb
@@ -22,7 +22,7 @@ class ResponseTest < Test::Unit::TestCase
   it 'writes to body' do
     @response.body = 'Hello'
     @response.write ' World'
-    assert_equal 'Hello World', @response.body
+    assert_equal 'Hello World', @response.body.join
   end
 
   [204, 304].each do |status_code|
@@ -37,6 +37,6 @@ class ResponseTest < Test::Unit::TestCase
     @response.body = ['Hello', 'World!', '✈']
     status, headers, body = @response.finish
     assert_equal '14', headers['Content-Length']
-    assert_equal @response.body, body
+    assert_equal @response.body, body.body
   end
 end
diff --git a/test/routing_test.rb b/test/routing_test.rb
index b8a1064a..9434bf2c 100644
--- a/test/routing_test.rb
+++ b/test/routing_test.rb
@@ -1047,7 +1047,7 @@ class RoutingTest < Test::Unit::TestCase
     mock_app do
       get '/foo' do
         status, headers, body = call env.merge("PATH_INFO" => '/bar')
-        [status, headers, body.map(&:upcase)]
+        [status, headers, body.each.map(&:upcase)]
       end
 
       get '/bar' do