diff --git a/sncf/api/serializers.py b/sncf/api/serializers.py index e61d56c..4307d84 100644 --- a/sncf/api/serializers.py +++ b/sncf/api/serializers.py @@ -33,6 +33,7 @@ class StopTimeSerializer(serializers.ModelSerializer): departure_date = serializers.DateField(required=False) arrival_time_24h = serializers.DurationField(required=False) departure_time_24h = serializers.DurationField(required=False) + departure_time_real = serializers.CharField(required=False) class Meta: model = StopTime diff --git a/sncf/api/views.py b/sncf/api/views.py index 30d15ad..45c4bb4 100644 --- a/sncf/api/views.py +++ b/sncf/api/views.py @@ -1,6 +1,6 @@ from datetime import datetime, timedelta, date -from django.db.models import F, Q, Value +from django.db.models import F, Q, Value, When, Case, Exists, OuterRef from django.utils.decorators import method_decorator from django.views.decorators.cache import cache_control from django.views.decorators.http import last_modified @@ -171,27 +171,41 @@ class NextDeparturesViewSet(viewsets.ReadOnlyModelViewSet): .values_list('service_id', flat=True))) .values_list('id')) + def stop_time_update_qs(d: date): + return StopTimeUpdate.objects.filter(trip_update__start_date=d) \ + .exclude(departure_time=datetime.fromtimestamp(0)).filter(stop_time_id=OuterRef('pk')) + + def departure_time_real(d: date): + return Case( + When( + condition=Exists(stop_time_update_qs(d)), + then=F('departure_time') + stop_time_update_qs(d).values('departure_delay'), + ), + default=F('departure_time'), + ) + def canceled_filter(d: date): - return Q(Q(update__schedule_relationship=1) | Q(update__trip_update__schedule_relationship=3), - Q(update__trip_update__start_date=d), - ~Q(update__departure_time=datetime.fromtimestamp(0))) + return Exists(stop_time_update_qs(d).filter(Q(schedule_relationship=1) | Q(schedule_relationship=3))) qs_today = StopTime.objects.filter(stop_filter) \ - .filter(Q(departure_time__gte=query_time - F('update__departure_delay'))) \ + .annotate(departure_time_real=departure_time_real(query_date)) \ + .filter(departure_time_real__gte=query_time) \ .filter(Q(pickup_type=0) | canceled_filter(query_date)) \ .filter(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 - F('update__departure_delay'))) \ + .annotate(departure_time_real=departure_time_real(query_date)) \ + .filter(departure_time_real__gte=time_yesterday) \ .filter(Q(pickup_type=0) | canceled_filter(yesterday)) \ .filter(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))) \ + .annotate(departure_time_real=departure_time_real(query_date)) \ + .filter(departure_time_real__gte=timedelta(0)) \ .filter(Q(pickup_type=0) | canceled_filter(tomorrow)) \ .filter(calendar_filter(tomorrow)) \ .annotate(departure_date=Value(tomorrow)) \ @@ -243,27 +257,41 @@ class NextArrivalsViewSet(viewsets.ReadOnlyModelViewSet): .values_list('service_id', flat=True))) .values_list('id')) + def stop_time_update_qs(d: date): + return StopTimeUpdate.objects.filter(trip_update__start_date=d) \ + .exclude(arrival_time=datetime.fromtimestamp(0)).filter(stop_time_id=OuterRef('pk')) + + def arrival_time_real(d: date): + return Case( + When( + condition=Exists(stop_time_update_qs(d)), + then=F('arrival_time') + stop_time_update_qs(d).values('arrival_delay'), + ), + default=F('arrival_time'), + ) + def canceled_filter(d: date): - return Q(Q(update__schedule_relationship=1) | Q(update__trip_update__schedule_relationship=3), - Q(update__trip_update__start_date=d), - ~Q(update__arrival_time=datetime.fromtimestamp(0))) + return Exists(stop_time_update_qs(d).filter(Q(schedule_relationship=1) | Q(schedule_relationship=3))) qs_today = StopTime.objects.filter(stop_filter) \ - .filter(Q(departure_time__gte=query_time - F('update__arrival_delay'))) \ + .annotate(arrival_time_real=arrival_time_real(query_date)) \ + .filter(arrival_time_real__gte=query_time) \ .filter(Q(drop_off_type=0) | canceled_filter(query_date)) \ .filter(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(departure_time__gte=time_yesterday - F('update__arrival_delay'))) \ + .annotate(arrival_time_real=arrival_time_real(yesterday)) \ + .filter(arrival_time_real__gte=time_yesterday) \ .filter(Q(drop_off_type=0) | canceled_filter(yesterday)) \ .filter(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(departure_time__gte=timedelta(0))) \ + .annotate(arrival_time_real=arrival_time_real(tomorrow)) \ + .filter(arrival_time_real__gte=timedelta(0)) \ .filter(Q(drop_off_type=0) | canceled_filter(tomorrow)) \ .filter(calendar_filter(tomorrow)) \ .annotate(arrival_date=Value(tomorrow)) \