From 8293ade8ecdc83e3c9a785be5953351e441b2b8b Mon Sep 17 00:00:00 2001 From: Evelyn Date: Wed, 21 Dec 2016 15:50:45 +0000 Subject: [PATCH] NR: Add better filtering, add arrivals capablity, add train joining information --- modules/nr.py | 149 ++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 115 insertions(+), 34 deletions(-) diff --git a/modules/nr.py b/modules/nr.py index 95c2a85b..006adfd8 100644 --- a/modules/nr.py +++ b/modules/nr.py @@ -29,11 +29,49 @@ class Module(object): help="Get information for a given headcode/UID/RID (Powered by NRE)", usage="") + + def filter(self, args, defaults): + args = re.findall(r"[^\s,]+", args) + params = {} + + for arg in args: + if ":" in arg: + params[arg.split(":", 1)[0]] = arg.split(":", 1)[1] + elif "=" in arg: + params[arg.split("=", 1)[0]] = arg.split("=", 1)[1] + else: + params[arg] = True + + ret = {k: v[0] for k,v in defaults.items()} + ret["default"] = True + + for k,v in params.items(): + if not k in defaults.keys(): + raise Exception("Not a valid parameter") + if not defaults[k][1](v): + raise Exception("Not a valid value") + ret["default"] = False + ret[k] = v if len(defaults[k]) == 2 else defaults[k][2](v) + return ret + def arrivals(self, event): - colours = [Utils.COLOR_LIGHTBLUE, Utils.COLOR_GREEN, Utils.COLOR_RED] + colours = [Utils.COLOR_LIGHTBLUE, Utils.COLOR_GREEN, Utils.COLOR_RED, Utils.COLOR_CYAN] location_code = event["args_split"][0].upper() - filter = event["args_split"][1] if len(event["args_split"]) > 1 else "" + filter = self.filter(' '.join(event["args_split"][1:]) if len(event["args_split"]) > 1 else "", { + "dest": ('', lambda x: x.isalpha() and len(x)==3), + "origin":('', lambda x: x.isalpha() and len(x)==3), + "inter": ('', lambda x: x.isalpha() and len(x)==3, lambda x: x.upper()), + "toc": ('', lambda x: x.isalpha() and len(x) == 2), + "dedup": (False, lambda x: type(x)==type(True)), + "plat": ('', lambda x: len(x) <= 3), + "type": ("departures", lambda x: x in ["departures", "arrivals", "both"]), + "terminating": (False, lambda x: type(x)==type(True)) + }) + + if filter["inter"] and filter["type"]!="departures": + event["stderr"].write("Filtering by intermediate stations is only supported for departures.") + return token = self.bot.config["nre-api-key"] client = Client(URL) @@ -41,23 +79,33 @@ class Module(object): header_token.TokenValue = token client.set_options(soapheaders=header_token) - method = client.service.GetDepartureBoardByCRS if len(location_code) == 3 else client.service.GetDepartureBoardByTIPLOC - query = method(50, location_code, datetime.now().isoformat().split(".")[0], 240, - client.factory.create("filterList"), "to", '', "PBS", False) - if not "trainServices" in query: - event["stdout"].write("%s (%s): No services for the next 240 minutes" % (query["locationName"], query["crs"])) + nr_filterlist = client.factory.create("filterList") + if filter["inter"]: nr_filterlist.crs.append(filter["inter"]) + + method = client.service.GetArrivalDepartureBoardByCRS if len(location_code) == 3 else client.service.GetArrivalDepartureBoardByTIPLOC + query = method(100, location_code, datetime.now().isoformat().split(".")[0], 120, + nr_filterlist, "to", '', "PBS", False) + if not "trainServices" in query and not "busServices" in query: + event["stdout"].write("%s (%s): No services for the next 120 minutes" % (query["locationName"], query["crs"])) return trains = [] - for t in query["trainServices"][0]: - parsed = { "scheduled" : datetime.strptime(t["std"], "%Y-%m-%dT%H:%M:%S"), - "called" : "atd" in t, + for t in query["trainServices"][0] if "trainServices" in query else [] + query["busServices"][0] if "busServices" in query else []: + parsed = { + "scheduled": datetime.strptime(t["std"] if "std" in t else t["sta"], "%Y-%m-%dT%H:%M:%S"), + "scheduled_type" : "departure" if "std" in t else "arrival", + "scheduled_short": 'd' if "std" in t else "a", + "arrived" : "ata" in t, + "departed": "atd" in t, "rid" : t["rid"], "uid" : t["uid"], "head" : t["trainid"], - "via": '' if not "via" in t["destination"][0][0] else t["destination"][0][0]["via"], - "platform": "?" if not "platform" in t else t["platform"] + "platform": "?" if not "platform" in t else t["platform"], + "toc": t["operatorCode"], + "cancelled" : t["isCancelled"] if "isCancelled" in t else False, + "terminating" : not "std" in t and not "etd" in t and not "atd" in t, + "bus" : t["trainid"]=="0B00" } parsed["destinations"] = [{"name": a["locationName"], "tiploc": a["tiploc"], "crs": a["crs"] if "crs" in a else '', "code": a["crs"] if "crs" @@ -69,51 +117,73 @@ class Module(object): in a else a["tiploc"], "via": a["via"] if "via" in a else ''} for a in t["origin"][0]] - if "etd" in t or "atd" in t: - parsed["departure"] = datetime.strptime(t["etd"] if "etd" in t else t["atd"], "%Y-%m-%dT%H:%M:%S") - parsed["time"] = parsed["departure"].strftime("%H%M") - elif "isCancelled" in t and t["isCancelled"]: - parsed["departure"] = "Cancelled" - parsed["time"] = "Cancelled" - else: - parsed["departure"] = t["departureType"] - parsed["time"] = t["departureType"] + parsed["departure_only"] = location_code in [a["code"] for a in parsed["origins"]] - parsed["on_time"] = parsed["scheduled"] == parsed["departure"] + parsed["arrival"] = datetime.strptime(t["eta"] if "eta" in t else t["ata"], "%Y-%m-%dT%H:%M:%S") if "eta" in t or "ata" in t else None + parsed["departure"] = datetime.strptime(t["etd"] if "etd" in t else t["atd"], "%Y-%m-%dT%H:%M:%S") if "etd" in t or "atd" in t else None + parsed["time"], parsed["timeprefix"] = [a for a in [(parsed["departure"], "d"), (parsed["arrival"], "a"), (parsed["scheduled"], parsed["scheduled_short"] + "s")] if a[0] != None][0] + parsed["datetime"] = parsed["time"] + + if parsed["cancelled"]: + parsed["time"], parsed["timeprefix"], parsed["prediction"] = ("Cancelled", '', False) + else: + parsed["time"] = parsed["time"].strftime("%H%M") + + parsed["on_time"] = parsed["datetime"] == parsed["scheduled"] and not parsed["cancelled"] parsed["status"] = 1 if parsed["on_time"] else 2 - if parsed["called"]: parsed["status"] = 0 + if parsed["departed"]: parsed["status"] = 3 + if parsed["arrived"]: parsed["status"] = 0 trains.append(parsed) for t in trains: t["dest_summary"] = "/".join(["%s%s" %(a["name"], " " + a["via"] if a["via"] else '') for a in t["destinations"]]) + t["origin_summary"] = "/".join(["%s%s" %(a["name"], " " + a["via"] + if a["via"] else '') for a in t["origins"]]) trains = sorted(trains, key=lambda t: t["scheduled"]) trains_filtered = [] - train_dest_plat = [] + train_locs_toc = [] for train in trains: - if (train["destinations"], train["platform"]) in train_dest_plat and not filter: continue - train_dest_plat.append((train["destinations"], train["platform"])) - trains_filtered.append(train) + if not True in [ + (train["destinations"], train["toc"]) in train_locs_toc and (filter["dedup"] or filter["default"]), + filter["dest"] and not filter["dest"].upper() in [a["code"] for a in train["destinations"]], + filter["origin"] and not filter["origin"].upper() in [a["code"] for a in train["origins"]], + filter["toc"] and not filter["toc"].upper() == train["toc"], + filter["plat"] and not filter["plat"] == train["platform"], + filter["type"] == "departures" and train["terminating"], + filter["type"] == "arrivals" and train["departure_only"], + filter["terminating"] and not train["terminating"] + ]: + train_locs_toc.append((train["destinations"], train["toc"])) + trains_filtered.append(train) - trains_string = ", ".join(["%s (%s, %s, %s%s%s)" % (t["dest_summary"], t["uid"], t["platform"], + trains_string = ", ".join(["%s%s (%s, %s, %s%s%s%s)" % ( + "from " if not filter["type"] in ["arrivals", "departures"] and t["terminating"] else '', + t["origin_summary"] if t["terminating"] or filter["type"]=="arrivals" else t["dest_summary"], + t["uid"], "bus" if t["bus"] else t["platform"], + t["timeprefix"] if t["timeprefix"]!=filter["type"][0] else '', Utils.color(colours[t["status"]]), t["time"], Utils.color(Utils.FONT_RESET) ) for t in trains_filtered]) - event["stdout"].write("%s (%s): %s" % (query["locationName"], query["crs"], + event["stdout"].write("%s (%s)%s: %s" % (query["locationName"], query["crs"], + " departures calling at %s" % filter["inter"] if filter["inter"] else '', trains_string)) def service(self, event): colours = [Utils.COLOR_LIGHTBLUE, Utils.COLOR_GREEN, Utils.COLOR_RED, Utils.COLOR_CYAN] service_id = event["args_split"][0] - filter = event["args_split"][1] if len(event["args_split"]) > 1 else "" + + filter = self.filter(' '.join(event["args_split"][1:]) if len(event["args_split"]) > 1 else "", { + "passing": (False, lambda x: type(x)==type(True)) + }) token = self.bot.config["nre-api-key"] client = Client(URL) @@ -172,6 +242,7 @@ class Module(object): parsed["associations"] = {a["category"] : a for a in station["associations"][0]} if "associations" in station else {} parsed["divides"] = "divide" in parsed["associations"].keys() + parsed["joins"] = "join" in parsed["associations"].keys() if parsed["divides"]: divide = parsed["associations"]["divide"] parsed["divide_summary"] = "%sDividing as %s to %s (%s)%s at " % ( @@ -180,6 +251,14 @@ class Module(object): divide["destCRS"] if "destCRS" in divide else divide["destTiploc"], Utils.color(Utils.FONT_RESET) ) + if parsed["joins"]: + divide = parsed["associations"]["join"] + parsed["divide_summary"] = "%sJoining %s from %s (%s)%s at " % ( + Utils.color(Utils.FONT_BOLD), + divide["uid"], divide["origin"], + divide["originCRS"] if "originCRS" in divide else divide["originTiploc"], + Utils.color(Utils.FONT_RESET) + ) stations.append(parsed) @@ -205,18 +284,20 @@ class Module(object): stations_filtered = [] for station in stations: - if station["passing"] and filter != "*": continue - if station["called"] and filter != "*": + if station["passing"] and not filter["passing"]: continue + if station["called"] and filter["default"]: if not station["first"] and not station["last"]: continue + stations_filtered.append(station) - if station["first"] and not station["last"]: + if station["first"] and not station["last"] and filter["default"]: stations_filtered.append({"summary": "(...)"}) done_count = len([s for s in stations if s["called"]]) total_count = len(stations) - event["stdout"].write("%s%s train (%s%s%s/%s/%s): %s" % (disruptions, query["operator"], + event["stdout"].write("%s%s %s (%s%s%s/%s/%s): %s" % (disruptions, query["operator"], + query["serviceType"], Utils.color(Utils.COLOR_LIGHTBLUE), done_count, Utils.color(Utils.FONT_RESET), len(stations_filtered), total_count, ", ".join([s["summary"] for s in stations_filtered])))