Merge pull request 'rewrite to using argparse' (#2) from dev into main

Reviewed-on: #2

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

crystal_orb_cli.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+import argparse
+
+import dataclasses
+import sys
+from datetime import datetime
+
+from crystal_orb import get_departures
+from crystal_orb import STOP_PLACES
+
+def pad(s, total_length, justification="l"):
+    if len(s) >= total_length:
+        return s[0:total_length]
+    else:
+        surplus = total_length - len(s)
+        out_str = " " * total_length
+        if justification == "l":
+            return s + " "*surplus
+        elif justification == "m":
+            return " "*(surplus//2) + s + " "*(surplus - surplus//2)
+        elif justification == "r":
+            return " "*surplus + s
+        else:
+            raise ValueError(f"Justification was {justification}, which is not l, m or r.")
+
+@dataclasses.dataclass
+class Departure:
+    line: int
+    departure_time: datetime
+    destination: str
+    is_realtime: bool
+    towards_midtbyen: bool
+    
+    def __post_init__(self):
+        self.departure_time = datetime.fromisoformat(self.departure_time)
+
+    def colorize(self, timestamp = None):
+        out_str = "\033[0;37m"
+        if timestamp is not None and self.departure_time < timestamp:
+                out_str += "\033[9m"
+        out_str += "\033[2;35m" + str(self.departure_time.time()) + "\033[22;37m"
+        out_str += ": "
+        out_str += pad(str(self.line), 4, justification="r")
+        out_str += "  "
+        out_str += "\033[3;37m" + pad(self.destination, 35+19) + "\033[23;37m"
+        out_str += " ("
+        if self.is_realtime: out_str += "\033[32mR\033[37m"
+        else: out_str += " "
+        
+        out_str += ", "
+        if self.towards_midtbyen: out_str += "\033[32mTo  \033[37m"
+        else: out_str += "\033[31mFrom\033[37m"
+        
+        out_str += ")"
+        if timestamp is not None and self.departure_time < timestamp:
+            out_str += "\033[0m"
+        return out_str
+
+
+def get_departure_as_classes(stop_place: str):
+    departures = get_departures(STOP_PLACES[stop_place])
+    departures = [
+        Departure(
+            d["line"],
+            d["scheduledDepartureTime"],
+            d["destination"],
+            d["isRealtimeData"],
+            d["isGoingTowardsCentrum"])
+        for d in departures
+    ]
+    return sorted(departures, key = lambda d: d.departure_time)
+
+def print_colorized_departures(stop_place: str):
+    departures = get_departure_as_classes(stop_place)
+    for dep in departures:
+        print(dep.colorize())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    DEFAULT_STOP_PLACE = "gløshaugen"
+
+    # Requiring value
+    parser.add_argument("-b", "--bus", 
+                        help="limit output to only the supplied bus line")
+    parser.add_argument("-s", "--stop", 
+                        help="limit output to only the supplied stop place")
+
+    # Boolean flags
+    parser.add_argument("-l", "--list-stop-places", action="store_true",
+                        help="list admissible stop places")
+    parser.add_argument("-n", "--next", action="store_true",
+                        help="limit output to one response, the next matching the query")
+    parser.add_argument("-t", "--to", action="store_true",
+                        help="show only buses going into town")
+    parser.add_argument("-f", "--from", action="store_true", dest="from_",
+                        help="show only buses going away from town")
+    args = parser.parse_args()
+
+    # list places and exit
+    if args.list_stop_places:
+        for place in STOP_PLACES.keys():
+            print(place)
+        sys.exit(0)
+
+    # use stop input to fetch buses
+    if args.stop is None:
+        deps = get_departure_as_classes(DEFAULT_STOP_PLACE)
+    else:
+        try:
+            deps = get_departure_as_classes(args.stop)
+        except:
+            print(f"Stop {args.stop} not admissible. Run --list-stop-places for help.")
+            sys.exit(0)
+
+    # filter output
+    if args.bus is not None:
+        deps = [d for d in deps if d.line == args.bus]
+    if args.to:
+        deps = [d for d in deps if d.towards_midtbyen]
+    if args.from_:
+        deps = [d for d in deps if not d.towards_midtbyen]
+    current_time = datetime.now()
+    if args.next:
+        deps = list(filter(lambda d: d.departure_time > current_time, deps))
+        print(deps[0].colorize(current_time))
+    else:
+        for dep in deps:
+            print(dep.colorize(current_time))
+
+    sys.exit(0)