views.py

from rest_framework import viewsets

from common import permissions
from . import models
from . import serializers


class DocumentViewSet(viewsets.ModelViewSet):
    serializer_class = serializers.DocumentSerializer
    permission_classes = (permissions.IsStaffOrStudent, )

    def get_queryset(self):
        user = self.request.user
        if user.profile.role == 'Staff':
            queryset = self.staff_queryset()
        else:
            queryset = self.student_queryset(user.profile)
        return queryset

    def staff_queryset(self):
        queryset = models.Document.objects.all()
        profile_id = self.request.query_params.get('profileID', None)
        solution_id = self.request.query_params.get('solutionID', None)
        if profile_id is not None and solution_id is not None:
            return queryset.filter(uploaded_by=profile_id, solution=solution_id)
        if profile_id is not None:
            return queryset.filter(uploaded_by=profile_id)
        if solution_id is not None:
            return queryset.filter(solution=solution_id)
        return queryset

    def student_queryset(self, profile):
        queryset = models.Document.objects.filter(uploaded_by=profile)
        solution_id = self.request.query_params.get('solutionID', None)
        if solution_id is not None:
            return queryset.filter(solution=solution_id)
        return queryset