vrr-accountability/setup.py

#!/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()