200 lines
9.7 KiB
Python
200 lines
9.7 KiB
Python
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,
|
|
}
|
|
)
|