diff --git a/crystal_orb.py b/crystal_orb.py index b4325e2..a26c32d 100644 --- a/crystal_orb.py +++ b/crystal_orb.py @@ -13,6 +13,9 @@ STOP_PLACES = { "hesthagen": "41620" } +def get_departures(stop_place: str, atb_api_url: str = ATB_API_URL): + return requests.get(atb_api_url + stop_place).json()["departures"] + def format_departure(departure_dict: dict) -> str: """Formats the json supplied by mpolden's atb api to a single line of text quickly summarizing the scheduled bus. @@ -39,14 +42,17 @@ def write_to_socket(s: str, mpv_socket) -> None: if __name__ == "__main__": - input_stop_place = sys.argv[1] + try: + input_stop_place = sys.argv[1] + except: + input_stop_place = "no input was given" + if input_stop_place in STOP_PLACES.keys(): active_stop_place = STOP_PLACES[input_stop_place] else: active_stop_place = STOP_PLACES["gløshaugen"] - r = requests.get(ATB_API_URL + active_stop_place) - departures = r.json()["departures"] + departures = get_departures(active_stop_place) sorted_departures = sorted(departures, key = lambda d: d["scheduledDepartureTime"]) # ensure mpv is running with the ipc-server socket set in MPV_SOCKET