from datetime import datetime, timedelta, date from django.db.models import F, Q 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:])) 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_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_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_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:])) 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_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_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_time_24h=F('arrival_time') + timedelta(days=1)) return qs_today.union(qs_yesterday).union(qs_tomorrow).order_by("arrival_time_24h").all()