lib/rswag/api/middleware.rb in rswag-api-2.1.1 vs lib/rswag/api/middleware.rb in rswag-api-2.2.0

@@ -1,10 +1,12 @@
 require 'json'
+require 'yaml'
+require 'rack/mime'
 
 module Rswag
   module Api
-    class Middleware 
+    class Middleware
 
       def initialize(app, config)
         @app = app
         @config = config
       end
@@ -12,26 +14,48 @@
       def call(env)
         path = env['PATH_INFO']
         filename = "#{@config.resolve_swagger_root(env)}/#{path}"
 
         if env['REQUEST_METHOD'] == 'GET' && File.file?(filename)
-          swagger = load_json(filename)
+          swagger = parse_file(filename)
           @config.swagger_filter.call(swagger, env) unless @config.swagger_filter.nil?
+          mime = Rack::Mime.mime_type(::File.extname(path), 'text/plain')
+          body = unload_swagger(filename, swagger)
 
           return [
             '200',
-            { 'Content-Type' => 'application/json' },
-            [ JSON.dump(swagger) ]
+            { 'Content-Type' => mime },
+            [ body ]
           ]
         end
-          
+
         return @app.call(env)
       end
 
       private
 
+      def parse_file(filename)
+        if /\.ya?ml$/ === filename
+          load_yaml(filename)
+        else
+          load_json(filename)
+        end
+      end
+
+      def load_yaml(filename)
+        YAML.safe_load(File.read(filename))
+      end
+
       def load_json(filename)
         JSON.parse(File.read(filename))
+      end
+
+      def unload_swagger(filename, swagger)
+        if /\.ya?ml$/ === filename
+          YAML.dump(swagger)
+        else
+          JSON.dump(swagger)
+        end
       end
     end
   end
 end