def import_train(yaml, mode, dijkstra_lock, mysql_lock): train = Train() name = yaml['name'] category, number = name.split(' ') try: with mysql_lock: t, = train.search(name, variant = yaml['variant']) except ValueError: pass else: if mode == 'force' or True: t.delete() else: return with mysql_lock: train.category = TrainCategory.objects.get(name = category) train.number = number train.variant = yaml['variant'] with mysql_lock: train.save() operations = yaml['operations'] for oper in operations: if oper['mode'] == 'interval': start = oper['from'] end = oper['to'] dates = [ end ] while start != end: dates.append(start) start += DZIEN elif oper['mode'] == 'single': dates = [ oper['date'] ] elif oper['mode'] == 'multi': dates = oper['dates'] timetable = oper['timetable'] tts = [] for i in dates: tts.append(TrainTimetable(train = train, date =i)) with mysql_lock: TrainTimetable.objects.bulk_create(tts) tts = train.traintimetable_set.all() ttbd = {} for x in tts: ttbd[x.date] = x stops = [] tt = ttbd[dates[0]] i = 1 prev = None arr_over = False for description in timetable: stop = TrainStop(timetable = tt) try: stop.station = Station.search(description['station'])[0] except IndexError: print(description) raise if 'departure' in description: stop.departure = datetime.combine(tt.date, time_from_yaml(description.get('departure', None))) if 'arrival' in description: stop.arrival = datetime.combine(tt.date, time_from_yaml(description.get('arrival', None))) stop.track = description.get('track', None) stop.platform = description.get('platform', None) stop.order = i if prev is not None: if arr_over or stop.arrival < prev.departure: # stop.arrival_overnight = 1 arr_over = True if arr_over or (stop.arrival is not None and stop.departure is not None and stop.arrival > stop.departure): # stop.departure_overnight = 1 arr_over = True if prev is None: stop.distance = 0.0 else: dijkstra_lock.acquire() stop.distance = prev.distance + Dijkstra.length(prev.station.gskID, stop.station.gskID) dijkstra_lock.release() prev = stop stops.append(stop) i += 1 with mysql_lock: TrainStop.objects.bulk_create(stops) cursor = connection.cursor() for d, tt in ttbd.items(): if d == dates[0]: continue tto = ttbd[dates[0]] ttn = ttbd[d] with mysql_lock: cursor.execute(''' INSERT INTO kolstatapp_trainstop (timetable_id, station_id, arrival, departure, track, platform, distance, "order") SELECT {0}, station_id, arrival + '{2} days', departure + '{2} days', track, platform, distance, "order" FROM kolstatapp_trainstop WHERE timetable_id = {1}'''.format(ttn.pk, tto.pk, (d - dates[0]).days)) with mysql_lock: transaction.commit_unless_managed()