generateschema.py 2.5 KB

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