Mercurial > buscompiler / file comparison
comparison: gtfsc.py
gtfsc.py
- changeset 1
- f9788970fa46
- child 2
- 7378b802ddf8
equal
deleted
inserted
replaced
| 0:659ab465152e | 1:f9788970fa46 |
|---|---|
| 1 #!/usr/bin/env python3 | |
| 2 import io | |
| 3 import sys | |
| 4 import sqlalchemy | |
| 5 import sqlalchemy.orm | |
| 6 from datamodel import * | |
| 7 | |
| 8 ROUTE_TYPES = { | |
| 9 '0': 'tram', | |
| 10 '1': 'subway', | |
| 11 '2': 'rail', | |
| 12 '3': 'bus', | |
| 13 '4': 'ferry', | |
| 14 '5': 'cable-tram', | |
| 15 '6': 'aerial-lift', | |
| 16 '7': 'funicular', | |
| 17 '11': 'trolleybus', | |
| 18 '12': 'monorail', | |
| 19 } | |
| 20 | |
| 21 def read_csv(file): | |
| 22 import csv | |
| 23 reader = csv.reader(file) | |
| 24 keys = next(reader) | |
| 25 for i in range(len(keys)): | |
| 26 keys[i] = keys[i].replace('\ufeff', '').strip() | |
| 27 for row in reader: | |
| 28 yield dict(zip(keys, row)) | |
| 29 | |
| 30 def load_gtfs_routes(gtfs_zip): | |
| 31 with gtfs_zip.open('routes.txt') as file: | |
| 32 for row in read_csv(map(bytes.decode, file)): | |
| 33 route = GtfsRoute( | |
| 34 id = row['route_id'], | |
| 35 reference = row['route_short_name'], | |
| 36 description = row['route_long_name'], | |
| 37 type = int(row['route_type']), | |
| 38 ) | |
| 39 yield route.id, route | |
| 40 | |
| 41 def load_shapes(gtfs_zip): | |
| 42 from collections import defaultdict | |
| 43 shapes = dict() | |
| 44 with gtfs_zip.open('shapes.txt') as file: | |
| 45 for row in read_csv(map(bytes.decode, file)): | |
| 46 shape_id = row['shape_id'] | |
| 47 if shape_id not in shapes: | |
| 48 shapes[shape_id] = GtfsShape( | |
| 49 id = shape_id, | |
| 50 shape_coordinates = '', | |
| 51 length = 0, | |
| 52 ) | |
| 53 shape = shapes[shape_id] | |
| 54 if len(shape.shape_coordinates) > 0: | |
| 55 shape.shape_coordinates += ' ' | |
| 56 shape.shape_coordinates += str.format( | |
| 57 '{shape_pt_lat} {shape_pt_lon}', | |
| 58 **row, | |
| 59 ) | |
| 60 shape.length = max(shape.length, float(row['shape_dist_traveled'])) | |
| 61 return shapes.values() | |
| 62 | |
| 63 def trip_length(trip, *, shapes): | |
| 64 if trip.shape_id: | |
| 65 return dict.get(shapes, trip.shape_id).length * float(profile['metrics']['shape-modifier']) | |
| 66 else: | |
| 67 return 0 | |
| 68 | |
| 69 def load_trips(gtfs_zip): | |
| 70 services = set() | |
| 71 with gtfs_zip.open('trips.txt') as file: | |
| 72 for row in read_csv(map(bytes.decode, file)): | |
| 73 if row['service_id'] not in services: | |
| 74 set.add(services, row['service_id']) | |
| 75 yield GtfsService(id = row['service_id']) | |
| 76 yield GtfsTrip( | |
| 77 id = row['trip_id'], | |
| 78 route_id = row['route_id'], | |
| 79 service = row['service_id'], | |
| 80 shape_id = dict.get(row, 'shape_id') | |
| 81 ) | |
| 82 | |
| 83 def load_stops(gtfs_zip): | |
| 84 with gtfs_zip.open('stops.txt') as file: | |
| 85 for row in read_csv(map(bytes.decode, file)): | |
| 86 lat = float(row['stop_lat']) | |
| 87 lon = float(row['stop_lon']) | |
| 88 yield GtfsStop( | |
| 89 stop_id = row['stop_id'], | |
| 90 stop_name = row['stop_name'], | |
| 91 stop_latitude = lat, | |
| 92 stop_longitude = float(row['stop_lon']), | |
| 93 ) | |
| 94 | |
| 95 def gtfs_stop_spatial_testing(session, regions): | |
| 96 print('Finding out in which regions bus stops are...') | |
| 97 from compute_regions import RegionTester | |
| 98 regiontester = RegionTester(regions) | |
| 99 for bus_stop in session.query(GtfsStop): | |
| 100 classification = regiontester( | |
| 101 latitude = bus_stop.stop_latitude, | |
| 102 longitude = bus_stop.stop_longitude, | |
| 103 ) | |
| 104 if classification: | |
| 105 bus_stop.stop_region = classification.region | |
| 106 bus_stop.stop_region_major = classification.region_class == 'major' | |
| 107 | |
| 108 def load_with_loading_text(fn, what, device): | |
| 109 print( | |
| 110 str.format('Loading {}s... ', what), | |
| 111 file = device, | |
| 112 end = '', | |
| 113 flush = True, | |
| 114 ) | |
| 115 result = fn() | |
| 116 print( | |
| 117 str.format( | |
| 118 '{n} {what}s', | |
| 119 n = len(result if type(result) is not tuple else result[0]), | |
| 120 what = what, | |
| 121 ), | |
| 122 file = device, | |
| 123 ) | |
| 124 return result | |
| 125 | |
| 126 def load_gtfs( | |
| 127 gtfs_zip_path, | |
| 128 *, | |
| 129 profile, | |
| 130 session, | |
| 131 device = sys.stderr | |
| 132 ): | |
| 133 from zipfile import ZipFile | |
| 134 with ZipFile(gtfs_zip_path) as gtfs_zip: | |
| 135 print('Loading routes...') | |
| 136 for route_id, route in load_gtfs_routes(gtfs_zip): | |
| 137 session.add(route) | |
| 138 print('Loading stops...') | |
| 139 for stop in load_stops(gtfs_zip): | |
| 140 session.add(stop) | |
| 141 print('Loading shapes...') | |
| 142 for shape in load_shapes(gtfs_zip): | |
| 143 session.add(shape) | |
| 144 print('Loading trips...') | |
| 145 for trip_or_service in load_trips(gtfs_zip): | |
| 146 session.add(trip_or_service) | |
| 147 | |
| 148 def parse_yesno(value): | |
| 149 return value and value != 'no' | |
| 150 | |
| 151 def regions_to_db(regions): | |
| 152 from itertools import product | |
| 153 for region in regions.values(): | |
| 154 names = dict() | |
| 155 for prefix, language in product( | |
| 156 ['', 'short_', 'internal_'], | |
| 157 ['', ':sv', ':en', ':ja'], | |
| 158 ): | |
| 159 key = 'region_' + prefix + 'name' + str.replace(language, ':', '_') | |
| 160 value = dict.get(region, prefix + 'name' + language) | |
| 161 names[key] = value | |
| 162 yield GtfsRegion( | |
| 163 **names, | |
| 164 municipality = dict.get(region, 'municipality'), | |
| 165 external = parse_yesno(dict.get(region, 'external')), | |
| 166 ) | |
| 167 | |
| 168 if __name__ == '__main__': | |
| 169 import sys | |
| 170 from configparser import ConfigParser | |
| 171 from regions import parse_regions | |
| 172 profile = ConfigParser() | |
| 173 profile.read('föli.ini') | |
| 174 engine = sqlalchemy.create_engine('sqlite:///gtfs.db') | |
| 175 GtfsBase.metadata.create_all(engine) | |
| 176 session = sqlalchemy.orm.sessionmaker(bind = engine)() | |
| 177 regions = parse_regions('föli.osm') | |
| 178 for region in regions_to_db(regions): | |
| 179 session.add(region) | |
| 180 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() |
