generateschema.py 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  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. parser.add_argument('--file', dest="file", default=None, type=str)
  23. def handle(self, *args, **options):
  24. if options['generator_class']:
  25. generator_class = import_string(options['generator_class'])
  26. else:
  27. generator_class = self.get_generator_class()
  28. generator = generator_class(
  29. url=options['url'],
  30. title=options['title'],
  31. description=options['description'],
  32. urlconf=options['urlconf'],
  33. )
  34. schema = generator.get_schema(request=None, public=True)
  35. renderer = self.get_renderer(options['format'])
  36. output = renderer.render(schema, renderer_context={})
  37. if options['file']:
  38. with open(options['file'], 'wb') as f:
  39. f.write(output)
  40. else:
  41. self.stdout.write(output.decode())
  42. def get_renderer(self, format):
  43. if self.get_mode() == COREAPI_MODE:
  44. renderer_cls = {
  45. 'corejson': renderers.CoreJSONRenderer,
  46. 'openapi': renderers.CoreAPIOpenAPIRenderer,
  47. 'openapi-json': renderers.CoreAPIJSONOpenAPIRenderer,
  48. }[format]
  49. return renderer_cls()
  50. renderer_cls = {
  51. 'openapi': renderers.OpenAPIRenderer,
  52. 'openapi-json': renderers.JSONOpenAPIRenderer,
  53. }[format]
  54. return renderer_cls()
  55. def get_generator_class(self):
  56. if self.get_mode() == COREAPI_MODE:
  57. return coreapi.SchemaGenerator
  58. return SchemaGenerator