add utility functions

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]}')

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