diff --git a/sncfgtfs/management/commands/update_sncf_gtfs.py b/sncfgtfs/management/commands/update_sncf_gtfs.py index 6078b55..4718f85 100644 --- a/sncfgtfs/management/commands/update_sncf_gtfs.py +++ b/sncfgtfs/management/commands/update_sncf_gtfs.py @@ -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'])