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

View File

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