1+ """Fetches data from the Nysse GTFS API."""
2+
13import csv
2- from datetime import UTC , datetime
4+ from datetime import UTC , datetime , timedelta
35import logging
46import os
57import pathlib
@@ -39,12 +41,12 @@ async def _fetch_gtfs():
3941 try :
4042 path = _get_dir_path ()
4143 filename = "extended_gtfs_tampere.zip"
42- timestamp = _get_file_modified_time (path + filename )
43- if timestamp != datetime (1970 , 1 , 1 ) or datetime .now ().minute != 0 :
44+ if os .path .isfile (path + filename ) and datetime .now ().minute != 0 :
4445 _LOGGER .debug ("Skipped fetching GTFS data" )
4546 return # Skip fetching if the file exists or it's not the top of the hour
47+ timestamp = _get_file_modified_time (path + filename )
4648
47- _LOGGER .debug ("Fetching GTFS data" )
49+ _LOGGER .debug ("Fetching GTFS data from %s" , GTFS_URL )
4850 timeout = aiohttp .ClientTimeout (total = 30 )
4951 async with (
5052 aiohttp .ClientSession (timeout = timeout ) as session ,
@@ -72,7 +74,14 @@ async def _read_csv_to_db():
7274
7375 # Stops
7476 cursor .execute (
75- "CREATE TABLE IF NOT EXISTS stops (stop_id TEXT PRIMARY KEY, stop_name TEXT, stop_lat TEXT, stop_lon TEXT)"
77+ """
78+ CREATE TABLE IF NOT EXISTS stops (
79+ stop_id TEXT PRIMARY KEY,
80+ stop_name TEXT,
81+ stop_lat TEXT,
82+ stop_lon TEXT
83+ )
84+ """
7685 )
7786 stops = _parse_csv_file (_get_dir_path () + "stops.txt" )
7887 to_db = [
@@ -85,7 +94,15 @@ async def _read_csv_to_db():
8594
8695 # Routes
8796 cursor .execute (
88- "CREATE TABLE IF NOT EXISTS trips (trip_id TEXT PRIMARY KEY, route_id TEXT, service_id TEXT, trip_headsign TEXT, direction_id TEXT)"
97+ """
98+ CREATE TABLE IF NOT EXISTS trips (
99+ trip_id TEXT PRIMARY KEY,
100+ route_id TEXT,
101+ service_id TEXT,
102+ trip_headsign TEXT,
103+ direction_id TEXT
104+ )
105+ """
89106 )
90107 trips = _parse_csv_file (_get_dir_path () + "trips.txt" )
91108 to_db = [
@@ -99,13 +116,30 @@ async def _read_csv_to_db():
99116 for i in trips
100117 ]
101118 cursor .executemany (
102- "INSERT OR REPLACE INTO trips (trip_id, route_id, service_id, trip_headsign, direction_id) VALUES (?, ?, ?, ?, ?)" ,
119+ """
120+ INSERT OR REPLACE INTO trips
121+ (trip_id, route_id, service_id, trip_headsign, direction_id)
122+ VALUES (?, ?, ?, ?, ?)
123+ """ ,
103124 to_db ,
104125 )
105126
106127 # Calendar
107128 cursor .execute (
108- "CREATE TABLE IF NOT EXISTS calendar (service_id TEXT PRIMARY KEY, monday TEXT, tuesday TEXT, wednesday TEXT, thursday TEXT, friday TEXT, saturday TEXT, sunday TEXT)"
129+ """
130+ CREATE TABLE IF NOT EXISTS calendar (
131+ service_id TEXT PRIMARY KEY,
132+ monday TEXT,
133+ tuesday TEXT,
134+ wednesday TEXT,
135+ thursday TEXT,
136+ friday TEXT,
137+ saturday TEXT,
138+ sunday TEXT,
139+ start_date TEXT,
140+ end_date TEXT
141+ )
142+ """
109143 )
110144 calendar = _parse_csv_file (_get_dir_path () + "calendar.txt" )
111145 to_db = [
@@ -118,17 +152,32 @@ async def _read_csv_to_db():
118152 i ["friday" ],
119153 i ["saturday" ],
120154 i ["sunday" ],
155+ i ["start_date" ],
156+ i ["end_date" ],
121157 )
122158 for i in calendar
123159 ]
124160 cursor .executemany (
125- "INSERT OR REPLACE INTO calendar (service_id, monday, tuesday, wednesday, thursday, friday, saturday, sunday) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" ,
161+ """
162+ INSERT OR REPLACE INTO calendar
163+ (service_id, monday, tuesday, wednesday, thursday, friday, saturday, sunday, start_date, end_date)
164+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
165+ """ ,
126166 to_db ,
127167 )
128168
129169 # Stop times
130170 cursor .execute (
131- "CREATE TABLE IF NOT EXISTS stop_times (trip_id TEXT, arrival_time TIME, departure_time TIME, stop_id TEXT, stop_sequence TEXT, PRIMARY KEY(trip_id, arrival_time))"
171+ """
172+ CREATE TABLE IF NOT EXISTS stop_times (
173+ trip_id TEXT,
174+ arrival_time TIME,
175+ departure_time TIME,
176+ stop_id TEXT,
177+ stop_sequence TEXT,
178+ PRIMARY KEY(trip_id, arrival_time)
179+ )
180+ """
132181 )
133182 stop_times = _parse_csv_file (_get_dir_path () + "stop_times.txt" )
134183 to_db = [
@@ -143,7 +192,11 @@ async def _read_csv_to_db():
143192 ]
144193 to_db .sort (key = lambda x : x [2 ]) # Sort by departure_time
145194 cursor .executemany (
146- "INSERT OR REPLACE INTO stop_times (trip_id, arrival_time, departure_time, stop_id, stop_sequence) VALUES (?, ?, ?, ?, ?)" ,
195+ """
196+ INSERT OR REPLACE INTO stop_times
197+ (trip_id, arrival_time, departure_time, stop_id, stop_sequence)
198+ VALUES (?, ?, ?, ?, ?)
199+ """ ,
147200 to_db ,
148201 )
149202
@@ -198,7 +251,15 @@ async def get_route_ids(stop_id):
198251 await _fetch_gtfs ()
199252 conn , cursor = _get_database ()
200253 cursor .execute (
201- "SELECT DISTINCT route_id FROM trips WHERE trip_id IN (SELECT trip_id FROM stop_times WHERE stop_id = ?)" ,
254+ """
255+ SELECT DISTINCT route_id
256+ FROM trips
257+ WHERE trip_id IN (
258+ SELECT trip_id
259+ FROM stop_times
260+ WHERE stop_id = ?
261+ )
262+ """ ,
202263 (stop_id ,),
203264 )
204265 route_ids = [row [0 ] for row in cursor .fetchall ()]
@@ -221,18 +282,38 @@ async def get_stop_times(stop_id, route_ids, amount, from_time):
221282 """
222283 await _fetch_gtfs ()
223284 conn , cursor = _get_database ()
224- today = datetime .now ().strftime ("%Y-%m-%d" )
225- weekday = datetime .strptime (today , "%Y-%m-%d" ).strftime ("%A" ).lower ()
226- if from_time :
227- cursor .execute (
228- f"SELECT route_id, trip_headsign, departure_time FROM stop_times JOIN trips ON stop_times.trip_id = trips.trip_id JOIN calendar ON trips.service_id = calendar.service_id WHERE stop_id = ? AND trips.route_id IN ({ ',' .join (['?' ]* len (route_ids ))} ) AND calendar.{ weekday } = '1' AND departure_time > ? LIMIT ?" ,
229- [stop_id , * route_ids , from_time .strftime ("%H:%M:%S" ), amount ],
230- )
231- else :
285+ today = datetime .now ().strftime ("%Y%m%d" )
286+ weekday = datetime .strptime (today , "%Y%m%d" ).strftime ("%A" ).lower ()
287+ stop_times = []
288+ delta_days = 0
289+ while len (stop_times ) < amount :
232290 cursor .execute (
233- f"SELECT route_id, trip_headsign, departure_time FROM stop_times JOIN trips ON stop_times.trip_id = trips.trip_id JOIN calendar ON trips.service_id = calendar.service_id WHERE stop_id = ? AND trips.route_id IN ({ ',' .join (['?' ]* len (route_ids ))} ) AND calendar.{ weekday } = '1' LIMIT ?" ,
234- [stop_id , * route_ids , amount ],
291+ f"""
292+ SELECT stop_times.trip_id, route_id, trip_headsign, departure_time, { delta_days } as delta_days
293+ FROM stop_times
294+ JOIN trips ON stop_times.trip_id = trips.trip_id
295+ JOIN calendar ON trips.service_id = calendar.service_id
296+ WHERE stop_id = ?
297+ AND trips.route_id IN ({ ',' .join (['?' ]* len (route_ids ))} )
298+ AND calendar.{ weekday } = '1'
299+ AND calendar.start_date < ?
300+ AND departure_time > ?
301+ LIMIT ?
302+ """ ,
303+ [stop_id , * route_ids , today , from_time .strftime ("%H:%M:%S" ), amount ],
235304 )
236- stop_times = cursor .fetchall ()
305+ stop_times += cursor .fetchall ()
306+ if len (stop_times ) >= amount :
307+ break
308+ # If there are no more stop times for today, move to the next day
309+ delta_days += 1
310+ if delta_days == 7 :
311+ _LOGGER .debug (
312+ "Not enough departures found. Consider decreasing the amount of requested departures"
313+ )
314+ break
315+ next_day = datetime .strptime (today , "%Y%m%d" ) + timedelta (days = 1 )
316+ today = next_day .strftime ("%Y%m%d" )
317+ weekday = next_day .strftime ("%A" ).lower ()
237318 conn .close ()
238- return stop_times
319+ return stop_times [: amount ]
0 commit comments