Keep arrival trains during 5 minutes

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

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),
+            return Exists(stop_time_update_qs(d).filter(Q(schedule_relationship=1) | Q(schedule_relationship=3)))
-                     Q(update__trip_update__start_date=d),
-                     ~Q(update__departure_time=datetime.fromtimestamp(0)))
 
         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)) \
@@ -215,6 +229,7 @@ class NextArrivalsViewSet(viewsets.ReadOnlyModelViewSet):
         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))
+        query_time -= timedelta(minutes=5)  # Keep the last trains of the 5 previous minutes
 
         yesterday = query_date - timedelta(days=1)
         time_yesterday = query_time + timedelta(days=1)
@@ -243,27 +258,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),
+            return Exists(stop_time_update_qs(d).filter(Q(schedule_relationship=1) | Q(schedule_relationship=3)))
-                     Q(update__trip_update__start_date=d),
-                     ~Q(update__arrival_time=datetime.fromtimestamp(0)))
 
         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)) \