Add --generator-class argument to management command

Closes #140

src/drf_yasg/app_settings.py

@@ -2,6 +2,7 @@ from django.conf import settings
 from rest_framework.settings import perform_import
 
 SWAGGER_DEFAULTS = {
+    'DEFAULT_GENERATOR_CLASS': 'drf_yasg.generators.OpenAPISchemaGenerator',
     'DEFAULT_AUTO_SCHEMA_CLASS': 'drf_yasg.inspectors.SwaggerAutoSchema',
 
     'DEFAULT_FIELD_INSPECTORS': [
@@ -68,6 +69,7 @@ REDOC_DEFAULTS = {
 }
 
 IMPORT_STRINGS = [
+    'DEFAULT_GENERATOR_CLASS',
     'DEFAULT_AUTO_SCHEMA_CLASS',
     'DEFAULT_FIELD_INSPECTORS',
     'DEFAULT_FILTER_INSPECTORS',

src/drf_yasg/management/commands/generate_swagger.py

@@ -2,6 +2,7 @@ import json
 import logging
 import os
 from collections import OrderedDict
+from importlib import import_module
 
 from django.contrib.auth.models import User
 from django.core.exceptions import ImproperlyConfigured
@@ -12,7 +13,15 @@ from rest_framework.views import APIView
 from ... import openapi
 from ...app_settings import swagger_settings
 from ...codecs import OpenAPICodecJson, OpenAPICodecYaml
-from ...generators import OpenAPISchemaGenerator
+
+
+def import_class(import_string):
+    if not import_string:
+        return None
+
+    module_path, class_name = import_string.rsplit('.', 1)
+    module = import_module(module_path)
+    return getattr(module, class_name)
 
 
 class Command(BaseCommand):
@@ -64,6 +73,11 @@ class Command(BaseCommand):
                  'OpenAPISchemaGenerator.get_schema().\n'
                  'This option implies --mock-request.'
         )
+        parser.add_argument(
+            '-g', '--generator-class', dest='generator_class_name',
+            default='',
+            help='Import string pointing to an OpenAPISchemaGenerator subclass to use for schema generation.'
+        )
 
     def write_schema(self, schema, stream, format):
         if format == 'json':
@@ -89,7 +103,8 @@ class Command(BaseCommand):
         request = APIView().initialize_request(request)
         return request
 
-    def handle(self, output_file, overwrite, format, api_url, mock, user, private, *args, **options):
+    def handle(self, output_file, overwrite, format, api_url, mock, user, private, generator_class_name,
+               *args, **kwargs):
         # disable logs of WARNING and below
         logging.disable(logging.WARNING)
 
@@ -117,7 +132,8 @@ class Command(BaseCommand):
 
         request = self.get_mock_request(api_url, format, user) if mock else None
 
-        generator = OpenAPISchemaGenerator(
+        generator_class = import_class(generator_class_name) or swagger_settings.DEFAULT_GENERATOR_CLASS
+        generator = generator_class(
             info=info,
             url=api_url
         )

src/drf_yasg/views.py

@@ -11,7 +11,6 @@ from rest_framework.settings import api_settings
 from rest_framework.views import APIView
 
 from .app_settings import swagger_settings
-from .generators import OpenAPISchemaGenerator
 from .renderers import OpenAPIRenderer, ReDocRenderer, SwaggerJSONRenderer, SwaggerUIRenderer, SwaggerYAMLRenderer
 
 SPEC_RENDERERS = (SwaggerYAMLRenderer, SwaggerJSONRenderer, OpenAPIRenderer)
@@ -46,7 +45,7 @@ def deferred_never_cache(view_func):
 
 
 def get_schema_view(info=None, url=None, patterns=None, urlconf=None, public=False, validators=None,
-                    generator_class=OpenAPISchemaGenerator,
+                    generator_class=swagger_settings.DEFAULT_GENERATOR_CLASS,
                     authentication_classes=api_settings.DEFAULT_AUTHENTICATION_CLASSES,
                     permission_classes=api_settings.DEFAULT_PERMISSION_CLASSES):
     """Create a SchemaView class with default renderers and generators.