myModelViewSet.py 1.4 KB
Newer Older
Florent Chehab's avatar
Florent Chehab committed
1
from backend_app.permissions import DEFAULT_VIEWSET_PERMISSIONS
Florent Chehab's avatar
Florent Chehab committed
2
from rest_framework import viewsets
3

Florent Chehab's avatar
Florent Chehab committed
4 5
from .myModelSerializer import MyModelSerializer

6

Florent Chehab's avatar
Florent Chehab committed
7
class MyModelViewSet(viewsets.ModelViewSet):
Florent Chehab's avatar
Florent Chehab committed
8 9 10 11
    """
    Custom default viewset
    """

12
    serializer_class = MyModelSerializer
13
    permission_classes = DEFAULT_VIEWSET_PERMISSIONS
14

Florent Chehab's avatar
Florent Chehab committed
15 16 17
    # Class attribute to tell that return 1 element and not a list of element
    # For example when querying userData, we don't specify the userId and query the viewset
    # in list mode; still we want to return only one
Florent Chehab's avatar
Florent Chehab committed
18 19 20
    LIST_SHOULD_BE_DETAIL = False

    def list(self, request, *args, **kwargs):
Florent Chehab's avatar
Florent Chehab committed
21 22 23 24
        """
        Extend the list function of the viewset class
        """
        response = super().list(request, *args, **kwargs)  # call the original 'list'
Florent Chehab's avatar
Florent Chehab committed
25 26 27 28 29 30 31 32 33 34 35 36 37

        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 in the queryset, check it !"
                )

        return response

38 39 40 41 42
    def get_queryset(self):
        """
        Extended default rest framework behavior
        to prefetch some table and enhance performances
        """
43
        return self.queryset.prefetch_related("moderated_by", "updated_by")