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

10
    BYPASS_DICT_MODE = False
11
    LIST_SHOULD_BE_DETAIL = False
12
13
14

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

        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(
25
26
                    "There should be no more than one element here check your queryset !"
                )
27
28

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

32
        return response