crystal-orb/crystal_orb_cli.py

#!/usr/bin/python3
import dataclasses
import sys
from datetime import datetime

from crystal_orb import get_departures
from crystal_orb import STOP_PLACES

"""
Ideas:
- If time is less than 15 min from now, time text turns red?
    - Maybe far-away departures are faint?
- If realtime, then color green
- Different 

"""

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):
        out_str = ""
        out_str += "\033[2;20;35m" + str(self.departure_time.time()) + "\033[0m"
        out_str += ": "
        out_str += pad(str(self.line), 4, justification="r")
        out_str += "  "
        out_str += "\033[3;37m" + pad(self.destination, 35+19) + "\033[0m"
        out_str += " ("
        if self.is_realtime: out_str += "\033[0;32mR\033[0m"
        else: out_str += " "
        
        out_str += ", "
        if self.towards_midtbyen: out_str += "\033[0;32mTo  \033[0m"
        else: out_str += "\033[0;31mFrom\033[0m"
        
        out_str += ")"
        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())

def print_help_info():
    print(
"""Help info for the crystal orb cli.
Possible commands:
<stop place> : query the api for all busses from the given stop place
-h or --help : Show this screen
--list-stop-places : List the admissible stop places,
    This is configured in the STOP_PLACES dict in crystal_orb.py""")

if __name__ == "__main__":
    args = sys.argv
    # If no arguments are passed, we just print the default stop
    if len(args) < 2:
        print_colorized_departures("gløshaugen")
        sys.exit(0)

    if "-h" in args or "--help" in args:
        print_help_info()
        sys.exit(0)

    # Values for control flow and processing on the available departures
    stop_place = "gløshaugen"
    bus_line = None
    towards_city_center = None
    take_one = False

    # Parse input arguments to the program
    for i, arg in enumerate(args):
        if i == 1 and arg[0] != "-":  # First argument is taken as the stop place
            stop_place = arg.lower()
        elif arg == "--next":  # Limits to only giving "the next bus" matching the query
            take_one = True
        elif arg == "-b" or arg == "--bus":  # Limit to specific bus line
            try:
                bus_line = args[i+1]
            except:
                print("Bus line not supplied to -b or --bus.")
        elif arg == "--list-stop-places" or arg == "-l":  # Just print stop places
            print("Admissible stop places:")
            for place in STOP_PLACES.keys():
                print(f"\t{place}")
            sys.exit()
        elif arg == "-t" or arg == "--to":  # Limit to buses going to town
            towards_city_center = True
        elif arg == "-f" or arg == "--from":  # Limit to buses leaving town
            towards_city_center = False
        elif arg == "-s" or arg == "--stop":  # Change stop
            try:
                stop_place = args[i+1].lower()
            except:
                print("Stop place not supplied to -s or --stop")
        else:
            pass

    # Arguments are parsed, we can execute the query
    try:
        deps = get_departure_as_classes(stop_place)
    except:
        print(f"Stop {stop_place} not admissible. Run --list-stop-places for help.")
        sys.exit(0)

    if bus_line is not None:
        deps = list(filter(lambda d: str(d.line) == bus_line, deps))
    if towards_city_center is not None:
        if towards_city_center: 
            deps = list(filter(lambda d: d.towards_midtbyen, deps))
        else: 
            deps = list(filter(lambda d: not d.towards_midtbyen, deps))

    if take_one:
        print(deps[0].colorize())
    else:
        for dep in deps:
            print(dep.colorize())
    
    sys.exit(0)