DRF flex field optimization

views.py

import coreapi
from django.core.exceptions import FieldDoesNotExist
from rest_flex_fields.utils import is_expanded
from rest_framework.exceptions import APIException
from rest_framework.schemas.inspectors import AutoSchema


class FlexFieldSchema(AutoSchema):
    def __init__(self, view, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._view = view

    def get_manual_fields(self, path, method):
        fields = super().get_manual_fields(path, method)
        try:
            serializer = self.view.get_serializer()
        except APIException:
            pass
        else:
            for field in getattr(serializer, 'expandable_fields', {}).keys():
                fields.append(coreapi.Field(
                    field,
                    location='expand',
                ))
        return fields


class FlexFieldsViewMixin:
    @property
    def schema(self):
        return FlexFieldSchema(self)

    def get_queryset(self):
        qs = super().get_queryset()
        select = []
        prefetch = []
        for name, settings in getattr(self.serializer_class, 'expandable_fields', {}).items():
            if is_expanded(self.request, name):
                source = settings[1].get('source', name)
                try:
                    field = qs.model._meta.get_field(source)
                except FieldDoesNotExist:
                    continue
                if field.is_relation and field.concrete:
                    if settings[1].get('many', False):
                        prefetch.append(settings[1].get('source', name))
                    else:
                        select.append(settings[1].get('source', name))
        if select:
            qs = qs.select_related(*select)
        if prefetch:
            qs = qs.prefetch_related(*prefetch)
        return qs