|
1 | 1 | from django.core.management.base import BaseCommand
|
2 | 2 |
|
3 |
| -from rest_framework.compat import coreapi |
4 |
| -from rest_framework.renderers import ( |
5 |
| - CoreJSONRenderer, JSONOpenAPIRenderer, OpenAPIRenderer |
6 |
| -) |
7 |
| -from rest_framework.schemas.generators import SchemaGenerator |
| 3 | +from rest_framework import renderers |
| 4 | +from rest_framework.schemas import coreapi |
| 5 | +from rest_framework.schemas.openapi import SchemaGenerator |
| 6 | + |
| 7 | +OPENAPI_MODE = 'openapi' |
| 8 | +COREAPI_MODE = 'coreapi' |
8 | 9 |
|
9 | 10 |
|
10 | 11 | class Command(BaseCommand):
|
11 | 12 | help = "Generates configured API schema for project."
|
12 | 13 |
|
| 14 | + def get_mode(self): |
| 15 | + return COREAPI_MODE if coreapi.is_enabled() else OPENAPI_MODE |
| 16 | + |
13 | 17 | def add_arguments(self, parser):
|
14 |
| - parser.add_argument('--title', dest="title", default=None, type=str) |
| 18 | + parser.add_argument('--title', dest="title", default='', type=str) |
15 | 19 | parser.add_argument('--url', dest="url", default=None, type=str)
|
16 | 20 | parser.add_argument('--description', dest="description", default=None, type=str)
|
17 |
| - parser.add_argument('--format', dest="format", choices=['openapi', 'openapi-json', 'corejson'], default='openapi', type=str) |
| 21 | + if self.get_mode() == COREAPI_MODE: |
| 22 | + parser.add_argument('--format', dest="format", choices=['openapi', 'openapi-json', 'corejson'], default='openapi', type=str) |
| 23 | + else: |
| 24 | + parser.add_argument('--format', dest="format", choices=['openapi', 'openapi-json'], default='openapi', type=str) |
18 | 25 |
|
19 | 26 | def handle(self, *args, **options):
|
20 |
| - assert coreapi is not None, 'coreapi must be installed.' |
21 |
| - |
22 |
| - generator = SchemaGenerator( |
| 27 | + generator_class = self.get_generator_class() |
| 28 | + generator = generator_class( |
23 | 29 | url=options['url'],
|
24 | 30 | title=options['title'],
|
25 | 31 | description=options['description']
|
26 | 32 | )
|
27 |
| - |
28 | 33 | schema = generator.get_schema(request=None, public=True)
|
29 |
| - |
30 | 34 | renderer = self.get_renderer(options['format'])
|
31 | 35 | output = renderer.render(schema, renderer_context={})
|
32 | 36 | self.stdout.write(output.decode())
|
33 | 37 |
|
34 | 38 | def get_renderer(self, format):
|
| 39 | + if self.get_mode() == COREAPI_MODE: |
| 40 | + renderer_cls = { |
| 41 | + 'corejson': renderers.CoreJSONRenderer, |
| 42 | + 'openapi': renderers.CoreAPIOpenAPIRenderer, |
| 43 | + 'openapi-json': renderers.CoreAPIJSONOpenAPIRenderer, |
| 44 | + }[format] |
| 45 | + return renderer_cls() |
| 46 | + |
35 | 47 | renderer_cls = {
|
36 |
| - 'corejson': CoreJSONRenderer, |
37 |
| - 'openapi': OpenAPIRenderer, |
38 |
| - 'openapi-json': JSONOpenAPIRenderer, |
| 48 | + 'openapi': renderers.OpenAPIRenderer, |
| 49 | + 'openapi-json': renderers.JSONOpenAPIRenderer, |
39 | 50 | }[format]
|
40 |
| - |
41 | 51 | return renderer_cls()
|
| 52 | + |
| 53 | + def get_generator_class(self): |
| 54 | + if self.get_mode() == COREAPI_MODE: |
| 55 | + return coreapi.SchemaGenerator |
| 56 | + return SchemaGenerator |
0 commit comments