launch_viewset.py

from requests import Response
from api.endpoints.common.views import PolymorphicModelViewSet
from api.models import Launch
from rest_framework.exceptions import ParseError
from api.utils.auto_prefetching import prefetch
from api.utils.docstrings import format_endpoint_docstring
from rest_framework.decorators import action
from datetime import datetime, timedelta

from ...common.serializers import LaunchBasicSerializer
from ...launch.filters import LaunchFilterSet
from ...launch.serializers import LaunchDetailedSerializer, LaunchNormalSerializer
from ...launch.filters import UpcomingLaunchFilterSet

class LaunchViewSet(PolymorphicModelViewSet):
    default_mode = PolymorphicModelViewSet.MODE_LIST

    path = "launches"
    model = Launch
    serializers = {
        "list": LaunchBasicSerializer,
        "normal": LaunchNormalSerializer,
        "detailed": LaunchDetailedSerializer,
    }
    filter_class = LaunchFilterSet
    search_fields = (
        "name",
        "rocket__configuration__name",
        "rocket__configuration__manufacturer__name",
        "launch_service_provider__name",
        "rocket__configuration__manufacturer__abbrev",
        "mission__name",
        "pad__location__name",
        "pad__name",
        "rocket__spacecraftflight__spacecraft__name",
        "launch_designator",
    )
    ordering_fields = (
        "id",
        "name",
        "net",
        "last_updated",
    )
    permission_groups = {"retrieve": ["_Public"], "list": ["_Public"], "previous": ["_Public"], "upcoming": ["_Public"], "this_day_in_history": ["_Public"]}

    # Override the get_queryset method for ordering
    def get_queryset(self):
        if self.action == "retrieve":
            qs = prefetch(self.model.objects.all(), list(self.serializers.values())[-1])
        else:
            mode = self.request.query_params.get("mode", "normal")
            qs = prefetch(self.model.objects.all(), self.serializers[mode])
        return qs.order_by("net", "net_precision", "name")
    
    @action(detail=False, methods=["get"], url_path="this_day_in_history")
    def this_day_in_history(self, request):
        """
        Custom endpoint to retrieve launches that occurred on this day in history.
        """
        today = datetime.now()
        month = today.month
        day = today.day

        launches = self.model.objects.filter(net__month=month, net__day=day)

        mode = self.request.query_params.get("mode", self.default_mode)
        if mode not in self.serializers:
            raise ParseError(f"mode '{mode}' not in {list(self.serializers.keys())}")

        page = self.paginate_queryset(launches)
        if page is not None:
            serializer = self.serializers[mode](page, many=True, context={"request": request})
            return self.get_paginated_response(serializer.data)

        serializer = self.serializers[mode](launches, many=True, context={"request": request})
        return Response(serializer.data)
    
    @action(detail=False, methods=["get"], url_path="upcoming")
    def upcoming(self, request):
        __doc__ = format_endpoint_docstring(
            self.path,
            modes=list(self.serializers.keys()),
            filters=self.filter_class,
            search_fields=self.search_fields,
            ordering_fields=self.ordering_fields,
            filter_example="pad__location=13",
            search_example="Starlink",
            ordering_example="-last_updated",
        )

        self.filter_class = UpcomingLaunchFilterSet
        day_ago = datetime.now() - timedelta(days=1)
        launches = self.model.objects.filter(net__gte=day_ago)
        

        mode = self.request.query_params.get("mode", self.default_mode)
        if mode not in self.serializers:
            raise ParseError(f"mode '{mode}' not in {list(self.serializers.keys())}")

        page = self.paginate_queryset(launches)
        if page is not None:
            serializer = self.serializers[mode](page, many=True, context={"request": request})
            return self.get_paginated_response(serializer.data)

        serializer = self.serializers[mode](launches, many=True, context={"request": request})
        return Response(serializer.data)
    
    @action(detail=False, methods=["get"], url_path="previous")
    def previous(self, request):
        __doc__ = format_endpoint_docstring(
            self.path,
            modes=list(self.serializers.keys()),
            filters=self.filter_class,
            search_fields=self.search_fields,
            ordering_fields=self.ordering_fields,
            filter_example="pad__location=13",
            search_example="Starlink",
            ordering_example="-last_updated",
        )
            
        launches = self.model.objects.filter(net__lt=datetime.now())

        mode = self.request.query_params.get("mode", self.default_mode)
        if mode not in self.serializers:
            raise ParseError(f"mode '{mode}' not in {list(self.serializers.keys())}")

        page = self.paginate_queryset(launches)
        if page is not None:
            serializer = self.serializers[mode](page, many=True, context={"request": request})
            return self.get_paginated_response(serializer.data)

        serializer = self.serializers[mode](launches, many=True, context={"request": request})
        return Response(serializer.data)

    __doc__ = format_endpoint_docstring(
        path,
        modes=list(serializers.keys()),
        filters=filter_class,
        search_fields=search_fields,
        ordering_fields=ordering_fields,
        filter_example="pad__location=13",
        search_example="Starlink",
        ordering_example="-last_updated",
    )