From 8b0da2607fa1bc4986fd73e105c2c475c595fd97 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cristi=20V=C3=AEjdea?= <cristi@cvjd.me>
Date: Mon, 24 Dec 2018 07:55:17 +0200
Subject: [PATCH] Add pretty and media_type options to codecs

---
 src/drf_yasg/codecs.py                        | 27 ++++++++++++++++---
 .../management/commands/generate_swagger.py   | 11 ++++----
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/drf_yasg/codecs.py b/src/drf_yasg/codecs.py
index a9cf557..2ed6980 100644
--- a/src/drf_yasg/codecs.py
+++ b/src/drf_yasg/codecs.py
@@ -88,7 +88,7 @@ class _OpenAPICodec(object):
 
         :param dict spec: a python dict
         :return: string representation of ``spec``
-        :rtype: str
+        :rtype: str or bytes
         """
         raise NotImplementedError("override this method")
 
@@ -105,9 +105,22 @@ class _OpenAPICodec(object):
 class OpenAPICodecJson(_OpenAPICodec):
     media_type = 'application/json'
 
+    def __init__(self, validators, pretty=False, media_type='application/json'):
+        super(OpenAPICodecJson, self).__init__(validators)
+        self.pretty = pretty
+        self.media_type = media_type
+
     def _dump_dict(self, spec):
-        """Dump ``spec`` into JSON."""
-        return json.dumps(spec)
+        """Dump ``spec`` into JSON.
+
+        :rtype: str"""
+        if self.pretty:
+            out = json.dumps(spec, indent=4, separators=(',', ': '))
+            if out[-1] != '\n':
+                out += '\n'
+            return out
+        else:
+            return json.dumps(spec)
 
 
 YAML_MAP_TAG = u'tag:yaml.org,2002:map'
@@ -201,6 +214,12 @@ def yaml_sane_load(stream):
 class OpenAPICodecYaml(_OpenAPICodec):
     media_type = 'application/yaml'
 
+    def __init__(self, validators, media_type='application/yaml'):
+        super(OpenAPICodecYaml, self).__init__(validators)
+        self.media_type = media_type
+
     def _dump_dict(self, spec):
-        """Dump ``spec`` into YAML."""
+        """Dump ``spec`` into YAML.
+
+        :rtype: bytes"""
         return yaml_sane_dump(spec, binary=True)
diff --git a/src/drf_yasg/management/commands/generate_swagger.py b/src/drf_yasg/management/commands/generate_swagger.py
index fef3e7f..3f1db30 100644
--- a/src/drf_yasg/management/commands/generate_swagger.py
+++ b/src/drf_yasg/management/commands/generate_swagger.py
@@ -56,7 +56,7 @@ class Command(BaseCommand):
         parser.add_argument(
             '-m', '--mock-request', dest='mock',
             default=False, action='store_true',
-            help='Use a mock request when generating the swagger schema. This is useful if your views or serializers'
+            help='Use a mock request when generating the swagger schema. This is useful if your views or serializers '
                  'depend on context from a request in order to function.'
         )
         parser.add_argument(
@@ -85,11 +85,9 @@ class Command(BaseCommand):
 
     def write_schema(self, schema, stream, format):
         if format == 'json':
-            codec = OpenAPICodecJson(validators=[])
-            swagger_json = codec.encode(schema)
-            swagger_json = json.loads(swagger_json.decode('utf-8'), object_pairs_hook=OrderedDict)
-            pretty_json = json.dumps(swagger_json, indent=4, ensure_ascii=True)
-            stream.write(pretty_json)
+            codec = OpenAPICodecJson(validators=[], pretty=True)
+            swagger_json = codec.encode(schema).decode('utf-8')
+            stream.write(swagger_json)
         elif format == 'yaml':
             codec = OpenAPICodecYaml(validators=[])
             swagger_yaml = codec.encode(schema).decode('utf-8')
@@ -142,6 +140,7 @@ class Command(BaseCommand):
         if mock:
             request = self.get_mock_request(api_url, format, user)
 
+
         if request and api_version:
             request.version = api_version