118 lines
4.8 KiB
Python
118 lines
4.8 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 NOT NULL,
|
||
days TEXT NOT NULL,
|
||
route TEXT NOT NULL,
|
||
UNIQUE(train_number, route)
|
||
);
|
||
''')
|
||
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(entries: List[Dict]):
|
||
"""
|
||
Зберігає повний розклад без видалення попередніх записів.
|
||
entries — список словників з полями:
|
||
'train_number', 'days', 'route', 'times'.
|
||
"""
|
||
today = date.today().isoformat()
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
train_ids = []
|
||
# Додаємо або оновлюємо поїзди
|
||
for e in entries:
|
||
tn = e['train_number']
|
||
days = e['days']
|
||
route = e.get('route', '')
|
||
con.execute('''
|
||
INSERT INTO trains (train_number, days, route)
|
||
VALUES (?, ?, ?)
|
||
ON CONFLICT(train_number, route) DO UPDATE SET days = excluded.days
|
||
''', (tn, days, route))
|
||
tid = con.execute(
|
||
'SELECT id FROM trains WHERE train_number = ? AND route = ?',
|
||
(tn, route)
|
||
).fetchone()[0]
|
||
train_ids.append(tid)
|
||
|
||
# Додаємо або оновлюємо станції та розклад
|
||
for idx, e in enumerate(entries):
|
||
tid = train_ids[idx]
|
||
for t in e['times']:
|
||
st = t['station']
|
||
km = t.get('km')
|
||
con.execute('''
|
||
INSERT INTO stations (name, km)
|
||
VALUES (?, ?)
|
||
ON CONFLICT(name) DO UPDATE SET km = COALESCE(excluded.km, stations.km)
|
||
''', (st, km))
|
||
sid = 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, sid, 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()]
|