Wed, 29 Jul 2020 23:45:53 +0300
begin work on bus compiler
1 | 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() |