diff --git a/trainvel/gtfs/management/commands/update_trainvel_gtfs.py b/trainvel/gtfs/management/commands/update_trainvel_gtfs.py index bd9e201..862436d 100644 --- a/trainvel/gtfs/management/commands/update_trainvel_gtfs.py +++ b/trainvel/gtfs/management/commands/update_trainvel_gtfs.py @@ -225,21 +225,23 @@ class Command(BaseCommand): calendars.clear() calendar_dates = [] + all_calendars = {calendar.id: calendar for calendar in Calendar.objects.filter(gtfs_feed_id=gtfs_code)} + new_calendars = {} with transaction.atomic(): for calendar_date_dict in read_csv("calendar_dates.txt"): calendar_date_dict: dict 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( id=f"{gtfs_code}-{calendar_date_dict['service_id']}-{calendar_date_dict['date']}", service_id=service_id, - date=date, + date=calendar_date_dict['date'], exception_type=calendar_date_dict['exception_type'], ) 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( id=service_id, monday=False, @@ -249,26 +251,27 @@ class Command(BaseCommand): friday=False, saturday=False, sunday=False, - start_date=calendar_date_dict['date'], - end_date=calendar_date_dict['date'], + start_date=date, + end_date=date, gtfs_feed_id=gtfs_code, ) - calendar.save() + all_calendars[service_id] = calendar + new_calendars[service_id] = calendar else: - calendar = Calendar.objects.get(id=service_id) - if calendar.start_date.isoformat() > date: + calendar = all_calendars[service_id] + if calendar.start_date > date: calendar.start_date = date - calendar.save() - if calendar.end_date.isoformat() < date: + if calendar.end_date < date: calendar.end_date = date - calendar.save() if len(calendar_dates) >= bulk_size and not dry_run: CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size) 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) + new_calendars.clear() calendar_dates.clear() trips = []