diff -r f9788970fa46 -r 7378b802ddf8 gtfsc.py --- a/gtfsc.py Wed Jul 29 23:45:53 2020 +0300 +++ b/gtfsc.py Thu Jul 30 21:52:31 2020 +0300 @@ -3,6 +3,7 @@ import sys import sqlalchemy import sqlalchemy.orm +import datetime from datamodel import * ROUTE_TYPES = { @@ -92,6 +93,22 @@ stop_longitude = float(row['stop_lon']), ) +def parse_time(timetext): + hour, minute, second = map(int, timetext.split(':')) + return datetime.timedelta(hours = hour, minutes = minute, seconds = second) + +def load_stop_times(gtfs_zip): + with gtfs_zip.open('stop_times.txt') as file: + for row in read_csv(map(bytes.decode, file)): + yield GtfsStopTime( + trip_id = row['trip_id'], + stop_id = row['stop_id'], + arrival_time = parse_time(row['arrival_time']), + departure_time = parse_time(row['departure_time']), + stop_sequence = int(row['stop_sequence']), + shape_distance_traveled = float(row['shape_dist_traveled']), + ) + def gtfs_stop_spatial_testing(session, regions): print('Finding out in which regions bus stops are...') from compute_regions import RegionTester @@ -138,12 +155,22 @@ print('Loading stops...') for stop in load_stops(gtfs_zip): session.add(stop) + session.commit() print('Loading shapes...') for shape in load_shapes(gtfs_zip): session.add(shape) + session.commit() print('Loading trips...') for trip_or_service in load_trips(gtfs_zip): session.add(trip_or_service) + session.commit() + print('Loading stop times...') + for i, stop_time in enumerate(load_stop_times(gtfs_zip)): + if i & 0xffff == 0: + # commit every now and then to keep RAM usage under control + session.commit() + session.add(stop_time) + session.commit() def parse_yesno(value): return value and value != 'no' @@ -154,31 +181,40 @@ names = dict() for prefix, language in product( ['', 'short_', 'internal_'], - ['', ':sv', ':en', ':ja'], + ['fi', 'sv', 'en', 'ja'], ): - key = 'region_' + prefix + 'name' + str.replace(language, ':', '_') - value = dict.get(region, prefix + 'name' + language) + key = 'region_' + prefix + 'name_' + language + value = dict.get(region, prefix + 'name:' + language) names[key] = value yield GtfsRegion( **names, + ref = region['ref'], municipality = dict.get(region, 'municipality'), external = parse_yesno(dict.get(region, 'external')), ) +def get_args(): + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('profile') + parser.add_argument('gtfs') + parser.add_argument('--process-only', action = 'store_true') + return parser.parse_args() + if __name__ == '__main__': - import sys from configparser import ConfigParser from regions import parse_regions + args = get_args() profile = ConfigParser() - profile.read('föli.ini') + profile.read(args.profile) engine = sqlalchemy.create_engine('sqlite:///gtfs.db') GtfsBase.metadata.create_all(engine) session = sqlalchemy.orm.sessionmaker(bind = engine)() regions = parse_regions('föli.osm') - for region in regions_to_db(regions): - session.add(region) + if not args.process_only: + for region in regions_to_db(regions): + session.add(region) + session.commit() + buses = load_gtfs(args.gtfs, profile = profile, session = session) + gtfs_stop_spatial_testing(session = session, regions = regions) session.commit() - buses = load_gtfs('gtfs.zip', profile = profile, session = session) - gtfs_stop_spatial_testing(session = session, regions = regions) - print('Committing to database...') - session.commit()