130 lines
5.0 KiB
Python
130 lines
5.0 KiB
Python
import sqlite3
|
||
from datetime import date
|
||
from typing import List, Dict, Optional
|
||
|
||
DB_PATH = 'schedules.db'
|
||
|
||
def init_db():
|
||
"""Створює таблиці trains, stations та schedules, якщо їх ще нема."""
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
con.execute('''
|
||
CREATE TABLE IF NOT EXISTS trains (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
tid TEXT UNIQUE NOT NULL,
|
||
train_number TEXT NOT NULL,
|
||
days TEXT NOT NULL,
|
||
route TEXT NOT NULL
|
||
);
|
||
''')
|
||
con.execute('''
|
||
CREATE TABLE IF NOT EXISTS stations (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT UNIQUE NOT NULL
|
||
);
|
||
''')
|
||
con.execute('''
|
||
CREATE TABLE IF NOT EXISTS schedules (
|
||
train_id INTEGER NOT NULL,
|
||
station_id INTEGER NOT NULL,
|
||
arrival_time TEXT,
|
||
departure_time TEXT,
|
||
travel_date DATE NOT NULL,
|
||
fetched_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
PRIMARY KEY (train_id, station_id, travel_date),
|
||
FOREIGN KEY(train_id) REFERENCES trains(id),
|
||
FOREIGN KEY(station_id) REFERENCES stations(id)
|
||
);
|
||
''')
|
||
con.commit()
|
||
|
||
|
||
def save_schedule(entries: List[Dict]):
|
||
"""
|
||
Зберігає повний розклад.
|
||
entries — список словників:
|
||
{
|
||
'tid': str,
|
||
'train_number': str,
|
||
'days': '1111111',
|
||
'route': str,
|
||
'times': [
|
||
{'station': str, 'arrival': str, 'departure': str}, ...
|
||
]
|
||
}
|
||
"""
|
||
today = date.today().isoformat()
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
train_ids = []
|
||
# Додаємо або оновлюємо trains по tid
|
||
for e in entries:
|
||
tid_val = e['tid']
|
||
tn = e['train_number']
|
||
days = e['days']
|
||
route = e['route']
|
||
con.execute('''
|
||
INSERT INTO trains (tid, train_number, days, route)
|
||
VALUES (?, ?, ?, ?)
|
||
ON CONFLICT(tid) DO UPDATE SET
|
||
train_number = excluded.train_number,
|
||
days = excluded.days,
|
||
route = excluded.route
|
||
''', (tid_val, tn, days, route))
|
||
train_id = con.execute(
|
||
'SELECT id FROM trains WHERE tid = ?',
|
||
(tid_val,)
|
||
).fetchone()[0]
|
||
train_ids.append(train_id)
|
||
|
||
# Вставляємо записи schedules
|
||
for idx, e in enumerate(entries):
|
||
tid = train_ids[idx]
|
||
for t in e['times']:
|
||
st = t['station']
|
||
# Додаємо лише станцію без колонки km
|
||
con.execute('''
|
||
INSERT INTO stations (name)
|
||
VALUES (?)
|
||
ON CONFLICT(name) DO NOTHING
|
||
''', (st,))
|
||
station_id = con.execute(
|
||
'SELECT id FROM stations WHERE name = ?',
|
||
(st,)
|
||
).fetchone()[0]
|
||
arr = t['arrival']
|
||
dep = t['departure']
|
||
con.execute('''
|
||
INSERT OR REPLACE INTO schedules
|
||
(train_id, station_id, arrival_time, departure_time, travel_date)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
''', (tid, station_id, arr, dep, today))
|
||
con.commit()
|
||
|
||
|
||
def get_schedule(route: Optional[str] = None, travel_date: Optional[str] = None) -> List[Dict]:
|
||
"""Повертає розклад поїздів. Якщо вказано route, фільтрує за ним."""
|
||
from datetime import date as _date
|
||
travel_date = travel_date or _date.today().isoformat()
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
if route:
|
||
rows = con.execute('''
|
||
SELECT tr.train_number, tr.route, st.name, sc.arrival_time, sc.departure_time
|
||
FROM schedules sc
|
||
JOIN trains tr ON sc.train_id = tr.id
|
||
JOIN stations st ON sc.station_id = st.id
|
||
WHERE sc.travel_date = ? AND tr.route = ?
|
||
ORDER BY tr.train_number, st.id
|
||
''', (travel_date, route)).fetchall()
|
||
else:
|
||
rows = con.execute('''
|
||
SELECT tr.train_number, tr.route, st.name, sc.arrival_time, sc.departure_time
|
||
FROM schedules sc
|
||
JOIN trains tr ON sc.train_id = tr.id
|
||
JOIN stations st ON sc.station_id = st.id
|
||
WHERE sc.travel_date = ?
|
||
ORDER BY tr.train_number, st.id
|
||
''', (travel_date,)).fetchall()
|
||
schedule: Dict[tuple, List[Dict]] = {}
|
||
for num, rt, station, arrival, departure in rows:
|
||
schedule.setdefault((num, rt), []).append({'station': station, 'arrival': arrival, 'departure': departure})
|
||
return [{'train_number': num, 'route': rt, 'times': times} for (num, rt), times in schedule.items()]
|