Use the new IDFM API to fetch trafic

modules/ratp.py

@@ -6,55 +6,167 @@ import urllib.request
 
 from PIL import Image, ImageDraw, ImageFont
 
-class RATPAPI:
+class IDFMAPI:
 
     fnt_R_path = "./fonts/Parisine-Regular.ttf"
     fnt_RB_path = "./fonts/Parisine-Bold.ttf"
 
-    def __init__(self):
+    lines = {
-        self.baseurl = "https://ratp.p0m.fr/api"
+        "metros": {
+            "1": "C01371",
+            "2": "C01372",
+            "3": "C01373",
+            "4": "C01374",
+            "5": "C01375",
+            "6": "C01376",
+            "7": "C01377",
+            "8": "C01378",
+            "9": "C01379",
+            "10": "C01380",
+            "11": "C01381",
+            "12": "C01382",
+            "13": "C01383",
+            "14": "C01384",
+            "3B": "C01386",
+            "7B": "C01387",
+        },
+        "buses": {
+            "57": "C01094",
+            "125": "C01154",
+            #"131": "C01159",
+            "184": "C01205",
+        },
+        "rers": {
+            "A": "C01742",
+            "B": "C01743",
+            "C": "C01727",
+            "D": "C01728",
+            "E": "C01729"
+        },
+        "tramways": {
+            "T2": "C01390",
+            "T3A": "C01391",
+            "T3B": "C01679",
+            "T7": "C01774",
+            "T9": "C02317",
+        },
+    }
+
+
+
+    def __init__(self, apikey=None):
+        self.baseurl = "https://prim.iledefrance-mobilites.fr/marketplace"
+        self.apikey = apikey or os.environ["TOKEN_IDFM"]
 
         self._cached_file = ".ratp-%s.cache"
+        self.cache_time = 5
 
     def get_weather(self):
+        ret = {}
+
+        for mode in IDFMAPI.lines:
+            ret[mode] = {}
+            for line in IDFMAPI.lines[mode]:
+                ret[mode][line] = self.get_line_weather(mode, line)
+
+        return ret
+
+    def get_line_weather(self, mode, line):
+        cache_file = self._cached_file % ("traffic-" + IDFMAPI.lines[mode][line])
         # Read the mod time
         statinfo = None
         try:
-            statinfo = os.stat(self._cached_file % "traffic")
+            statinfo = os.stat(cache_file)
         except:
             pass
 
-        if statinfo is None or datetime.fromtimestamp(statinfo.st_mtime, tz=timezone.utc) + timedelta(minutes=5) < datetime.now(tz=timezone.utc):
+        if statinfo is None or datetime.fromtimestamp(statinfo.st_mtime, tz=timezone.utc) + timedelta(minutes=self.cache_time) < datetime.now(tz=timezone.utc):
             # Do the request and save it
-            with urllib.request.urlopen(self.baseurl + "/traffic") as f:
+            req = urllib.request.Request(self.baseurl + "/general-message?LineRef=STIF:Line::" + IDFMAPI.lines[mode][line] + ":", headers={'apikey': self.apikey})
-                with open(self._cached_file % "traffic", 'wb') as fd:
+            with urllib.request.urlopen(req) as f:
+                with open(cache_file, 'wb') as fd:
                     fd.write(f.read())
 
         # Retrieve cached data
         res = {}
-        with open(self._cached_file % "traffic") as f:
+        with open(cache_file) as f:
             res = json.load(f)
 
-        return res["result"]
+        return res["Siri"]["ServiceDelivery"]["GeneralMessageDelivery"][0]["InfoMessage"]
+
+    def get_line_icon(mode, line, size, fill="gray"):
+        width = int(size * 1.38) if mode == "buses" or mode == "tramways" else size
+        image = Image.new('RGBA', (width, size), '#fff0')
+        draw = ImageDraw.Draw(image)
+
+        fnt_icon = ImageFont.truetype(IDFMAPI.fnt_RB_path, size-3)
+
+        if mode == "metros":
+            draw.ellipse((0, 0, width, size), fill=fill)
+        elif mode == "tramways":
+            draw.rectangle((0, 0, width, 1), fill=fill)
+            draw.rectangle((0, size-2, width, size), fill=fill)
+        else:
+            draw.rectangle((0, 0, width, size), fill=fill)
+        draw.text((int(width / 2), int(size / 2)), line, fill="white" if (fill == "black" and mode == "tramways") or (fill != "white" and mode != "tramways") else "black", anchor="mm", font=fnt_icon)
+
+        return image
+
 
 
 class RATPWeatherModule:
 
+    def __init__(self):
+        self.major_lines = ["M7", "M5", "M14", "RB", "T3A"]
+
+    def gen_alerts(self):
+        alerts = []
+
+        weather = IDFMAPI().get_weather()
+        for mode in weather:
+            for line in weather[mode]:
+                if mode[0].upper() + line not in self.major_lines:
+                    continue
+
+                def alert_icon(mode, line):
+                    def icon(size=64):
+                        image = Image.new('RGBA', (size, size), '#0000')
+
+                        white = Image.new('RGB', (int(size / 2), int(size / 2)), '#fff')
+                        mode_icon = Image.open("icons/" + mode + ".png").resize((int(size/2), int(size/2)))
+                        image.paste(white, (-5,0), mode_icon)
+
+                        line_icon = IDFMAPI.get_line_icon(mode, line, int(size/2), fill="white")
+                        image.paste(line_icon, (int(size/2) - 5,0), line_icon)
+
+                        return image
+                    return icon
+
+                for info in weather[mode][line]:
+                    if info["InfoChannelRef"]["value"] != "Perturbation":
+                        continue
+
+                    for msg in info["Content"]["Message"]:
+                        if msg["MessageType"] != "TEXT_ONLY":
+                            continue
+
+                        yield {
+                            "description": msg["MessageText"]["value"],
+                            "icon": alert_icon(mode, line),
+                        }
+
     def draw_module(self, config, width, height, line_height=19):
         image = Image.new('RGB', (width, height), '#fff')
         draw = ImageDraw.Draw(image)
 
-        fnt_icon = ImageFont.truetype(RATPAPI.fnt_RB_path, line_height-1)
+        weather = IDFMAPI().get_weather()
-
-        weather = RATPAPI().get_weather()
 
         align_x = 0
         align_y = 0
-        for mode in weather:
+        for mode in ["metros", "rers", "tramways", "buses"]:
-            if mode != "rers":
+            if mode != "rers" and mode != "buses":
                 align_x = 0
 
-
             # display mode icon
             icon = Image.open("icons/" + mode + ".png").resize((line_height, line_height))
             image.paste(icon, (align_x,align_y), icon)
@@ -66,20 +178,18 @@ class RATPWeatherModule:
                     align_x = line_height + 10
                     align_y += line_height + 6
 
-                fill = "gray" if line["slug"] == "normal" else "black"
+                states = []
+                for info in weather[mode][line]:
+                    states.append(info["InfoChannelRef"]["value"])
 
-                if mode == "metros":
+                fill = "gray" if "Perturbation" not in states else "black"
-                    draw.ellipse((align_x - 1, align_y - 1, align_x + line_height + 1, align_y + line_height + 1), fill=fill)
-                elif mode == "tramways":
-                    draw.rectangle((align_x - 2, align_y - 2, align_x + line_height, align_y), fill=fill)
-                    draw.rectangle((align_x - 1, align_y + line_height - 1, align_x + line_height + 1, align_y + line_height + 1), fill=fill)
-                else:
-                    draw.rectangle((align_x - 1, align_y - 1, align_x + line_height + 1, align_y + line_height + 1), fill=fill)
-                draw.text((align_x + 1 + line_height / 2, align_y + line_height / 2), line["line"], fill="white" if mode != "tramways" else "black", anchor="mm", font=fnt_icon)
 
-                align_x += line_height + 5
+                icon = IDFMAPI.get_line_icon(mode, line, line_height, fill=fill)
+                image.paste(icon, (align_x, align_y), icon)
 
-            if mode != "metros":
+                align_x += icon.width + 5
+
+            if mode != "metros" and mode != "tramways":
                 align_y += line_height + 10
             else:
                 align_y += 10