def parse_dlr_data(dlr_data, station): """ Takes a parsed XML elementTree dlr_data and the RailStation object for the station whose departures we are querying Returns a DepartureCollection object of all departures from the station in question, classified by platform """ train_info_regex = re.compile(r"[1-4] (\D+)(([0-9]+) mins?)?", flags=re.I) platforms_to_ignore = [('tog', 'P1'), ('wiq', 'P1')] platforms_to_ignore_if_empty = [('ban', 'P10'), ('str', 'P4B'), ('lew', 'P5')] # Go through each platform and get data about every train arriving, including which direction it's headed trains_by_platform = DepartureCollection() for platform in dlr_data.findall("div[@id='ttbox']"): # Get the platform number from image attached and the time published img = platform.find("div[@id='platformleft']/img") platform_name = img.attrib['src'].split('.')[0][:-1].upper() if (station.code, platform_name) in platforms_to_ignore: continue trains_by_platform[platform_name] = [] # Get trains for this platform info = platform.find("div[@id='platformmiddle']") publication_time = info.find("div[@id='time']").text.strip() publication_time = datetime.strptime(publication_time, "%H:%M") line1 = info.find("div[@id='line1']") line2 = info.find("div[@id='line23']/p") line3 = info.find("div[@id='line23']/p/br") trains = [line for line in (line1.text, line2.text, line3.tail) if line] # Go through trains, parse out the relevant data for train in trains: result = train_info_regex.search(train) if result: destination = capwords(result.group(1).strip()) if destination == 'Terminates Here': continue departure_delta = timedelta(minutes=(result.group(3) and int(result.group(3)) or 0)) departure_time = datetime.strftime(publication_time + departure_delta, "%H%M") trains_by_platform.add_to_slot(platform_name, DLRTrain(destination, departure_time)) logging.debug("Found a train going to %s at %s", destination, departure_time) else: logging.debug("Error - could not parse this line: %s", train) # If there are no trains in this platform to our specified stop, or it is a platform that can be ignored when it is empty # e.g. it is the "spare" platform at a terminus, then delete this platform entirely if not trains_by_platform[platform_name] and (station.code, platform_name) in platforms_to_ignore_if_empty: del trains_by_platform[platform_name] # If two platforms have exact same set of destinations, treat them as one by merging trains_by_platform.merge_common_slots() return trains_by_platform
def test_models(self): """ Unit tests for train, bus, station and bus stop objects """ # Location fundamentals location_name = "Trafalgar Square" location = Location(location_name) self.assertEqual(str(location), location_name) self.assertEqual(repr(location), location_name) self.assertEqual(len(location), len(location_name)) # BusStop fundamentals bus_stop = BusStop("TRAFALGAR SQUARE / CHARING CROSS STATION <> # [DLR] >T<", bus_stop_code='10000', distance=2.0, run=1) bus_stop2 = BusStop("CHARING CROSS STATION <> # [DLR} >T< / TRAFALGAR SQUARE", bus_stop_code='10001', distance=1.0, run=2) self.assertLess(bus_stop2, bus_stop) self.assertEqual(len(bus_stop), 26) self.assertEqual(hash(bus_stop), hash(BusStop("TRAFALGAR SQUARE / CHARING CROSS STATION <> # [DLR] >T<", run=1))) # BusStop complex functions for undesirable in ('<>', '#', r'\[DLR\]', '>T<'): self.assertNotRegexpMatches(bus_stop.get_clean_name(), undesirable) self.assertEqual(bus_stop.get_clean_name(), "Trafalgar Square / Charing Cross Station") self.assertEqual(bus_stop.get_normalised_name(), "TRAFALGARSQCHARINGCROSSSTN") self.assertEqual(bus_stop.get_similarity(bus_stop.name), 100) self.assertEqual(bus_stop.get_similarity("Charing Cross Station"), 94) self.assertEqual(bus_stop2.get_similarity("Charing Cross Station"), 95) self.assertEqual(bus_stop.get_similarity("Charing Cross"), 90) self.assertEqual(bus_stop2.get_similarity("Charing Cross"), 91) # RailStation complex functions station = RailStation("King's Cross St. Pancras", "KXX", 530237, 182944) station2 = RailStation("Earl's Court", "ECT") self.assertEqual(station.get_abbreviated_name(), "Kings X St P") self.assertEqual(station2.get_abbreviated_name(), "Earls Ct") self.assertEqual(station.get_similarity(station.name), 100) self.assertGreaterEqual(station.get_similarity("Kings Cross St Pancras"), 95) self.assertGreaterEqual(station.get_similarity("Kings Cross St Pancreas"), 90) self.assertGreaterEqual(station.get_similarity("Kings Cross"), 90) # Departure departure = Departure("Trafalgar Square", "2359") departure2 = Departure("Trafalgar Square", "0001") self.assertLess(departure, departure2) # Fails if test run at 0000-0059 self.assertEqual(hash(departure), hash(Departure("Trafalgar Square", "2359"))) self.assertEqual(str(departure), "Trafalgar Square 2359") self.assertEqual(departure.get_destination(), "Trafalgar Square") self.assertEqual(departure.get_departure_time(), "2359") # NullDeparture null_departure = NullDeparture("East") self.assertEqual(null_departure.get_destination(), "None shown going East") self.assertEqual(null_departure.get_departure_time(), "") # Bus bus = Bus("Blackwall", "2359") bus2 = Bus("Blackwall", "0001") self.assertLess(bus, bus2) # Fails if test run at 0000-0059 self.assertEqual(bus.get_destination(), "Blackwall") # Train train = Train("Charing Cross via Bank", "2359") train2 = Train("Charing Cross via Bank", "0001") self.assertLess(train, train2) # Fails if test run at 0000-0059 self.assertEqual(train.get_destination(), "Charing Cross via Bank") # TubeTrain tube_train = TubeTrain("Charing Cross via Bank", "Northbound", "2359", "N", "001") tube_train2 = TubeTrain("Charing Cross via Bank then depot", "Northbound", "2359", "N", "001") tube_train3 = TubeTrain("Northern Line", "Northbound", "2359", "N", "001") tube_train4 = TubeTrain("Heathrow T123 + 5", "Westbound", "2359", "P", "001") self.assertEqual(hash(tube_train), hash(tube_train2)) self.assertEqual(tube_train.get_destination(), "Charing Cross via Bank") self.assertEqual(tube_train3.get_destination(), "Northbound Train") self.assertEqual(tube_train4.get_destination(), "Heathrow Terminal 5") self.assertEqual(tube_train.get_destination_no_via(), "Charing Cross") self.assertEqual(tube_train.get_via(), "Bank") # DLRTrain dlr_train = DLRTrain("Beckton", "1200") self.assertEqual(dlr_train.line_code, "DLR") # DepartureCollection fundamentals departures = DepartureCollection() departures[bus_stop] = [bus] self.assertEqual(departures[bus_stop], [bus]) self.assertEqual(len(departures), 1) del departures[bus_stop] self.assertFalse(bus_stop in departures) # DepartureCollection for trains departures.add_to_slot(bus_stop, bus) departures.add_to_slot(bus_stop, bus2) self.assertEqual(str(departures), "Trafalgar Square / Charing Cross Station to Blackwall 2359 0001") departures[bus_stop2] = [] departures.cleanup(lambda stop: NullDeparture("West")) self.assertEqual(str(departures), "None shown going West; Trafalgar Square / Charing Cross Station to Blackwall 2359 0001") departures.add_to_slot(bus_stop, Bus("Leamouth", "2358")) self.assertEqual(str(departures), "None shown going West; Trafalgar Square / Charing Cross Station to Leamouth 2358, Blackwall 2359 0001") # DepartureCollection for trains departures = DepartureCollection() departures["P1"] = [Train("Bank", "1210"), Train("Tower Gateway", "1203"), Train("Bank", "1200")] departures["P2"] = [Train("Tower Gateway", "1205"), Train("Tower Gateway", "1212"), Train("Bank", "1207")] departures["P3"] = [Train("Lewisham", "1200"), Train("Lewisham", "1204"), Train("Lewisham", "1208")] departures["P4"] = [] departures.merge_common_slots() departures.cleanup(lambda platform: NullDeparture("from %s" % platform)) self.assertEqual(str(departures), "Bank 1200 1207 1210, Tower Gateway 1203 1205; Lewisham 1200 1204 1208; None shown going from P4") departures.filter(lambda train: train.get_destination() != "Lewisham") self.assertEqual(str(departures), "Bank 1200 1207 1210, Tower Gateway 1203 1205; None shown going from P4") departures["P4"] = [] departures.filter(lambda train: train.get_destination() != "Tower Gateway", True) self.assertEqual(str(departures), "Bank 1200 1207 1210")