Create a gist now

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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