103 lines
4.3 KiB
Python
103 lines
4.3 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,
|
||
train_number TEXT UNIQUE NOT NULL,
|
||
days TEXT NOT NULL -- бінарна маска (Mon→Sun)
|
||
);
|
||
''')
|
||
con.execute('''
|
||
CREATE TABLE IF NOT EXISTS stations (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT UNIQUE NOT NULL,
|
||
km REAL
|
||
);
|
||
''')
|
||
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(direction: str, entries: List[Dict]):
|
||
"""Зберігає повний розклад для заданого напряму."""
|
||
today = date.today().isoformat()
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
# Видаляємо старі записи за сьогоднішню дату
|
||
con.execute('DELETE FROM schedules WHERE travel_date = ?', (today,))
|
||
|
||
for e in entries:
|
||
train_number = e['train_number']
|
||
days = e['days']
|
||
# Додаємо або оновлюємо поїзд
|
||
con.execute('''
|
||
INSERT INTO trains (train_number, days)
|
||
VALUES (?, ?)
|
||
ON CONFLICT(train_number) DO UPDATE SET days = excluded.days
|
||
''', (train_number, days))
|
||
train_id = con.execute(
|
||
'SELECT id FROM trains WHERE train_number = ?', (train_number,)
|
||
).fetchone()[0]
|
||
|
||
for t in e['times']:
|
||
station = t['station']
|
||
km = t.get('km') # невідомо, можна None
|
||
# Додаємо або оновлюємо станцію
|
||
con.execute('''
|
||
INSERT INTO stations (name, km)
|
||
VALUES (?, ?)
|
||
ON CONFLICT(name) DO UPDATE SET km = COALESCE(excluded.km, stations.km)
|
||
''', (station, km))
|
||
station_id = con.execute(
|
||
'SELECT id FROM stations WHERE name = ?', (station,)
|
||
).fetchone()[0]
|
||
|
||
arrival = t['arrival']
|
||
departure = t['departure']
|
||
# Вставляємо або замінюємо розклад поїзда на станції
|
||
con.execute('''
|
||
INSERT OR REPLACE INTO schedules
|
||
(train_id, station_id, arrival_time, departure_time, travel_date)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
''', (train_id, station_id, arrival, departure, today))
|
||
con.commit()
|
||
|
||
def get_schedule(direction: str, travel_date: Optional[str] = None) -> List[Dict]:
|
||
"""Повертає розклад поїздів за датою."""
|
||
travel_date = travel_date or date.today().isoformat()
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
rows = con.execute('''
|
||
SELECT tr.train_number, 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[str, List[Dict]] = {}
|
||
for num, station, arrival, departure in rows:
|
||
schedule.setdefault(num, []).append({
|
||
'station': station,
|
||
'arrival': arrival,
|
||
'departure': departure
|
||
})
|
||
return [{'train_number': num, 'times': times} for num, times in schedule.items()]
|