From fd4157acbd44afe7f8f200be89791894b95eb2f7 Mon Sep 17 00:00:00 2001
From: Emmy D'Anello <ynerant@emy.lu>
Date: Sat, 10 Feb 2024 19:54:10 +0100
Subject: [PATCH] Better calendar import

---
 .../management/commands/update_sncf_gtfs.py   | 32 ++++++++-----------
 1 file changed, 13 insertions(+), 19 deletions(-)

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'])