ttcli.py

#                                        Problem Description
# ------------------------------------------------------------------------------------------------------
# Every day as I leave for school, I check the expected arrival times for 3 stops near my house.
# One way to check is Google Maps, but this requires zooming around a map looking for a little stop
# icon to tap on. Another way to check is to text the stop's number to [898882](tel:+1898882), but
# this has its own drawbacks. The output is slow and visually indistinct, making it hard to parse as
# I'm running for the bus. It also means I have to remember the stop numbers, which by now I know like
# the back of my hand, but remembering numbers (and using SMS) is very 20th century.
#
# A solution for the 21st century exists, though! The TTC doesn't seem to talk about this anywhere
# I can find, but transit agency consultant UmoIQ provides a public XML feed that (I believe)
# backs the TTC's SMS predictions service. The documentation can be found
# [here](https://retro.umoiq.com/xmlFeedDocs/NextBusXMLFeed.pdf). The aim of this project is to
# provide a simple CLI interface for this XML feed, as well as a config file to save & load
# stop numbers, making checking bus times in the morning as simple as one command.
# ------------------------------------------------------------------------------------------------------
# all of these are from the python standard library, so I'd expect they're allowed for this project
import argparse, os, sys # argparse for building the CLI, os for expanding '~' as user's home directory
from typing import Callable, Optional # I love type signatures so much. It pains me that python doesn't enforce them all the time
import urllib.request as request
import xml.parsers.expat as xml

# a pre-defined query from the user's config file
class Query:
    name: str # name of the query (prefixed with a '.' in the config file, specified with '-q' when running the 'query' command)
    agency: str # agency identifier, from list found with 'agencies' command
    stops: list[str] # stopIds can be non-numeric, therefore a string has to be used to represent them.

    def __init__(self, name, agency, stops):
        self.name = name
        self.agency = agency
        self.stops = stops

# a config file holds named queries that the user can execute with the query subcommand
class Config:
    queries: list[Query]

    def __init__(self, file_path: Optional[str] = None):
        # the default spot for the config file is in the user's config directory, usually /home/{user}/.config
        file_path = os.path.expanduser("~/.config/bus_times.conf") if file_path == None else file_path

        try:
            file = open(file_path)

            self.queries = [] # empty for now, will be filled as config file is parsed

            line = file.readline()
            while line != "":
                # this while loop looks through each line of the config file,
                # it follows this pattern:
                # A query begins with a name, prefix with '.'
                # then the next line contains an agency identifier
                # the following lines (until a new name, again prefixed with '.') are stop IDs
                # This can repeat as many times as the user likes
                # for example:
                # ```test.conf
                # # this is a test config file
                #
                # .test
                # ttc
                # 3197
                # ```
                # would be a valid config file that defines a 'test' query, which gets the predictions for stop 3197 from the TTC.
                # Comments can be added with "#" at the start of the line (nowhere else, though), and blank lines are skipped.

                if line[0] == "#" or line == "\n":
                    line = file.readline()
                    continue
                if line[0] == ".":
                    self.queries.append(Query(line[1:-1], "", [])) # new lines need to be removed from the end
                elif self.queries[-1].agency == "":
                    self.queries[-1].agency = line[:-1]
                else:
                    self.queries[-1].stops.append(line[:-1])

                line = file.readline()
        except:
            print(f"Failed to open config file: {file_path}")
            sys.exit(1)

def umoiq_api(
    command: str, # commands are defined on [page 6](https://retro.umoiq.com/xmlFeedDocs/NextBusXMLFeed.pdf#page=6), all requests use one.
    agency: Optional[str] = None, # optional because the 'agencies' command doesn't require it
    route: Optional[str] = None, # optional because the 'agencies' and 'routes' commands don't require it
    stop_id: Optional[str] = None, # optional, required only by the 'times' and 'query' commands
) -> str:
    base_url = "https://retro.umoiq.com/service/publicXMLFeed"

    # easy shorthand b/c I got tired of repeating myself, only used here ofc
    fmt = lambda tag, val: (f"&{tag}={val}" if val != None else "")

    req_url = f"{base_url}?command={command}" \
        + fmt("a", agency) \
        + fmt("r", route) \
        + fmt("stopId", stop_id) \
        + "&terse" # terse output is defined [here](https://retro.umoiq.com/xmlFeedDocs/NextBusXMLFeed.pdf#page=9)

    try:
        response = request.urlopen(req_url, timeout = 2.0).read()

        # the XML returned from UmoIQ has some seemingly broken newline characters
        # and python adds some extra formatting that needs to be stripped before
        # expat can parse it properly, hence:
        return str(response).replace("\\n", "\n")[2:-2]
    except:
        print(f"Request to UmoIQ api: \'{req_url}\' failed or timed out.")
        sys.exit(1)

def gen_parser(
    node_name: str, # XML node name that this parser looks for
    format: list[str], # list of the XML node's attributes that the parser will extract

    # Since the parser this function returns must have a return type of None (this is seemingly just how expat works)
    # the parser generated will have to write output to some external variable. Previously this was a global variable
    # that a function defined for each command would modify, but I refactored all of that into this function.
    output: list[tuple]
) -> Callable[[str, dict], None]: # returns a function to be used as expat's StartElementHandler
    def parser(name: str, attrs: dict):
        if name == node_name:
            try:
                output.append(tuple(map(lambda f: attrs[f], format))) # extracts each attribute specified in "format"
            except:
                return # some <stop> nodes don't have any attributes attached, but these are fine to skip as they're duplicates.

    return parser

def gen_predictions_parser(
    output: list[tuple] # same as gen_parser(), this is the variable the generated parser writes to
) -> Callable[[str, dict], None]:
    def parser(name: str, attrs: dict):
        if name == "predictions": # if the current node is a <predictions> node, start a new group/line of predictions
            output.append((attrs['routeTitle'], []))
        # <predictions> nodes contain multiple <prediction> nodes, so we need to append to the newest item
        # (file is parsed in order) instead of the whole list
        elif name == "prediction":
            # output[-1][1] selects the 2nd element of the last tuple, a.k.a the list element of the tuple we just added
            output[-1][1].append(
                int(attrs['seconds']) // 60 # round down seconds to closest minute, here's your arithmetic!
            )

    return parser

def print_response(
    resp: str, # XML returned from umoiq_api()
    node_name: str, format: list[str], # these two parameters are just passed to gen_parser(), see above
    formatter: Callable[[tuple], str], # formats each line that the parser returns
):
    resp_parser = xml.ParserCreate()
    parsed_response = [] # where resp_parser's StartElementHandler will write to

    resp_parser.StartElementHandler = gen_parser(node_name, format, parsed_response)
    resp_parser.Parse(resp)

    for r in parsed_response:
        print(formatter(r))

# same idea as print_response(), except this time the formatting style needs adjustments that
# I couldn't find a succinct way of adding to print_response()
def print_prediction_response(
    # these arguments are passed to umoiq_api() since it may be called multiple times for queries,
    # TODO: Refactor to use predictionsForMultiStop API command
    agency: str, stop_id: str
):
    resp = umoiq_api(
        'predictions',
        agency,
        stop_id = stop_id
    )

    resp_parser = xml.ParserCreate()
    parsed_response = []

    resp_parser.StartElementHandler = gen_predictions_parser(parsed_response)
    resp_parser.Parse(resp)

    def fmtr(r: tuple[str, list[int]]):
        o = f"{r[0]}: "
        for i in r[1]:
            o += f"{i} min" + (", " if i is not r[1][-1] else "") # if this isn't the last entry, add a comma and space

        # if there are no predictions for a given route, the API still returns an empty <predictions> object.
        return o + ("No predictions available." if len(r[1]) == 0 else "")

    for r in parsed_response:
        print(fmtr(r))

def main():
    # Root Parser
    # arg_parser is called on sys.argv, see argparse library documentation for more details, this is a fairly basic and self-explanatory setup
    arg_parser = argparse.ArgumentParser(
        prog = 'bus_times',
        description = 'get bus times from UmoIQ'
    )
    arg_parser.add_argument(
        '-c', '--config',
        help = 'path to config file',
    )
    arg_parser.add_argument(
        '-a', '--agency',
        default = 'ttc'
    )

    # Sub-parsers
    subparsers = arg_parser.add_subparsers(dest = 'cmd')

    ## Agency Parser
    agency_parser = subparsers.add_parser(
        'agencies',
        description = 'list agency codes'
    )

    ## List Parser
    list_parser = subparsers.add_parser(
        'routes',
        description = 'list all routes for a given agency'
    )

    ## Route Parser
    route_parser = subparsers.add_parser(
        'stops',
        description = 'stops for a given route'
    )
    route_parser.add_argument(
        'route',
        help = 'tag/\'stop-ID\' for a given stop',
    )

    ## Predictions Parser
    predictions_parser = subparsers.add_parser(
        'times',
        description = 'departure times for a given stop ID'
    )
    predictions_parser.add_argument(
        'stop_id',
        help = '\'stop-ID\' obtained from the \'routes\' subcommand'
    )

    ## Queries Parser
    queries_parser = subparsers.add_parser(
        'query',
        description = 'run a pre-defined query'
    )
    queries_parser.add_argument(
        '-q',
        help = 'name of pre-defined query, without \'.\' prefix'
    )

    args = arg_parser.parse_args()

    if args.cmd == None: # when no command is specified, return the help statement
        arg_parser.print_help()
        sys.exit(0)

    config = Config(args.config) # get a config, Config.__init__() returns a default in case of missing config
    agency = args.agency if args.agency != None else config.queries[0].agency # -a argument overrides the config's default agency, which is the first query's agency

    # all of these subcommands have their own '-h'/'--help' messages
    if args.cmd == 'agencies': # agencies subcommand, retrieves a list of agencies UmoIQ has info for
        print_response(
            umoiq_api('agencyList'),
            'agency',
            ['tag', 'title'],
            lambda r: f"{r[0]}: {r[1]}"
        )
    elif args.cmd == 'routes': # routes subcommand, retrieves a list of routes for a given agency
        print_response(
            umoiq_api(
                'routeList',
                agency
            ),
            'route',
            ['title'],
            lambda r: f"{r[0]}"
        )
    elif args.cmd == 'stops': # retrieves a list of stops for a given route
        print_response(
            umoiq_api(
                'routeConfig',
                agency,
                route = args.route
            ),
            'stop',
            ['stopId', 'title'],
            lambda r: f"{r[0]}: {r[1]}"
        )
    elif args.cmd == 'times':
        print_prediction_response(agency, args.stop_id)
    elif args.cmd == 'query':
        q = "default" if args.q == None else args.q

        query = Query("", "", [])

        for i in config.queries:
            if i.name == q:
                query = i

        if query.name != "":
            for s in query.stops:
                print_prediction_response(agency, s)
        else:
            print(f"Couldn't find a definition for query \'{args.q}\' in config file.")

    sys.exit(0) # argparse will handle invariants w/r/t the CLI, so this means one of the above is true

if __name__ == "__main__":
    main() # force of habit, srynotsry