trainvel/sncf/api/views.py

from datetime import datetime, timedelta, date

from django.db.models import F, Q, Value
from django_filters.rest_framework import DjangoFilterBackend
from rest_framework import viewsets
from rest_framework.filters import OrderingFilter

from sncf.api.serializers import AgencySerializer, StopSerializer, RouteSerializer, TripSerializer, \
    StopTimeSerializer, CalendarSerializer, CalendarDateSerializer, TransferSerializer, \
    FeedInfoSerializer
from sncfgtfs.models import Agency, Stop, Route, Trip, StopTime, Calendar, CalendarDate, \
    Transfer, FeedInfo


class AgencyViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Agency.objects.all()
    serializer_class = AgencySerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class StopViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Stop.objects.all()
    serializer_class = StopSerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class RouteViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Route.objects.all()
    serializer_class = RouteSerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class TripViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Trip.objects.all()
    serializer_class = TripSerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class StopTimeViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = StopTime.objects.order_by('id').all()
    serializer_class = StopTimeSerializer
    filter_backends = [DjangoFilterBackend, OrderingFilter]
    filterset_fields = '__all__'
    ordering_fields = ['arrival_time', 'departure_time', 'stop_sequence', ]
    ordering = ['stop_sequence', ]


class CalendarViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Calendar.objects.all()
    serializer_class = CalendarSerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class CalendarDateViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = CalendarDate.objects.all()
    serializer_class = CalendarDateSerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class TransferViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = Transfer.objects.all()
    serializer_class = TransferSerializer
    filter_backends = [DjangoFilterBackend]


class FeedInfoViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = FeedInfo.objects.all()
    serializer_class = FeedInfoSerializer
    filter_backends = [DjangoFilterBackend]
    filterset_fields = '__all__'


class NextDeparturesViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = StopTime.objects.none()
    serializer_class = StopTimeSerializer
    filter_backends = [DjangoFilterBackend]

    def get_queryset(self):
        now = datetime.now()

        stop_id = self.request.query_params.get('stop_id', None)
        stop_name = self.request.query_params.get('stop_name', None)
        query_date = date.fromisoformat(self.request.query_params.get('date', now.date().isoformat()))
        query_time = self.request.query_params.get('time', now.time().isoformat(timespec='seconds'))
        query_time = timedelta(seconds=int(query_time[:2]) * 3600
                                       + int(query_time[3:5]) * 60
                                       + (int(query_time[6:]) if len(query_time) > 6 else 0))

        yesterday = query_date - timedelta(days=1)
        time_yesterday = query_time + timedelta(days=1)
        tomorrow = query_date + timedelta(days=1)

        stop_filter = Q(stop__location_type=0)
        if stop_id:
            stop = Stop.objects.get(id=stop_id)
            stops = Stop.objects.filter(Q(id=stop_id)
                                        | Q(parent_station=stop_id))
            if stop.location_type == 0:
                stops |= Stop.objects.filter(parent_station=stop.parent_station_id)
            stop_filter = Q(stop__in=stops.values_list('id', flat=True))
        elif stop_name:
            stops = Stop.objects.filter(name__iexact=stop_name).values_list('id', flat=True)
            stop_filter = Q(stop__in=stops)

        def calendar_filter(d: date):
            return Q(trip__service_id__in=CalendarDate.objects.filter(date=d, exception_type=1)
                     .values_list('service_id')) \
                | Q(trip__service_id__in=Calendar.objects.filter(
                    start_date__lte=d,
                    end_date__gte=d,
                    **{f"{d:%A}".lower(): True})
                    .filter(~Q(id__in=CalendarDate.objects.filter(date=d, exception_type=2)
                               .values_list('service_id', flat=True)))
                    .values_list('id'))

        qs_today = StopTime.objects.filter(stop_filter) \
            .filter(Q(departure_time__gte=query_time, pickup_type=0), calendar_filter(query_date)) \
            .annotate(departure_date=Value(query_date)) \
            .annotate(departure_time_24h=F('departure_time'))

        qs_yesterday = StopTime.objects.filter(stop_filter) \
            .filter(Q(departure_time__gte=time_yesterday, pickup_type=0), calendar_filter(yesterday)) \
            .annotate(departure_date=Value(yesterday)) \
            .annotate(departure_time_24h=F('departure_time') - timedelta(days=1))

        qs_tomorrow = StopTime.objects.filter(stop_filter) \
            .filter(Q(departure_time__gte=timedelta(0), pickup_type=0), calendar_filter(tomorrow)) \
            .annotate(departure_date=Value(tomorrow)) \
            .annotate(departure_time_24h=F('departure_time') + timedelta(days=1))

        return qs_today.union(qs_yesterday).union(qs_tomorrow).order_by("departure_time_24h").all()


class NextArrivalsViewSet(viewsets.ReadOnlyModelViewSet):
    queryset = StopTime.objects.none()
    serializer_class = StopTimeSerializer
    filter_backends = [DjangoFilterBackend]

    def get_queryset(self):
        now = datetime.now()

        stop_id = self.request.query_params.get('stop_id', None)
        stop_name = self.request.query_params.get('stop_name', None)
        query_date = date.fromisoformat(self.request.query_params.get('date', now.date().isoformat()))
        query_time = self.request.query_params.get('time', now.time().isoformat(timespec='seconds'))
        query_time = timedelta(seconds=int(query_time[:2]) * 3600
                                       + int(query_time[3:5]) * 60
                                       + (int(query_time[6:]) if len(query_time) > 6 else 0))

        yesterday = query_date - timedelta(days=1)
        time_yesterday = query_time + timedelta(days=1)
        tomorrow = query_date + timedelta(days=1)

        stop_filter = Q(stop__location_type=0)
        if stop_id:
            stop = Stop.objects.get(id=stop_id)
            stops = Stop.objects.filter(Q(id=stop_id)
                                        | Q(parent_station=stop_id))
            if stop.location_type == 0:
                stops |= Stop.objects.filter(parent_station=stop.parent_station_id)
            stop_filter = Q(stop__in=stops.values_list('id', flat=True))
        elif stop_name:
            stops = Stop.objects.filter(name__iexact=stop_name).values_list('id', flat=True)
            stop_filter = Q(stop__in=stops)

        def calendar_filter(d: date):
            return Q(trip__service_id__in=CalendarDate.objects.filter(date=d, exception_type=1)
                     .values_list('service_id')) \
                | Q(trip__service_id__in=Calendar.objects.filter(
                    start_date__lte=d,
                    end_date__gte=d,
                    **{f"{d:%A}".lower(): True})
                    .filter(~Q(id__in=CalendarDate.objects.filter(date=d, exception_type=2)
                               .values_list('service_id', flat=True)))
                    .values_list('id'))

        qs_today = StopTime.objects.filter(stop_filter) \
            .filter(Q(arrival_time__gte=query_time, drop_off_type=0), calendar_filter(query_date)) \
            .annotate(arrival_date=Value(query_date)) \
            .annotate(arrival_time_24h=F('arrival_time'))

        qs_yesterday = StopTime.objects.filter(stop_filter) \
            .filter(Q(arrival_time__gte=time_yesterday, drop_off_type=0), calendar_filter(yesterday)) \
            .annotate(arrival_date=Value(yesterday)) \
            .annotate(arrival_time_24h=F('arrival_time') - timedelta(days=1))

        qs_tomorrow = StopTime.objects.filter(stop_filter) \
            .filter(Q(arrival_time__gte=timedelta(0), drop_off_type=0), calendar_filter(tomorrow)) \
            .annotate(arrival_date=Value(tomorrow)) \
            .annotate(arrival_time_24h=F('arrival_time') + timedelta(days=1))

        return qs_today.union(qs_yesterday).union(qs_tomorrow).order_by("arrival_time_24h").all()