def test_tubeutils(self): """ Unit test for helper functions in WhensMyTrain """ self.assertEqual(get_line_code('Central'), 'C') self.assertEqual(get_line_code('Circle'), 'O') self.assertEqual(get_line_name('C'), 'Central') self.assertEqual(get_line_name('O'), 'Circle') for (line_code, line_name) in LINE_NAMES.keys(): self.assertEqual(line_name, get_line_name(get_line_code(line_name))) self.assertEqual(line_code, get_line_code(get_line_name(line_code)))
def import_network_data_to_graph(): """ Import data from a file describing the edges of the Tube network and turn it into a graph object which we pickle and save """ database = WMTDatabase("whensmytrain.geodata.db") # Adapted from https://github.com/smly/hubigraph/blob/fa23adc07c87dd2a310a20d04f428f819d43cbdb/test/LondonUnderground.txt # which is a CSV of all edges in the network reader = csv.reader(open('./sourcedata/tube-connections.csv')) reader.next() # First we organise our data so that each station knows which lines it is on, and which stations it connects to stations_neighbours = {} interchanges_by_foot = [] for (station1, station2, line) in reader: if line in ("National Rail", "East London"): continue if line == "Walk": interchanges_by_foot.append((station1, station2)) else: # When a line splits into two branches, we don't want people being able to travel from one branch to another without # changing. So for these special cases, we mark the transitions as being in a particular direction in the CSV, with the # direction coming after a colon (e.g. "Leytonstone:Northbound","Wanstead","Central" and "Snaresbrook","Leytonstone:Southbound","Central" # Effectively the Central Line station has become two nodes, and now you cannot go directly from Snaresbrook to Wanstead. direction = station1.partition(':')[2] # Blank for most station1 = station1.partition(':')[0] # So station name becomes just e.g. Leytonstone station_data = stations_neighbours.get(station1, []) if (station2, direction, line) not in station_data: station_data += [(station2, direction, line)] stations_neighbours[station1] = station_data # Sanity-check our data and make sure it matches database canonical_data = database.get_rows("SELECT * FROM locations") canonical_station_names = unique_values([canonical['name'] for canonical in canonical_data]) for station in sorted(stations_neighbours.keys()): if station not in canonical_station_names: print "Error! %s is not in the canonical database of station names" % station for (neighbour, direction, line) in stations_neighbours[station]: line_code = get_line_code(line) if not database.get_value("SELECT name FROM locations WHERE name=? AND line=?", (station, line_code)): print "Error! %s is mistakenly labelled as being on the %s line in list of nodes" % (station, line) for station in sorted(canonical_station_names): if station not in stations_neighbours.keys(): print "Error! %s is not in the list of station nodes" % station continue database_lines = database.get_rows("SELECT line FROM locations WHERE name=?", (station,)) for row in database_lines: if row['line'] not in [get_line_code(line) for (neighbour, direction, line) in stations_neighbours[station]]: print "Error! %s is not shown as being on the %s line in the list of nodes" % (station, row['line']) # Produce versions of the graphs for unique lines graphs = {} lines = unique_values([line for station in stations_neighbours.values() for (neighbour, direction, line) in station]) for line in lines: this_line_only = {} for (station_name, neighbours) in stations_neighbours.items(): neighbours_for_this_line = [neighbour for neighbour in neighbours if neighbour[2] == line] if neighbours_for_this_line: this_line_only[station_name] = neighbours_for_this_line graphs[get_line_code(line)] = create_graph_from_dict(this_line_only, database, interchanges_by_foot) graphs['All'] = create_graph_from_dict(stations_neighbours, database, interchanges_by_foot) pickle.dump(graphs, open("./db/whensmytrain.network.gr", "w"))