1 files changed, 17 insertions, 9 deletions

diff --git a/db b/db
index 0a87ab2..d66401f 100755
--- a/db
+++ b/db
@@ -19,20 +19,28 @@ elif argv[-1] == "Mettingen":
 
 
 def get(date):
-    potential = []
     headers = {
         "Accept": "application/xml",
         "Authorization": "Bearer " + token,
     }
     url = "https://api.deutschebahn.com/timetables/v1/plan/" + location + "/" + date
-    root = ET.fromstring(requests.get(url, headers=headers).text)
-    for child in root:
-        for c in child:
-            if c.tag == "dp":
-                if "Bad Cannstatt" in str(c.attrib):
-                    potential.append(c.attrib)
-
-    return potential
+
+    timetablestoplists = [
+        timetablestoplist
+        for timetablestoplists in ET.fromstring(requests.get(url, headers=headers).text)
+        for timetablestoplist in timetablestoplists
+    ]
+    departures = [
+        timetablestoplist
+        for timetablestoplist in timetablestoplists
+        if timetablestoplist.tag == "dp"
+    ]
+
+    return [
+        departure.attrib
+        for departure in departures
+        if "Bad Cannstatt" in str(departure.attrib)
+    ]
 
 
 potential = get(datetime.now().strftime("%-y%m%d/%-H"))