Codility solution for gas station problem

task.py

#!/usr/bin/env python3


class Pump(object):
    def __init__(self, fuel_capacity):
        self.fuel_capacity = fuel_capacity
        self.car = None


class GasStation(object):
    def __init__(self, pumps):
        self.pumps = pumps

    def pumps_have_capacity(self, fuel_required):
        return any([pump for pump in self.pumps if pump.fuel_capacity >= fuel_required])

    def free_pump(self, fuel_required):
        for pump in self.pumps:
            if pump.fuel_capacity >= fuel_required and not pump.car:
                return pump
        else:
            return None

    def resolve_fuel_up(self):
        # get next pump to be free (car with the least amount of fuel to fill up)
        min_fuel_need = min(
            [pump.car.fuel_need for pump in self.pumps if pump.car])

        cars_all_fueled_up = []

        for pump in self.pumps:
            if pump.car:
                pump.car.fuel_need -= min_fuel_need
                pump.fuel_capacity -= min_fuel_need

                if pump.car.fuel_need <= 0:
                    cars_all_fueled_up.append(pump.car)
                    pump.car = None

        # Return the car, and the amount of fuel that was used
        return cars_all_fueled_up, min_fuel_need

    def are_cars_fueling_up(self):
        for pump in self.pumps:
            if pump.car:
                return True
        else:
            return False


class Car(object):
    def __init__(self, fuel_need):
        self.fuel_need = fuel_need
        self.wait_time = 0


def solution(A, X, Y, Z):
    gas_station = GasStation([Pump(X), Pump(Y), Pump(Z)])
    line = [Car(fuel_need) for fuel_need in A]
    road = []

    while len(line) > 0 or gas_station.are_cars_fueling_up():
        if len(line) > 0:
            # Check if pumps have capacity for next car in line
            if not gas_station.pumps_have_capacity(fuel_required=line[0].fuel_need):
                return -1

            free_pump = gas_station.free_pump(fuel_required=line[0].fuel_need)

        if free_pump and len(line) > 0:
            free_pump.car = line.pop(0)

        else:
            resolved_cars, fuel_amount = gas_station.resolve_fuel_up()
            for car in line:
                car.wait_time += fuel_amount
            road.extend(resolved_cars)

    return max([car.wait_time for car in road])


if __name__ == "__main__":
    assert solution([2, 8, 4, 3, 2], 7, 11, 3) == 8
    assert solution([5], 4, 0, 3) == -1