DictModeViewSet.py 1.04 KB
Newer Older
1 2 3 4 5 6 7 8
from rest_framework import viewsets


class DictModeViewSet(viewsets.ModelViewSet):
    """
    ViewSet that renders data as dict with keys corresponding to the model
    primary key. Instead of list.
    """
9
    BYPASS_DICT_MODE = False
10
    LIST_SHOULD_BE_DETAIL = False
11 12 13 14

    def list(self, request, *args, **kwargs):
        response = super(viewsets.ModelViewSet, self).list(  # pylint: disable=E1003
            request, *args, **kwargs)  # call the original 'list'
15 16 17 18 19 20 21 22 23 24 25

        if self.LIST_SHOULD_BE_DETAIL:
            if len(response.data) == 0:
                response.data = dict()
            elif len(response.data) == 1:
                response.data = response.data[0]
            else:
                raise Exception(
                    "There should be no more than one element here check your queryset !")

        elif not self.BYPASS_DICT_MODE:
26 27 28 29
            pk_attr_name = self.serializer_class.Meta.model._meta.pk.name
            response.data = {
                d[pk_attr_name]: d for d in response.data
            }
30

31
        return response