You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

165 lines
5.5 KiB

#!/usr/bin/env python3.6
import requests
import sys
from typing import List
import configparser
import mysql.connector
HTTP = requests.session()
VRR_STATIONS_TABLE = """
CREATE TABLE IF NOT EXISTS vrr_stations (
station_id int not null primary key,
station_name text);
"""
def yn(s: str) -> bool:
return s in ['y', 'Y', 'J', 'j', '']
def search_station(search: str) -> List or None:
params = {
'query': search.replace(" ", "+")
}
resp = HTTP.get("https://abfahrtsmonitor.vrr.de/backend/api/stations/search", params=params)
try:
resp.raise_for_status()
except requests.exceptions.HTTPError:
return None
return resp.json()['suggestions']
def add_stations_to_db(results: list, db: dict) -> None:
cx = mysql.connector.connect(**db)
cr = cx.cursor()
cr.execute(VRR_STATIONS_TABLE)
for r in results:
cr.execute('REPLACE INTO vrr_stations'
'(station_id, station_name)'
'VALUES (%s, %s)', (r['data'], r['value']))
cx.commit()
cr.close()
cx.close()
def get_station(db: dict) -> int:
station_id = None
while station_id is None:
search = input("Which station would you like to monitor? ")
print("Getting suggestions...")
results = search_station(search)
if results: # empty lists and None are False
add_stations_to_db(results, db)
for i, result in enumerate(results):
print(str(i) + ". " + result['value'] + "\t" + result['data'])
choice_ptr = None
if len(results) > 1:
while choice_ptr is None:
parse = input("Which one? ")
try:
parse = int(parse)
test = results[parse]
except ValueError:
print("You did not input a number!")
except IndexError:
print("Number out of range.")
choice_ptr = parse
else:
choice_ptr = 0
station_id = results[choice_ptr]['data']
else:
print('Got no results for "' + search + '".')
return int(station_id)
def get_lines(station_id: int) -> List[str]:
print("Getting lines...")
resp = HTTP.get("https://abfahrtsmonitor.vrr.de/backend/api/lines/" + str(station_id) + "/search")
resp.raise_for_status()
results = resp.json()
choices = ""
if len(results) > 20:
page = yn(input("There are more than 20 results. Would you like to view them page-by-page? (Y/n)"))
else:
page = False
if page:
print("Paging in batches of 20 results.")
for i, result in enumerate(results):
print(str(i) + ". " + result['name'])
if (i % 20 == 0 or i == len(results) - 1) and i != 0:
choices += input("Please input your choices as a space-separated list (e.g. '0 2 7 15'):\n") + ' '
else:
for i, result in enumerate(results):
print(str(i) + ". " + result['name'])
choices += input("Please input your choices as a space-separated list (e.g. '0 2 7 15'):\n") + ' '
filt_arr = []
for ptr in choices.split(" "):
if ptr == '':
continue
try:
filt_arr.append(results[int(ptr)])
except ValueError:
print(ptr + " is not a number")
except IndexError:
print(ptr + " is out of range")
return ["{0}:{1}:{2}:{3}".format(r['network'], r['line'], r['supplement'], r['directionCode'])
for r in filt_arr]
def config_db() -> dict:
def _enter_details() -> dict:
r = dict()
r['host'] = input("Please enter the database hostname: ")
r['user'] = input("Please enter the database user: ")
r['password'] = input("Please enter the database password: ")
r['database'] = input("Please enter the database name: ")
return r
successful = False
while not successful:
r = _enter_details()
try:
cx = mysql.connector.connect(**r)
except:
print("The database settings seem incorrect. Please try again.")
else:
successful = True
return r
def setup() -> None:
db_config = config_db()
station_id = get_station(db_config)
lines_ch = input("Would you like to choose specific lines? (Y/n)", )
if yn(lines_ch):
lines = get_lines(station_id)
else:
lines = None
cfg = configparser.ConfigParser(allow_no_value=True)
cfg.add_section('crawl')
cfg['crawl']['; VRR_ACCOUNTABILITY CONFIGURATION FILE'] = None
cfg['crawl']['; ====================================='] = None
cfg['crawl']['; You may generate a configuration file using ./setup.py'] = None
cfg['crawl']['station_id'] = str(station_id)
cfg['crawl']['use_long_distance'] = 'yes'
cfg['crawl']['use_regional_trains'] = 'yes'
cfg['crawl']['use_commuter_trains'] = 'yes'
cfg['crawl']['use_underground_trains'] = 'yes'
cfg['crawl']['use_trams'] = 'yes'
cfg['crawl']['use_buses'] = 'yes'
cfg['crawl']['use_elevated_trains'] = 'yes'
cfg['crawl']['use_lines'] = ",".join(lines) if lines is not None else ""
cfg.add_section('db')
cfg['db']['user'] = db_config['user']
cfg['db']['pass'] = db_config['password']
cfg['db']['host'] = db_config['host']
cfg['db']['database'] = db_config['database']
print("Please save the following output to 'vrr.ini' and adjust any further settings:")
print("\n" * 3)
cfg.write(sys.stdout)
if __name__ == '__main__':
setup()