from datetime import timedelta, datetime, date, time from zoneinfo import ZoneInfo import requests from django.core.management import BaseCommand from django.db.models import Q from sncfgtfs.gtfs_realtime_pb2 import FeedMessage from sncfgtfs.models import Agency, Calendar, CalendarDate, ExceptionType, LocationType, PickupType, \ Route, RouteType, Stop, StopScheduleRelationship, StopTime, StopTimeUpdate, \ Trip, TripUpdate, TripScheduleRelationship class Command(BaseCommand): help = "Update the SNCF GTFS Realtime database." GTFS_RT_FEEDS = { "TGV": "https://proxy.transport.data.gouv.fr/resource/sncf-tgv-gtfs-rt-trip-updates", "IC": "https://proxy.transport.data.gouv.fr/resource/sncf-ic-gtfs-rt-trip-updates", "TER": "https://proxy.transport.data.gouv.fr/resource/sncf-ter-gtfs-rt-trip-updates", "TI": "https://thello.axelor.com/public/gtfs/GTFS-RT.bin", } def add_arguments(self, parser): pass def handle(self, *args, **options): for feed_type, feed_url in self.GTFS_RT_FEEDS.items(): self.stdout.write(f"Updating {feed_type} feed...") feed_message = FeedMessage() feed_message.ParseFromString(requests.get(feed_url).content) stop_times_updates = [] for entity in feed_message.entity: if entity.HasField("trip_update"): trip_update = entity.trip_update trip_id = trip_update.trip.trip_id if feed_type in ["TGV", "IC", "TER"]: trip_id = trip_id.split(":", 1)[0] start_date = date(year=int(trip_update.trip.start_date[:4]), month=int(trip_update.trip.start_date[4:6]), day=int(trip_update.trip.start_date[6:])) start_dt = datetime.combine(start_date, time(0), tzinfo=ZoneInfo("Europe/Paris")) if trip_update.trip.schedule_relationship == TripScheduleRelationship.ADDED: # C'est un trajet nouveau. On crée le trajet associé. self.create_trip(trip_update, trip_id, start_dt, feed_type) if not Trip.objects.filter(id=trip_id).exists(): self.stdout.write(f"Trip {trip_id} does not exist in the GTFS feed.") continue # Création du TripUpdate tu, _created = TripUpdate.objects.update_or_create( trip_id=trip_id, start_date=trip_update.trip.start_date, start_time=trip_update.trip.start_time, defaults=dict( schedule_relationship=trip_update.trip.schedule_relationship, ) ) for stop_sequence, stop_time_update in enumerate(trip_update.stop_time_update): stop_id = stop_time_update.stop_id if stop_id.startswith('StopArea:'): # On est dans le cadre d'une gare. On cherche le quai associé. if StopTime.objects.filter(trip_id=trip_id, stop__parent_station_id=stop_id).exists(): # U stop = StopTime.objects.get(trip_id=trip_id, stop__parent_station_id=stop_id).stop else: stops = [s for s in Stop.objects.filter(parent_station_id=stop_id).all() for s2 in StopTime.objects.filter(trip_id=trip_id).all() if s.stop_type in s2.stop.stop_type or s2.stop.stop_type in s.stop_type] stop = stops[0] if stops else Stop.objects.get(id=stop_id) st, _created = StopTime.objects.update_or_create( id=f"{trip_id}-{stop.id}", trip_id=trip_id, stop_id=stop.id, defaults={ "stop_sequence": stop_sequence, "arrival_time": datetime.fromtimestamp(stop_time_update.arrival.time, tz=ZoneInfo("Europe/Paris")) - start_dt, "departure_time": datetime.fromtimestamp(stop_time_update.departure.time, tz=ZoneInfo("Europe/Paris")) - start_dt, "pickup_type": (PickupType.REGULAR if stop_time_update.departure.time else PickupType.NONE), "drop_off_type": (PickupType.REGULAR if stop_time_update.arrival.time else PickupType.NONE), } ) elif stop_time_update.schedule_relationship == StopScheduleRelationship.SKIPPED: st = StopTime.objects.get(Q(stop=stop_id) | Q(stop__parent_station_id=stop_id), trip_id=trip_id) if st.pickup_type != PickupType.NONE or st.drop_off_type != PickupType.NONE: st.pickup_type = PickupType.NONE st.drop_off_type = PickupType.NONE st.save() else: qs = StopTime.objects.filter(Q(stop=stop_id) | Q(stop__parent_station_id=stop_id), trip_id=trip_id) if qs.count() == 1: st = qs.first() else: st = qs.get(stop_sequence=stop_sequence) if st.stop_sequence != stop_sequence: st.stop_sequence = stop_sequence st.save() st_update = StopTimeUpdate( trip_update=tu, stop_time=st, arrival_delay=timedelta(seconds=stop_time_update.arrival.delay), arrival_time=datetime.fromtimestamp(stop_time_update.arrival.time, tz=ZoneInfo("Europe/Paris")), departure_delay=timedelta(seconds=stop_time_update.departure.delay), departure_time=datetime.fromtimestamp(stop_time_update.departure.time, tz=ZoneInfo("Europe/Paris")), schedule_relationship=stop_time_update.schedule_relationship or StopScheduleRelationship.SCHEDULED, ) stop_times_updates.append(st_update) else: self.stdout.write(str(entity)) StopTimeUpdate.objects.bulk_create(stop_times_updates, update_conflicts=True, update_fields=['arrival_delay', 'arrival_time', 'departure_delay', 'departure_time'], unique_fields=['trip_update', 'stop_time']) def create_trip(self, trip_update, trip_id, start_dt, feed_type): headsign = trip_id[5:-1] trip_qs = Trip.objects.all() trip_ids = trip_qs.values_list('id', flat=True) first_stop_queryset = StopTime.objects.filter( stop__parent_station_id=trip_update.stop_time_update[0].stop_id, ).values('trip_id') last_stop_queryset = StopTime.objects.filter( stop__parent_station_id=trip_update.stop_time_update[-1].stop_id, ).values('trip_id') trip_ids = trip_ids.intersection(first_stop_queryset).intersection(last_stop_queryset) # print(trip_id, trip_ids) for stop_sequence, stop_time_update in enumerate(trip_update.stop_time_update): stop_id = stop_time_update.stop_id st_queryset = StopTime.objects.filter(stop__parent_station_id=stop_id) if stop_sequence == 0: st_queryset = st_queryset.filter(stop_sequence=0) # print(stop_sequence, Stop.objects.get(id=stop_id).name, stop_time_update) # print(trip_ids) # print(st_queryset.values('trip_id').all()) trip_ids_restrict = trip_ids.intersection(st_queryset.values('trip_id')) if trip_ids_restrict: trip_ids = trip_ids_restrict else: stop = Stop.objects.get(id=stop_id) self.stdout.write(self.style.WARNING(f"Warning: No trip is found passing by stop " f"{stop.name} ({stop_id})")) trip_ids = set(trip_ids) route_ids = set(Trip.objects.filter(id__in=trip_ids).values_list('route_id', flat=True)) self.stdout.write(f"{len(route_ids)} routes found on trip for new train {headsign}") if not route_ids: origin_id = trip_update.stop_time_update[0].stop_id origin = Stop.objects.get(id=origin_id) destination_id = trip_update.stop_time_update[-1].stop_id destination = Stop.objects.get(id=destination_id) trip_name = f"{origin.name} - {destination.name}" trip_reverse_name = f"{destination.name} - {origin.name}" route_qs = Route.objects.filter(long_name=trip_name, transport_type=feed_type) route_reverse_qs = Route.objects.filter(long_name=trip_reverse_name, transport_type=feed_type) if route_qs.exists(): route_ids = set(route_qs.values_list('id', flat=True)) elif route_reverse_qs.exists(): route_ids = set(route_reverse_qs.values_list('id', flat=True)) else: self.stdout.write(f"Route not found for trip {trip_id} ({trip_name}). Creating new one") route = Route.objects.create( id=f"CREATED-{trip_name}", agency=Agency.objects.filter(routes__transport_type=feed_type).first(), transport_type=feed_type, type=RouteType.RAIL, short_name=trip_name, long_name=trip_name, ) route_ids = {route.id} self.stdout.write(f"Route {route.id} created for trip {trip_id} ({trip_name})") elif len(route_ids) > 1: self.stdout.write(f"Multiple routes found for trip {trip_id}.") self.stdout.write(", ".join(route_ids)) route_id = route_ids.pop() Calendar.objects.update_or_create( id=f"{feed_type}-new-{headsign}", defaults={ "transport_type": feed_type, "monday": False, "tuesday": False, "wednesday": False, "thursday": False, "friday": False, "saturday": False, "sunday": False, "start_date": start_dt.date(), "end_date": start_dt.date(), } ) CalendarDate.objects.update_or_create( id=f"{feed_type}-{headsign}-{trip_update.trip.start_date}", defaults={ "service_id": f"{feed_type}-new-{headsign}", "date": trip_update.trip.start_date, "exception_type": ExceptionType.ADDED, } ) Trip.objects.update_or_create( id=trip_id, defaults={ "route_id": route_id, "service_id": f"{feed_type}-new-{headsign}", "headsign": headsign, "direction_id": trip_update.trip.direction_id, } ) sample_trip = Trip.objects.filter(id__in=trip_ids, route_id=route_id) sample_trip = sample_trip.first() if sample_trip.exists() else None for stop_sequence, stop_time_update in enumerate(trip_update.stop_time_update): stop_id = stop_time_update.stop_id stop = Stop.objects.get(id=stop_id) if stop.location_type == LocationType.STATION: if not StopTime.objects.filter(trip_id=trip_id).exists(): if sample_trip: stop = StopTime.objects.get(trip_id=sample_trip.id, stop__parent_station_id=stop_id).stop elif StopTime.objects.filter(trip_id=trip_id, stop__parent_station_id=stop_id).exists(): stop = StopTime.objects.get(trip_id=trip_id, stop__parent_station_id=stop_id).stop else: stops = [s for s in Stop.objects.filter(parent_station_id=stop_id).all() for s2 in StopTime.objects.filter(trip_id=trip_id).all() if s.stop_type in s2.stop.stop_type or s2.stop.stop_type in s.stop_type] stop = stops[0] if stops else stop stop_id = stop.id arr_time = datetime.fromtimestamp(stop_time_update.arrival.time, tz=ZoneInfo("Europe/Paris")) - start_dt dep_time = datetime.fromtimestamp(stop_time_update.departure.time, tz=ZoneInfo("Europe/Paris")) - start_dt pickup_type = PickupType.REGULAR if stop_time_update.departure.time and stop_sequence > 0 \ else PickupType.NONE drop_off_type = PickupType.REGULAR if stop_time_update.arrival.time \ and stop_sequence < len(trip_update.stop_time_update) - 1 else PickupType.NONE StopTime.objects.update_or_create( id=f"{trip_id}-{stop_id}", trip_id=trip_id, defaults={ "stop_id": stop_id, "stop_sequence": stop_sequence, "arrival_time": arr_time, "departure_time": dep_time, "pickup_type": pickup_type, "drop_off_type": drop_off_type, } )