diff --git a/main.py b/main.py index 227f6b0..25f7e0d 100644 --- a/main.py +++ b/main.py @@ -1,19 +1,5 @@ from stcp.api import get_stop_real_times, get_lines, get_line_directions, get_stop_data, get_line_stops - - -def follow_line(line_code: str, direction_code: str) -> None: - """ - Print the current times for a given line - :param line_code: the line to get the current times for - :param direction_code: line direction, usually '0' or '1', can be found using get_line_directions - :return: None - """ - stops = get_line_stops(line_code, direction_code) - for stop in stops: - stop_data = [bus for bus in get_stop_real_times(stop['stop_code']) if bus['line_code'] == line_code] - - print(f'{stop["name"]: <25} {stop_data[0]["time"] if len(stop_data) > 0 else ""}') - +from stcp.util import follow_line, stop_departures # usage examples all_lines = get_lines() @@ -22,4 +8,8 @@ def follow_line(line_code: str, direction_code: str) -> None: stop_data = get_stop_data('TRD1') next_buses = get_stop_real_times('TRD1') -follow_line('903', '0') +next_departures = stop_departures('TRD1') + +times = follow_line('903', '1') +for time in times: + print(f'{time[0]: <25} {time[1]}') diff --git a/stcp/util.py b/stcp/util.py index e678822..f1b79de 100644 --- a/stcp/util.py +++ b/stcp/util.py @@ -1,7 +1,10 @@ -def get_all_stops(): +from typing import Dict, List, Set, Tuple + + +def get_all_stops() -> Set[str]: """ - Returns a list of all STCP stops across all lines. - :return: a list of all STCP stops + Returns a set of all STCP stop codes across all lines. + :return: a set of all STCP stop codes """ from stcp.api import get_lines, get_line_stops, get_line_directions @@ -13,3 +16,44 @@ def get_all_stops(): all_stops.update([stop['stop_code'] for stop in stops]) # TODO there is a stop in Maia called . with code . return all_stops + + +def stop_departures(stop_code: str, use_hash_cache=True) -> Dict[str, List[str]]: + """ + Returns a map of upcoming departures of a stop, grouped by line. + :param use_hash_cache: use a local cache to avoid doing extra requests per invocation + :param stop_code: the code of the stop + :return: a map of upcoming departures of a stop + """ + from stcp.api import get_stop_real_times + + line_departures = {} + for bus in get_stop_real_times(stop_code, use_hash_cache): + line_code = bus['line_code'] + if line_code not in line_departures: + line_departures[line_code] = [] + + line_departures[line_code].append(bus['time']) + + return line_departures + + +def follow_line(line_code: str, direction_code: str) -> List[Tuple[str, str]]: + """ + Gets the current times for a given line + :param line_code: the line to get the current times for + :param direction_code: line direction, usually '0' or '1', can be found using get_line_directions + :return: a list of the current times for a given line + """ + from stcp.api import get_line_stops + + table = [] + + for stop in get_line_stops(line_code, direction_code): + departures = stop_departures(stop['stop_code']) + if line_code in departures: + table.append((stop['name'], departures[line_code][0])) + else: + table.append((stop['name'], None)) + + return table