added rails support

makevoid · makevoid · commit 9c1190ef1ec0 · 2012-04-10T02:13:10.000+02:00
diff --git a/lib/rack-tidy-ffi.rb b/lib/rack-tidy-ffi.rb
@@ -3,23 +3,27 @@
 
 class RackTidyFFI
   attr_accessor :headers
-  
+
   def initialize(app)
-    @app = app       
-  end             
-  
-  def call(env)   
+    @app = app
+  end
+
+  def call(env)
     @status, @headers, @response = @app.call(env)
-    # TODO: include all response attributes... needed?
-    # response_body = ""
-    # response_body = @response.each { |part| response_body += part }
     if @headers["Content-Type"].include? "text/html"
-      response = TidyMachine.new(@response.first).tidy
+      body = if @response.is_a? Array
+        @response.first
+      elsif @response.is_a? Rack::BodyProxy
+        @response.body.first
+      else
+        @response.body
+      end
+      response = TidyMachine.new(body).tidy
       response = "<!DOCTYPE HTML>\n#{response}"
       @headers["Content-Length"] = response.length.to_s
       @response = [response]
     end
     [@status, @headers, @response]
-  end            
-  
-end
+  end
+
+end
