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

Better calendar import

This commit is contained in:
Emmy D'Anello 2024-02-10 19:54:10 +01:00
parent 11cf6bbf55
commit fd4157acbd
Signed by: ynerant
GPG Key ID: 3A75C55819C8CF85
1 changed files with 13 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
sncfgtfs/management/commands/update_sncf_gtfs.py
View File

@ -160,9 +160,8 @@ class Command(BaseCommand):
routes.clear() routes.clear()
Calendar.objects.filter(transport_type=transport_type).delete() Calendar.objects.filter(transport_type=transport_type).delete()
calendar_ids = [] calendars = {}
if "calendar.txt" in zipfile.namelist(): if "calendar.txt" in zipfile.namelist():
calendars = []
for calendar_dict in csv.DictReader(read_file("calendar.txt")): for calendar_dict in csv.DictReader(read_file("calendar.txt")):
calendar_dict: dict calendar_dict: dict
calendar = Calendar( calendar = Calendar(
@ -197,7 +196,6 @@ class Command(BaseCommand):
unique_fields=['id']) unique_fields=['id'])
calendars.clear() calendars.clear()
calendars = []
calendar_dates = [] calendar_dates = []
for calendar_date_dict in csv.DictReader(read_file("calendar_dates.txt")): for calendar_date_dict in csv.DictReader(read_file("calendar_dates.txt")):
calendar_date_dict: dict calendar_date_dict: dict
@ -209,7 +207,7 @@ class Command(BaseCommand):
) )
calendar_dates.append(calendar_date) calendar_dates.append(calendar_date)
if calendar_date.service_id not in calendar_ids: if calendar_date.service_id not in calendars:
calendar = Calendar( calendar = Calendar(
id=f"{transport_type}-{calendar_date_dict['service_id']}", id=f"{transport_type}-{calendar_date_dict['service_id']}",
monday=False, monday=False,
@ -223,26 +221,22 @@ class Command(BaseCommand):
end_date=calendar_date_dict['date'], end_date=calendar_date_dict['date'],
transport_type=transport_type, transport_type=transport_type,
) )
calendars.append(calendar) calendars[calendar.id] = calendar
else:
if len(calendar_dates) >= bulk_size and not dry_run: calendar = calendars[f"{transport_type}-{calendar_date_dict['service_id']}"]
Calendar.objects.bulk_create(calendars, if calendar.start_date > calendar_date.date:
update_conflicts=True, calendar.start_date = calendar_date.date
update_fields=['end_date'], if calendar.end_date < calendar_date.date:
unique_fields=['id']) calendar.end_date = calendar_date.date
CalendarDate.objects.bulk_create(calendar_dates,
update_conflicts=True,
update_fields=['service_id', 'date', 'exception_type'],
unique_fields=['id'])
calendars.clear()
calendar_dates.clear()
if calendar_dates and not dry_run: if calendar_dates and not dry_run:
Calendar.objects.bulk_create(calendars, Calendar.objects.bulk_create(calendars.values(),
batch_size=bulk_size,
update_conflicts=True, update_conflicts=True,
update_fields=['end_date'], update_fields=['start_date', 'end_date'],
unique_fields=['id']) unique_fields=['id'])
CalendarDate.objects.bulk_create(calendar_dates, CalendarDate.objects.bulk_create(calendar_dates,
batch_size=bulk_size,
update_conflicts=True, update_conflicts=True,
update_fields=['service_id', 'date', 'exception_type'], update_fields=['service_id', 'date', 'exception_type'],
unique_fields=['id']) unique_fields=['id'])