131 lines
5.3 KiB
Python
131 lines
5.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,
|
||
route TEXT
|
||
);
|
||
''')
|
||
# Таблиця станцій
|
||
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]):
|
||
"""
|
||
Зберігає повний розклад для заданого напряму.
|
||
direction не використовується напряму — маршрут береться з поля 'route' у entries.
|
||
entries — список словників з полями:
|
||
{
|
||
'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:
|
||
# Видаляємо старі записи за сьогоднішню дату
|
||
con.execute('DELETE FROM schedules WHERE travel_date = ?', (today,))
|
||
|
||
for e in entries:
|
||
train_number = e['train_number']
|
||
days = e['days']
|
||
# Беріть маршрут із поля entry['route']
|
||
route = e.get('route', '')
|
||
# Додаємо або оновлюємо поїзд
|
||
con.execute('''
|
||
INSERT INTO trains (train_number, days, route)
|
||
VALUES (?, ?, ?)
|
||
ON CONFLICT(train_number) DO UPDATE SET
|
||
days = excluded.days,
|
||
route = excluded.route
|
||
''', (train_number, days, route))
|
||
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 = None, travel_date: Optional[str] = None) -> List[Dict]:
|
||
"""
|
||
Повертає розклад поїздів за датою.
|
||
direction поки ігнорується; можна додати фільтрацію за route.
|
||
"""
|
||
travel_date = travel_date or date.today().isoformat()
|
||
with sqlite3.connect(DB_PATH) as con:
|
||
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, route, station, arrival, departure in rows:
|
||
schedule.setdefault((num, route), []).append({
|
||
'station': station,
|
||
'arrival': arrival,
|
||
'departure': departure
|
||
})
|
||
return [
|
||
{'train_number': num, 'route': route, 'times': times}
|
||
for (num, route), times in schedule.items()
|
||
]
|