Skip to content

Instantly share code, notes, and snippets.

@nvbn
Created July 9, 2018 22:23
Show Gist options
  • Star 6 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save nvbn/05225aa1f55e5d57b71824010c3ba892 to your computer and use it in GitHub Desktop.
Save nvbn/05225aa1f55e5d57b71824010c3ba892 to your computer and use it in GitHub Desktop.
trip planner
import json
from datetime import date, datetime, timedelta
from collections import defaultdict, namedtuple
from multiprocessing import Pool, cpu_count
import csv
from operator import itemgetter
from heapq import merge
from itertools import islice
from dateutil.parser import parse
MIN_START = date(2018, 9, 1)
MAX_START = date(2018, 10, 1)
MIN_STAY = 3
MAX_STAY = 5
MAX_TRIP = 28
MIN_VISITED = 4
MAX_FLIGHT_PRICE = 500
MAX_HOME_FLIGHT_PRICE = 700
MAX_TRIP_PRICE = 2500
RESULT_SIZE = 100_000
HOME_CITY = 'Amsterdam'
CITIES = [
HOME_CITY,
'Rio de Janeiro',
'Montevideo',
'Asunción',
'Lima',
'Mexico City',
'Santiago',
'La Paz',
'Quito',
'Bogota',
'Buenos Aires',
]
Flight = namedtuple('Flight', ('from_id', 'to_id', 'day_number', 'price'))
Trip = namedtuple('Trip', ('price', 'flights'))
id2city = dict(enumerate(CITIES))
city2id = {city: id for id, city in id2city.items()}
city_ids = set(id2city.keys())
home_city_id = city2id[HOME_CITY]
def read_flights(path):
with open(path) as f:
data = json.load(f)
for flight_data in data:
if MIN_START <= parse(flight_data['date']).date() < MAX_START:
yield Flight(from_id=city2id[flight_data['from']],
to_id=city2id[flight_data['to']],
day_number=(parse(flight_data['date']).date() - MIN_START).days,
price=flight_data['price'])
flights_with_stop = [flight for flight in read_flights('flights.json')
if flight.price < MAX_FLIGHT_PRICE
and (flight.to_id == home_city_id or flight.from_id == home_city_id)
and flight.price < MAX_HOME_FLIGHT_PRICE]
flights_nonstop = [flight for flight in read_flights('flights_nonstop.json')
if flight.price < MAX_FLIGHT_PRICE
or ((flight.to_id == home_city_id or flight.from_id == home_city_id)
and flight.price < MAX_HOME_FLIGHT_PRICE)]
flights = flights_with_stop + flights_nonstop
print(f"Read {len(flights)} flights with tolerable price")
from_id2day_number2to_id2flight = defaultdict(
lambda: defaultdict(
lambda: {}))
for flight in flights:
from_id2day_number2to_id2flight[flight.from_id] \
[flight.day_number][flight.to_id] = flight
def _generate_trips(can_visit, can_travel, can_spent, current_id,
current_day, trip_flights):
if trip_flights[-1].to_id == home_city_id:
yield Trip(
price=sum(flight.price for flight in trip_flights),
flights=trip_flights)
return
if not can_visit or can_travel < MIN_STAY or can_spent == 0:
return
if len(trip_flights) >= MIN_VISITED and home_city_id not in can_visit:
can_visit.add(home_city_id)
for to_id in can_visit:
can_visit_next = can_visit.difference({to_id})
for stay in range(MIN_STAY, min(MAX_STAY, can_travel) + 1):
current_day_next = current_day + stay
flight_next = from_id2day_number2to_id2flight \
.get(current_id, {}).get(current_day_next, {}).get(to_id)
if not flight_next:
continue
can_spent_next = can_spent - flight_next.price
if can_spent_next < 0:
continue
yield from _generate_trips(
can_visit_next, can_travel - stay, can_spent_next, to_id,
current_day + stay, trip_flights + [flight_next])
def _generator_stage(params):
return sorted(_generate_trips(*params), key=itemgetter(0))
def generate_trips():
generators_params = [(
city_ids.difference({start_id, home_city_id}),
MAX_TRIP,
MAX_TRIP_PRICE - from_id2day_number2to_id2flight[home_city_id][start_day][start_id].price,
start_id,
start_day,
[from_id2day_number2to_id2flight[home_city_id][start_day][start_id]])
for start_day in range((MAX_START - MIN_START).days)
for start_id in from_id2day_number2to_id2flight[home_city_id][start_day].keys()]
start = datetime.now()
with Pool(cpu_count() * 2) as pool:
for n, stage_result in enumerate(pool.imap_unordered(_generator_stage, generators_params)):
print(f'Generated {n + 1} of {len(generators_params)} in {datetime.now() - start}')
yield stage_result
def format_trip(trip):
days = trip.flights[-1].day_number - trip.flights[0].day_number
cities = len(trip.flights) - 1
start_city = id2city[trip.flights[0].to_id]
start_date = MIN_START + timedelta(days=trip.flights[0].day_number)
end_city = id2city[trip.flights[-1].from_id]
end_date = MIN_START + timedelta(days=trip.flights[-1].day_number)
details = ' & '.join(id2city[flight.from_id] + ' -> ' + id2city[flight.to_id]
+ ' ' + str(MIN_START + timedelta(days=flight.day_number))
+ ' ' + str(flight.price)
for flight in trip.flights)
return (trip.price,
days, cities,
start_city, str(start_date),
end_city, str(end_date),
details)
trips = [*merge(*generate_trips(), key=itemgetter(0))]
print('Generated all possible trips')
with open('trips.json', 'w') as f:
json.dump(trips, f)
with open('trips.csv', 'w', newline='') as f:
writer = csv.writer(f)
writer.writerow(('price', 'days', 'cities', 'start city', 'start date', 'end city', 'end date', 'details'))
for trip in trips:
writer.writerow(format_trip(trip))
const HOME_CITY = 'Amsterdam';
const CITIES = [
'Rio de Janeiro',
'Montevideo',
'Asunción',
'Lima',
'Mexico City',
'Santiago',
'La Paz',
'Quito',
'Bogota',
'Buenos Aires',
];
const FROM = 0;
const TO = 1;
const getElements = (selector, intervalSize = 300) => new Promise((resolve) => {
const interval = setInterval(() => {
console.log(`Getting ${selector}`);
const elements = [...document.querySelectorAll(selector)];
if (elements.length) {
resolve(elements);
clearInterval(interval);
}
}, intervalSize);
});
const delay = (delaySize) => new Promise((resolve) => setTimeout(() => resolve(), delaySize));
const changeInput = (input, value) => {
input.value = value;
const keypressEvent = document.createEvent("HTMLEvents");
keypressEvent.initEvent('keypress');
input.dispatchEvent(keypressEvent);
};
const setDestination = async (type, value) => {
const holder = (await getElements('.gws-flights-form__location-text'))[type];
holder.click();
const [input] = await getElements('#flt-modaldialog input');
changeInput(input, value);
const [destination] = await getElements('#flt-modaldialog .fsapp-option-content');
await delay(500);
destination.click();
await delay(3000);
};
const getPrices = async () => {
const [holder] = await getElements('.gws-flights-form__calendar-input>div');
holder.click();
await delay(5000);
await getElements('.gws-travel-calendar__active [data-price]');
const dates = await getElements('[data-day] [data-price]');
const prices = dates
.filter(({textContent}) => textContent.length)
.map(({textContent, parentElement}) => [
parentElement.parentElement.getAttribute('data-day'),
parseInt(textContent.replace(/[^\d]*/g, ''), 10)
]);
const [underlay] = await getElements('.gws-flights__modal-underlay');
underlay.click();
await delay(1000);
return prices;
};
const getFlightsData = async ([from, to]) => {
await setDestination(FROM, from);
await setDestination(TO, to);
const prices = await getPrices();
return prices.map(([date, price]) => ({
date, price, from, to,
}));
};
function* getAllPossibleFlights() {
for (let from of CITIES) {
for (let to of CITIES) {
if (from !== to) {
yield [from, to];
}
}
yield [HOME_CITY, from];
yield [from, HOME_CITY];
}
}
const collectData = async () => {
let result = [];
for (let flight of getAllPossibleFlights()) {
console.log(`Fetching data for ${flight}`);
const flightsData = await getFlightsData(flight);
console.log(flightsData);
result = result.concat(flightsData);
}
return result;
};
const win = window.open('');
collectData().then(
(data) => win.document.write(JSON.stringify(data)),
(error) => console.error("Can't get flights", error),
);
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment