yasg.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. from django.conf import settings
  2. from drf_yasg.generators import EndpointEnumerator, OpenAPISchemaGenerator
  3. from drf_yasg.inspectors import SwaggerAutoSchema
  4. # Work around incompatibility with dj rest auth endpoint
  5. # https://github.com/axnsan12/drf-yasg/issues/435
  6. class CustomEndpointEnumerator(EndpointEnumerator):
  7. """
  8. Add custom setting to exclude views
  9. """
  10. def should_include_endpoint(
  11. self, path, callback, app_name="", namespace="", url_name=None
  12. ):
  13. view = "{}.{}".format(callback.__module__, callback.__qualname__)
  14. if view in settings.DRF_YASG_EXCLUDE_VIEWS:
  15. return False
  16. return super().should_include_endpoint(
  17. path, callback, app_name, namespace, url_name
  18. )
  19. class CustomOpenAPISchemaGenerator(OpenAPISchemaGenerator):
  20. """
  21. We want change default endpoint enumerator class
  22. """
  23. endpoint_enumerator_class = CustomEndpointEnumerator
  24. # Work around to group viewset views together instead of on giant "api" group
  25. # https://github.com/axnsan12/drf-yasg/issues/489
  26. class SquadSwaggerAutoSchema(SwaggerAutoSchema):
  27. def get_tags(self, operation_keys=None):
  28. tags = super().get_tags(operation_keys)
  29. if "api" in tags and operation_keys:
  30. # NOTE: `operation_keys` is a list like ["api", "v1", "token", "read"].
  31. tags[0] = operation_keys[2]
  32. return tags