Mercurial > buscompiler / file comparison
comparison: gtfsc.py
gtfsc.py
- changeset 2
- 7378b802ddf8
- parent 1
- f9788970fa46
equal
deleted
inserted
replaced
| 1:f9788970fa46 | 2:7378b802ddf8 |
|---|---|
| 1 #!/usr/bin/env python3 | 1 #!/usr/bin/env python3 |
| 2 import io | 2 import io |
| 3 import sys | 3 import sys |
| 4 import sqlalchemy | 4 import sqlalchemy |
| 5 import sqlalchemy.orm | 5 import sqlalchemy.orm |
| 6 import datetime | |
| 6 from datamodel import * | 7 from datamodel import * |
| 7 | 8 |
| 8 ROUTE_TYPES = { | 9 ROUTE_TYPES = { |
| 9 '0': 'tram', | 10 '0': 'tram', |
| 10 '1': 'subway', | 11 '1': 'subway', |
| 88 yield GtfsStop( | 89 yield GtfsStop( |
| 89 stop_id = row['stop_id'], | 90 stop_id = row['stop_id'], |
| 90 stop_name = row['stop_name'], | 91 stop_name = row['stop_name'], |
| 91 stop_latitude = lat, | 92 stop_latitude = lat, |
| 92 stop_longitude = float(row['stop_lon']), | 93 stop_longitude = float(row['stop_lon']), |
| 94 ) | |
| 95 | |
| 96 def parse_time(timetext): | |
| 97 hour, minute, second = map(int, timetext.split(':')) | |
| 98 return datetime.timedelta(hours = hour, minutes = minute, seconds = second) | |
| 99 | |
| 100 def load_stop_times(gtfs_zip): | |
| 101 with gtfs_zip.open('stop_times.txt') as file: | |
| 102 for row in read_csv(map(bytes.decode, file)): | |
| 103 yield GtfsStopTime( | |
| 104 trip_id = row['trip_id'], | |
| 105 stop_id = row['stop_id'], | |
| 106 arrival_time = parse_time(row['arrival_time']), | |
| 107 departure_time = parse_time(row['departure_time']), | |
| 108 stop_sequence = int(row['stop_sequence']), | |
| 109 shape_distance_traveled = float(row['shape_dist_traveled']), | |
| 93 ) | 110 ) |
| 94 | 111 |
| 95 def gtfs_stop_spatial_testing(session, regions): | 112 def gtfs_stop_spatial_testing(session, regions): |
| 96 print('Finding out in which regions bus stops are...') | 113 print('Finding out in which regions bus stops are...') |
| 97 from compute_regions import RegionTester | 114 from compute_regions import RegionTester |
| 136 for route_id, route in load_gtfs_routes(gtfs_zip): | 153 for route_id, route in load_gtfs_routes(gtfs_zip): |
| 137 session.add(route) | 154 session.add(route) |
| 138 print('Loading stops...') | 155 print('Loading stops...') |
| 139 for stop in load_stops(gtfs_zip): | 156 for stop in load_stops(gtfs_zip): |
| 140 session.add(stop) | 157 session.add(stop) |
| 158 session.commit() | |
| 141 print('Loading shapes...') | 159 print('Loading shapes...') |
| 142 for shape in load_shapes(gtfs_zip): | 160 for shape in load_shapes(gtfs_zip): |
| 143 session.add(shape) | 161 session.add(shape) |
| 162 session.commit() | |
| 144 print('Loading trips...') | 163 print('Loading trips...') |
| 145 for trip_or_service in load_trips(gtfs_zip): | 164 for trip_or_service in load_trips(gtfs_zip): |
| 146 session.add(trip_or_service) | 165 session.add(trip_or_service) |
| 166 session.commit() | |
| 167 print('Loading stop times...') | |
| 168 for i, stop_time in enumerate(load_stop_times(gtfs_zip)): | |
| 169 if i & 0xffff == 0: | |
| 170 # commit every now and then to keep RAM usage under control | |
| 171 session.commit() | |
| 172 session.add(stop_time) | |
| 173 session.commit() | |
| 147 | 174 |
| 148 def parse_yesno(value): | 175 def parse_yesno(value): |
| 149 return value and value != 'no' | 176 return value and value != 'no' |
| 150 | 177 |
| 151 def regions_to_db(regions): | 178 def regions_to_db(regions): |
| 152 from itertools import product | 179 from itertools import product |
| 153 for region in regions.values(): | 180 for region in regions.values(): |
| 154 names = dict() | 181 names = dict() |
| 155 for prefix, language in product( | 182 for prefix, language in product( |
| 156 ['', 'short_', 'internal_'], | 183 ['', 'short_', 'internal_'], |
| 157 ['', ':sv', ':en', ':ja'], | 184 ['fi', 'sv', 'en', 'ja'], |
| 158 ): | 185 ): |
| 159 key = 'region_' + prefix + 'name' + str.replace(language, ':', '_') | 186 key = 'region_' + prefix + 'name_' + language |
| 160 value = dict.get(region, prefix + 'name' + language) | 187 value = dict.get(region, prefix + 'name:' + language) |
| 161 names[key] = value | 188 names[key] = value |
| 162 yield GtfsRegion( | 189 yield GtfsRegion( |
| 163 **names, | 190 **names, |
| 191 ref = region['ref'], | |
| 164 municipality = dict.get(region, 'municipality'), | 192 municipality = dict.get(region, 'municipality'), |
| 165 external = parse_yesno(dict.get(region, 'external')), | 193 external = parse_yesno(dict.get(region, 'external')), |
| 166 ) | 194 ) |
| 167 | 195 |
| 196 def get_args(): | |
| 197 import argparse | |
| 198 parser = argparse.ArgumentParser() | |
| 199 parser.add_argument('profile') | |
| 200 parser.add_argument('gtfs') | |
| 201 parser.add_argument('--process-only', action = 'store_true') | |
| 202 return parser.parse_args() | |
| 203 | |
| 168 if __name__ == '__main__': | 204 if __name__ == '__main__': |
| 169 import sys | |
| 170 from configparser import ConfigParser | 205 from configparser import ConfigParser |
| 171 from regions import parse_regions | 206 from regions import parse_regions |
| 207 args = get_args() | |
| 172 profile = ConfigParser() | 208 profile = ConfigParser() |
| 173 profile.read('föli.ini') | 209 profile.read(args.profile) |
| 174 engine = sqlalchemy.create_engine('sqlite:///gtfs.db') | 210 engine = sqlalchemy.create_engine('sqlite:///gtfs.db') |
| 175 GtfsBase.metadata.create_all(engine) | 211 GtfsBase.metadata.create_all(engine) |
| 176 session = sqlalchemy.orm.sessionmaker(bind = engine)() | 212 session = sqlalchemy.orm.sessionmaker(bind = engine)() |
| 177 regions = parse_regions('föli.osm') | 213 regions = parse_regions('föli.osm') |
| 178 for region in regions_to_db(regions): | 214 if not args.process_only: |
| 179 session.add(region) | 215 for region in regions_to_db(regions): |
| 216 session.add(region) | |
| 217 session.commit() | |
| 218 buses = load_gtfs(args.gtfs, profile = profile, session = session) | |
| 219 gtfs_stop_spatial_testing(session = session, regions = regions) | |
| 180 session.commit() | 220 session.commit() |
| 181 buses = load_gtfs('gtfs.zip', profile = profile, session = session) | |
| 182 gtfs_stop_spatial_testing(session = session, regions = regions) | |
| 183 print('Committing to database...') | |
| 184 session.commit() |
