ynerant
/
trainvel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

More calendar optimizations

This commit is contained in:
Emmy D'Anello 2024-05-12 14:03:48 +02:00
parent 2b6523c728
commit b65dc10bc6
Signed by: ynerant
GPG Key ID: 3A75C55819C8CF85
1 changed files with 15 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
trainvel/gtfs/management/commands/update_trainvel_gtfs.py
View File

@ -225,21 +225,23 @@ class Command(BaseCommand):
calendars.clear() calendars.clear()
calendar_dates = [] calendar_dates = []
all_calendars = {calendar.id: calendar for calendar in Calendar.objects.filter(gtfs_feed_id=gtfs_code)}
new_calendars = {}
with transaction.atomic(): with transaction.atomic():
for calendar_date_dict in read_csv("calendar_dates.txt"): for calendar_date_dict in read_csv("calendar_dates.txt"):
calendar_date_dict: dict calendar_date_dict: dict
service_id = f"{gtfs_code}-{calendar_date_dict['service_id']}" service_id = f"{gtfs_code}-{calendar_date_dict['service_id']}"
date = calendar_date_dict['date'] date = datetime.fromisoformat(calendar_date_dict['date']).date()
calendar_date = CalendarDate( calendar_date = CalendarDate(
id=f"{gtfs_code}-{calendar_date_dict['service_id']}-{calendar_date_dict['date']}", id=f"{gtfs_code}-{calendar_date_dict['service_id']}-{calendar_date_dict['date']}",
service_id=service_id, service_id=service_id,
date=date, date=calendar_date_dict['date'],
exception_type=calendar_date_dict['exception_type'], exception_type=calendar_date_dict['exception_type'],
) )
calendar_dates.append(calendar_date) calendar_dates.append(calendar_date)
if not Calendar.objects.filter(id=calendar_date.service_id).exists(): if service_id not in all_calendars:
calendar = Calendar( calendar = Calendar(
id=service_id, id=service_id,
monday=False, monday=False,
@ -249,26 +251,27 @@ class Command(BaseCommand):
friday=False, friday=False,
saturday=False, saturday=False,
sunday=False, sunday=False,
start_date=calendar_date_dict['date'], start_date=date,
end_date=calendar_date_dict['date'], end_date=date,
gtfs_feed_id=gtfs_code, gtfs_feed_id=gtfs_code,
) )
calendar.save() all_calendars[service_id] = calendar
new_calendars[service_id] = calendar
else: else:
calendar = Calendar.objects.get(id=service_id) calendar = all_calendars[service_id]
if calendar.start_date.isoformat() > date: if calendar.start_date > date:
calendar.start_date = date calendar.start_date = date
calendar.save() if calendar.end_date < date:
if calendar.end_date.isoformat() < date:
calendar.end_date = date calendar.end_date = date
calendar.save()
if len(calendar_dates) >= bulk_size and not dry_run: if len(calendar_dates) >= bulk_size and not dry_run:
CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size) CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size)
calendar_dates.clear() calendar_dates.clear()
if calendar_dates and not dry_run: if (calendar_dates or new_calendars) and not dry_run:
Calendar.objects.bulk_create(new_calendars.values(), batch_size=bulk_size)
CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size) CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size)
new_calendars.clear()
calendar_dates.clear() calendar_dates.clear()
trips = [] trips = []