diff --git a/modules/nr.py b/modules/nr.py index 1596811a..d0c6e294 100644 --- a/modules/nr.py +++ b/modules/nr.py @@ -23,11 +23,11 @@ class Module(object): usage="") bot.events.on("received").on("command").on("nrservice" ).hook(self.service, min_args=1, - help="Get train service information for a headcode or RID (Powered by NRE)", + help="Get train service information for a UID, headcode or RID (Powered by NRE)", usage="") bot.events.on("received").on("command").on("nrhead" ).hook(self.head, min_args=1, - help="Get information for a given headcode (Powered by NRE)", + help="Get information for a given headcode/UID/RID (Powered by NRE)", usage="") def time_compare(self, one, two): @@ -50,14 +50,16 @@ class Module(object): else: return time_until def arrivals(self, event): - token = self.bot.config["nre-api-key"] + colours = [Utils.COLOR_LIGHTBLUE, Utils.COLOR_GREEN, Utils.COLOR_RED] + location_code = event["args_split"][0].upper() + token = self.bot.config["nre-api-key"] client = Client(URL) - header_token = client.factory.create('ns2:AccessToken') 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], 120, client.factory.create("filterList"), "to", '', "PBS", False) @@ -87,6 +89,9 @@ class Module(object): parsed["on_time"] = parsed["scheduled"] == parsed["departure"] + parsed["status"] = 1 if parsed["on_time"] else 2 + if parsed["called"]: parsed["status"] = 0 + trains.append(parsed) for t in trains: @@ -94,7 +99,6 @@ class Module(object): trains = sorted(trains, key=lambda t: t["scheduled"]) - trains_filtered = [] train_dest_plat = [] @@ -106,7 +110,7 @@ class Module(object): self.result_map[event["target"].id] = trains_filtered trains_string = ", ".join(["%s (%s, %s, %s%s%s)" % (t["dest_via"], t["uid"], t["platform"], - Utils.color(Utils.COLOR_GREEN if t["on_time"] else Utils.COLOR_RED), + Utils.color(colours[t["status"]]), t["time"], Utils.color(Utils.FONT_RESET) ) for t in trains_filtered]) @@ -169,7 +173,7 @@ class Module(object): parsed["time"], parsed["timeprefix"], parsed["prediction"] = ("Cancelled", '', False) else: parsed["time"] = parsed["time"].strftime("%H%M") - parsed["on_time"] = parsed["datetime"] == parsed["scheduled"] + parsed["on_time"] = parsed["datetime"] == parsed["scheduled"] and not parsed["cancelled"] stations.append(parsed)