NR: Add better filtering, add arrivals capablity, add train joining information

1 files changed, 114 insertions, 33 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="<headcode>")
 
+
+    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"
+            parsed["departure_only"] = location_code in [a["code"] for a in parsed["origins"]]
+
+            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["departure"] = t["departureType"]
-                parsed["time"] = t["departureType"]
+                parsed["time"] = parsed["time"].strftime("%H%M")
 
-            parsed["on_time"] = parsed["scheduled"] == parsed["departure"]
+            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])))