From e49a50c6c38a5a8c60e1e0adb1f0f1ed8bfaa7cd Mon Sep 17 00:00:00 2001 From: Alan Crosswell Date: Thu, 6 Jun 2019 12:42:43 -0400 Subject: [PATCH 1/2] add --generator_class CLI option to generateschema --- docs/community/release-notes.md | 1 + rest_framework/management/commands/generateschema.py | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/community/release-notes.md b/docs/community/release-notes.md index a27de2b641..0f88adfc3f 100644 --- a/docs/community/release-notes.md +++ b/docs/community/release-notes.md @@ -45,6 +45,7 @@ You can determine your currently installed version using `pip show`: **Date**: [Unreleased][3.10.0-milestone] * Resolve DeprecationWarning with markdown. [#6317][gh6317] +* Add `generateschema --generator_class` CLI option ## 3.9.x series diff --git a/rest_framework/management/commands/generateschema.py b/rest_framework/management/commands/generateschema.py index 631f402908..46077c3313 100644 --- a/rest_framework/management/commands/generateschema.py +++ b/rest_framework/management/commands/generateschema.py @@ -1,4 +1,5 @@ from django.core.management.base import BaseCommand +from django.utils.module_loading import import_string from rest_framework import renderers from rest_framework.schemas import coreapi @@ -22,9 +23,13 @@ def add_arguments(self, parser): parser.add_argument('--format', dest="format", choices=['openapi', 'openapi-json', 'corejson'], default='openapi', type=str) else: parser.add_argument('--format', dest="format", choices=['openapi', 'openapi-json'], default='openapi', type=str) + parser.add_argument('--generator_class', dest="generator_class", default=None, type=str) def handle(self, *args, **options): - generator_class = self.get_generator_class() + if options['generator_class']: + generator_class = import_string(options['generator_class']) + else: + generator_class = self.get_generator_class() generator = generator_class( url=options['url'], title=options['title'], From 8d9ae9a77790776dcc6a84dd858374e0017bfaef Mon Sep 17 00:00:00 2001 From: Carlton Gibson Date: Sun, 9 Jun 2019 14:22:02 +0200 Subject: [PATCH 2/2] =?UTF-8?q?Add=20test=20for=20generateschema=20?= =?UTF-8?q?=E2=80=94generator=5Fclass=20flag.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/schemas/test_managementcommand.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/schemas/test_managementcommand.py b/tests/schemas/test_managementcommand.py index 42e0e99d16..6cdf7f8b16 100644 --- a/tests/schemas/test_managementcommand.py +++ b/tests/schemas/test_managementcommand.py @@ -22,6 +22,16 @@ def get(self, request): ] +class CustomSchemaGenerator: + SCHEMA = {"key": "value"} + + def __init__(self, *args, **kwargs): + pass + + def get_schema(self, **kwargs): + return self.SCHEMA + + @override_settings(ROOT_URLCONF=__name__) @pytest.mark.skipif(not uritemplate, reason='uritemplate is not installed') class GenerateSchemaTests(TestCase): @@ -56,6 +66,13 @@ def test_renders_openapi_json_schema(self): out_json = json.loads(self.out.getvalue()) assert out_json['openapi'] == '3.0.2' + def test_accepts_custom_schema_generator(self): + call_command('generateschema', + '--generator_class={}.{}'.format(__name__, CustomSchemaGenerator.__name__), + stdout=self.out) + out_json = yaml.safe_load(self.out.getvalue()) + assert out_json == CustomSchemaGenerator.SCHEMA + @pytest.mark.skipif(yaml is None, reason='PyYAML is required.') @override_settings(REST_FRAMEWORK={'DEFAULT_SCHEMA_CLASS': 'rest_framework.schemas.AutoSchema'}) def test_coreapi_renders_default_schema_with_custom_title_url_and_description(self):