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