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, TripUpdate as GTFSTripUpdate from sncfgtfs.models import Agency, Calendar, CalendarDate, ExceptionType, GTFSFeed, LocationType, PickupType, \ Route, RouteType, Stop, StopScheduleRelationship, StopTime, StopTimeUpdate, \ Trip, TripUpdate, TripScheduleRelationship class Command(BaseCommand): help = "Update the SNCF GTFS Realtime database." def add_arguments(self, parser): parser.add_argument('--debug', '-d', action='store_true', help="Activate debug mode") def handle(self, debug: bool = False, verbosity: int = 1, *args, **options): for gtfs_feed in GTFSFeed.objects.all(): if not gtfs_feed.rt_feed_url: if verbosity >= 2: self.stdout.write(self.style.WARNING(f"No GTFS-RT feed found for {gtfs_feed}.")) continue self.stdout.write(f"Updating GTFS-RT feed for {gtfs_feed}…") gtfs_code = gtfs_feed.code feed_message = FeedMessage() feed_message.ParseFromString(requests.get(gtfs_feed.rt_feed_url, allow_redirects=True).content) stop_times_updates = [] if debug: with open(f'feed_message-{gtfs_code}.txt', 'w') as f: f.write(str(feed_message)) for entity in feed_message.entity: if entity.HasField("trip_update"): trip_update = entity.trip_update trip_id = trip_update.trip.trip_id trip_id = f"{gtfs_code}-{trip_id}" 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, gtfs_feed) 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 stop_id = f"{gtfs_code}-{stop_id}" if StopTime.objects.filter(trip_id=trip_id, stop=stop_id).exists(): st = StopTime.objects.filter(trip_id=trip_id, stop=stop_id) if st.count() > 1: st = st.get(stop_sequence=stop_sequence) else: st = st.first() else: # Stop is added st = StopTime.objects.create( id=f"{trip_id}-{stop_time_update.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), } ) if stop_time_update.schedule_relationship == StopScheduleRelationship.SKIPPED: 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() 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: GTFSTripUpdate, trip_id: str, start_dt: datetime, gtfs_feed: GTFSFeed) -> None: headsign = trip_id[5:-1] gtfs_code = gtfs_feed.code route, _created = Route.objects.get_or_create( id=f"{gtfs_code}-ADDED-{headsign}", gtfs_feed=gtfs_feed, type=RouteType.RAIL, short_name="ADDED", long_name="ADDED ROUTE", ) Calendar.objects.update_or_create( id=f"{gtfs_code}-ADDED-{headsign}", defaults={ "gtfs_feed": gtfs_feed, "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"{gtfs_code}-ADDED-{headsign}-{trip_update.trip.start_date}", defaults={ "service_id": f"{gtfs_code}-ADDED-{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"{gtfs_code}-ADDED-{headsign}", "headsign": headsign, "direction_id": trip_update.trip.direction_id, "gtfs_feed": gtfs_feed, } ) for stop_sequence, stop_time_update in enumerate(trip_update.stop_time_update): stop_id = stop_time_update.stop_id stop_id = f"{gtfs_code}-{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_time_update.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, } )