diff --git a/db.py b/db.py index 00e950b..6a373c8 100644 --- a/db.py +++ b/db.py @@ -7,13 +7,16 @@ 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) + days TEXT NOT NULL, + route TEXT ); ''') + # Таблиця станцій con.execute(''' CREATE TABLE IF NOT EXISTS stations ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -21,6 +24,7 @@ def init_db(): km REAL ); ''') + # Таблиця розкладу: зв'язок поїзда зі станцією та часи con.execute(''' CREATE TABLE IF NOT EXISTS schedules ( train_id INTEGER NOT NULL, @@ -36,8 +40,21 @@ def init_db(): ''') 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: # Видаляємо старі записи за сьогоднішню дату @@ -46,19 +63,23 @@ def save_schedule(direction: str, entries: List[Dict]): 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) - VALUES (?, ?) - ON CONFLICT(train_number) DO UPDATE SET days = excluded.days - ''', (train_number, days)) + 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 + km = t.get('km') # може бути None # Додаємо або оновлюємо станцію con.execute(''' INSERT INTO stations (name, km) @@ -79,12 +100,16 @@ def save_schedule(direction: str, entries: List[Dict]): ''', (train_id, station_id, arrival, departure, today)) con.commit() -def get_schedule(direction: str, travel_date: Optional[str] = None) -> List[Dict]: - """Повертає розклад поїздів за датою.""" + +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, st.name, sc.arrival_time, sc.departure_time + 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 @@ -92,11 +117,14 @@ def get_schedule(direction: str, travel_date: Optional[str] = None) -> List[Dict 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({ + 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, 'times': times} for num, times in schedule.items()] + return [ + {'train_number': num, 'route': route, 'times': times} + for (num, route), times in schedule.items() + ] diff --git a/parser.py b/parser.py index 5fb5d91..c14608a 100644 --- a/parser.py +++ b/parser.py @@ -13,12 +13,11 @@ DAY_INDEX = { } def parse_days(text: str) -> str: - """Перетворює текст днів курсування на бінарний рядок, довжиною 7.""" + """Перетворює текст днів курсування на бінарний рядок довжиною 7.""" text = text.lower().strip() if 'щоденно' in text: return '1111111' if text.startswith('крім'): - # виключні дні days_part = text.split('крім', 1)[1] days = [d.strip(' .') for d in days_part.split(',')] mask = [1] * 7 @@ -28,7 +27,6 @@ def parse_days(text: str) -> str: mask[idx] = 0 return ''.join(str(b) for b in mask) if text.startswith('по'): - # тільки вказані дні days_part = text.split('по', 1)[1] days = [d.strip(' .') for d in days_part.split(',')] mask = [0] * 7 @@ -37,25 +35,14 @@ def parse_days(text: str) -> str: if idx is not None: mask[idx] = 1 return ''.join(str(b) for b in mask) - # за замовчуванням — щоденно return '1111111' def fetch_schedule(use_local: bool = False) -> List[Dict]: - """Повертає список поїздів зі станціями та часами: - [ - { - 'train_number': '6902', - 'days': '1111111', - 'times': [ - {'station': 'Київ-Волинський', 'arrival': '', 'departure': '05:10'}, - ... 35 записів ... - ] - }, - ... - ] """ - # Завантажуємо HTML + Повертає список поїздів зі станціями, часами і маршрутом для кожного потяга. + """ + # Завантаження HTML if use_local: with open(LOCAL_HTML, 'r', encoding='utf-8') as f: html = f.read() @@ -66,49 +53,56 @@ def fetch_schedule(use_local: bool = False) -> List[Dict]: soup = BeautifulSoup(html, 'html.parser') - # Список станцій + # Таблиця з часами руху + times_table = soup.select_one('div#tabs-trains1 table.td_center') + if not times_table: + raise RuntimeError('Не знайдено таблицю розкладу') + + # Маршрути: кожен відповідає одному потягу + route_tags = times_table.select('td.course') + routes = [tag.get_text(strip=True) for tag in route_tags] + + # Список станцій (35) station_tags = soup.select( - 'div#tabs-trains1 table.left tr.on a.et, div#tabs-trains1 table.left tr.onx a.et' + 'div#tabs-trains1 table.left tr.on a.et, ' + 'div#tabs-trains1 table.left tr.onx a.et' ) stations = [a.get_text(strip=True) for a in station_tags] - # Таблиця з часами - times_table = soup.select_one('div#tabs-trains1 table.td_center') - trs = times_table.find_all('tr') - - # Рядок з заголовками поїздів - header_row = next(r for r in trs if r.find('td', class_='on_right_t')) - tds = header_row.find_all('td', class_='on_right_t') + # Заголовок з номерами потягів + rows = times_table.find_all('tr') + header_row = next(r for r in rows if r.find('td', class_='on_right_t')) + train_cells = header_row.find_all('td', class_='on_right_t') entries: List[Dict] = [] - for td in tds: - text = td.get_text(separator='|', strip=True) + for idx, cell in enumerate(train_cells): + text = cell.get_text(separator='|', strip=True) parts = text.split('|') num = parts[0].rstrip(',').strip() days_text = parts[1].strip() if len(parts) > 1 else 'щоденно' days = parse_days(days_text) + route = routes[idx] if idx < len(routes) else '' entries.append({ 'train_number': num, 'days': days, + 'route': route, 'times': [] }) - # Рядки з часами руху (35 рядків) - time_rows = [r for r in trs if r.find('td', class_='q0') or r.find('td', class_='q1')] + # Рядки з часами руху + time_rows = [r for r in rows if r.find('td', class_='q0') or r.find('td', class_='q1')] - # Заповнюємо часи для кожного поїзда та станції + # Збирання часу для кожного поїзда та станції for idx, entry in enumerate(entries): base = idx * 3 - times_list = [] for si, row in enumerate(time_rows): cells = row.find_all('td') arrival = cells[base + 1].get_text(strip=True) if base + 1 < len(cells) else '' departure = cells[base + 2].get_text(strip=True) if base + 2 < len(cells) else '' - times_list.append({ + entry['times'].append({ 'station': stations[si], 'arrival': arrival, 'departure': departure }) - entry['times'] = times_list return entries