From bb95ff037832ae8317388abf1eb1044e5839397f Mon Sep 17 00:00:00 2001
From: Dmitry Afanasyev <Balshbox@gmail.com>
Date: Mon, 8 Aug 2022 20:16:37 +0300
Subject: [PATCH] add station buttons

---
 main.py    | 15 +++++++++++++--
 mos_gor.py |  8 +++-----
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/main.py b/main.py
index 8c3262e..cada9c5 100644
--- a/main.py
+++ b/main.py
@@ -35,7 +35,12 @@ def get_keyboard() -> types.InlineKeyboardMarkup:
 @dp.callback_query_handler(stations_cb.filter(direction='home->office'))
 async def home_office(query: types.CallbackQuery, callback_data: dict[str, str]) -> None:
 
-    text = parse_site(driver=driver)
+    text = parse_site(
+        driver=driver,
+        url='https://yandex.ru/maps/213/moscow/stops/stop__9640740/'
+            '?l=masstransit&ll=37.527754%2C55.823507&tab=overview&z=21',
+        message='Остановка Б. Академическая ул, д. 15'
+    )
 
     # or reply INTO webhook
     return await query.message.edit_text(text)
@@ -45,7 +50,13 @@ async def home_office(query: types.CallbackQuery, callback_data: dict[str, str])
 async def office_home(query: types.CallbackQuery, callback_data: dict[str, str]) -> Message:
 
     # or reply INTO webhook
-    return await query.message.edit_text('Hello World')
+    text = parse_site(
+        driver=driver,
+        url='https://yandex.ru/maps/213/moscow/stops/stop__9640289/?'
+            'l=masstransit&ll=37.507947%2C55.802094&tab=timetable&z=21',
+        message='Остановка Улица Алабяна'
+    )
+    return await query.message.edit_text(text)
 
 
 @dp.message_handler(commands=['chatid'])
diff --git a/mos_gor.py b/mos_gor.py
index 3ddbd8e..fac15e1 100644
--- a/mos_gor.py
+++ b/mos_gor.py
@@ -47,10 +47,8 @@ def configure_firefox_driver(private_window: bool = False) -> WebDriver:
     return firefox_driver
 
 
-def parse_site(driver: WebDriver) -> str:
-    driver.get(
-        'https://yandex.ru/maps/213/moscow/stops/stop__9640740/?l=masstransit&ll=37.527754%2C55.823507&tab=overview&z=21'
-    )
+def parse_site(driver: WebDriver, url: str, message: str) -> str:
+    driver.get(url)
     time.sleep(4)
     elements = driver.find_elements(by='class name', value='masstransit-vehicle-snippet-view')
 
@@ -68,6 +66,6 @@ def parse_site(driver: WebDriver) -> str:
             bus_t19_arrival = element.find_element(by='class name', value='masstransit-prognoses-view__title-text')
         except NoSuchElementException:
             pass
-    return ('Остановка Б. Академическая ул, д. 15 в сторону Москва Сити\n\n'
+    return (f'{message}\n\n'
             f'Автобус {bus_300.text} - {bus_300_arrival.text}\nАвтобус {bus_t19.text} - {bus_t19_arrival.text}'
             )