trainvel/sncfgtfs/management/commands/update_sncf_gtfs_rt.py

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,
                }
            )