Add DEFAULT_SCHEMA_GENERATOR_CLASS setting

carltongibson · carltongibson · commit 84e8f4d7fa9b · 2018-08-09T10:03:49.000+02:00
diff --git a/rest_framework/management/commands/generate_schema.py b/rest_framework/management/commands/generate_schema.py
@@ -1,8 +1,8 @@
 from django.core.management.base import BaseCommand
 
-from rest_framework.renderers import CoreJSONRenderer
-from rest_framework.schemas import SchemaGenerator
-
+from rest_framework.renderers import CoreJSONRenderer, JSONRenderer
+from rest_framework.schemas.generators import OpenAPISchemaGenerator, SchemaGenerator
+from rest_framework.settings import api_settings
 
 class Command(BaseCommand):
     help = "Generates configured API schema for project."
@@ -26,11 +26,21 @@ def add_arguments(self, parser):
         pass
 
     def handle(self, *args, **options):
+        generator_class = self._get_generator_class()
+        generator = generator_class()
 
-        renderer = CoreJSONRenderer()
-        generator = SchemaGenerator()
         schema = generator.get_schema(request=None, public=True)
 
+        renderer = self._get_renderer(generator)
         rendered_schema = renderer.render(schema, renderer_context={}).decode('utf8')
 
         self.stdout.write(rendered_schema)
+
+    def _get_generator_class(self):
+        return api_settings.DEFAULT_SCHEMA_GENERATOR_CLASS
+
+    def _get_renderer(self, generator):
+        if isinstance(generator, OpenAPISchemaGenerator):
+            return JSONRenderer()
+        else:
+            return CoreJSONRenderer()
diff --git a/rest_framework/settings.py b/rest_framework/settings.py
@@ -57,6 +57,7 @@
 
     # Schema
     'DEFAULT_SCHEMA_CLASS': 'rest_framework.schemas.AutoSchema',
+    'DEFAULT_SCHEMA_GENERATOR_CLASS': 'rest_framework.schemas.generators.SchemaGenerator',
 
     # Throttling
     'DEFAULT_THROTTLE_RATES': {
@@ -144,6 +145,7 @@
     'DEFAULT_PAGINATION_CLASS',
     'DEFAULT_FILTER_BACKENDS',
     'DEFAULT_SCHEMA_CLASS',
+    'DEFAULT_SCHEMA_GENERATOR_CLASS',
     'EXCEPTION_HANDLER',
     'TEST_REQUEST_RENDERER_CLASSES',
     'UNAUTHENTICATED_USER',
