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()
