<commit_before>import os
import json
from will.logger import log
from voluptuous import Schema, Required, MultipleInvalid
import API
import subprocess
import threading

#TODO: get these working

def plugin_loader(file_path):
    json_loader = JsonLoader(file_path)

    try:
        plugin_data = json_loader.load()
        plugin_data.build_plugin()

    except IOError:
        return

class JsonLoader:

    def __init__(self, file_path):
        self.file_path = file_path

    def is_json_file(self, fs_tools=os.path):
        if self.file_path.endswith('.json'):
            if fs_tools.exists(self.file_path) and fs_tools.isfile(self.file_path):  # noqa
                return True
        return False

    def load(self):
        if self.is_json_file():
            log.info("Loading plugin: {0}".format(self.file_path))
            with open(os.path.abspath(self.file_path), 'r') as file_stream:
                return JsonData(json.load(file_stream))
        raise IOError

class PluginBuilder:

    def __init__(self, plugin_data):
        self.plugin_data = plugin_data

    @staticmethod
    def thread_shell_call(command):
        thread = threading.Thread(target=lambda: subprocess.call(
            command, shell=True
        ))
        thread.daemon = True
        thread.start()

    def build(self):
        self.build_init_callback()
        self.build_subscribed_callback()
        self.build_shutdown_callback()

    def build_init_callback(self):
        if "init" in self.plugin_data:
            @API.init
            def plugin():
                PluginBuilder.thread_shell_call(
                    self.plugin_data["init"]
                )

    def build_subscribed_callback(self):
        '''Build the api data for the subscribe decorator'''
        data = False
        api_data = {"name" : '',
                    "ents_needed" : {},
                    "structure" : {"needed" : False},
                    "questions_needed" : False,
                    "key_words" : []
                    }
        if "name" in self.plugin_data:
            api_data["name"] = self.plugin_data["name"]
        if "ents_needed" in self.plugin_data:
            plugin_ents = self.plugin_data["ents_needed"]
            if plugin_ents and plugin_ents != {}:
                api_data['ents_needed'].update(plugin_ents)
                data = True
        if "structure" in self.plugin_data:
            struct_needed = self.plugin_data['structure']['needed']
            if struct_needed:
                api_data['structure']['needed'].append(struct_needed)
                data = True
        if "questions_needed" in self.plugin_data:
            q_needed = self.plugin_data["questions_needed"]
            if q_needed:
                api_data["questions_needed"] = q_needed
                data = True
        if "key_words" in self.plugin_data:
            plugin_key_words = self.plugin_data["key_words"]
            if plugin_key_words:
                for key_word in plugin_key_words:
                    api_data["key_words"].append(key_word)
                data = True
        if data:
            @API.subscribe_to(api_data)
            def plugin(*args):
                leader = args[0]
                full_text = args[1]
                arg_data = args[2][0]
                text = full_text.split(" ")[0]
                plugin_data = {"name": arg_data["name"],
                            "ents": arg_data["ents"],
                            "struct_needed" : arg_data["struct_needed"],
                            "full_text" : full_text,
                            "leader" : leader,
                            "headless_text" : text}
                command_str = self.plugin_data["command"]
                for ent in plugin_data["ents"]():
                    ent_place = "{%s}" % plugin_data["ents"][ent]
                    replace_str = ent
                    if isinstance(replace_str, list):
                        replace_str = replace_str[0]
                    if ent_place in command_str:
                        command_str = command_str.replace(ent_place, replace_str)
                for word_type in plugin_data["struct"].keys():
                    word_place = "{%s}" % word_type
                    replace_str = plugin_data["struct"][word_type]
                    if isinstance(replace_str, list):
                        replace_str = replace_str[0]
                    if word_place in command_str:
                        command_str = command_str.replace(word_place, replace_str)
                if "{full_text}" in command_str:
                    command_str = command_str.replace("{full_text}", full_text)
                if "{leader}" in command_str:
                    command_str = command_str.replace("{leader}", leader)
                if "{headless_text}" in command_str:
                    command_str = command_str.replace("{headless_text}", full_text)
                PluginBuilder.thread_shell_call(
                    command_str
                )
                return "Done"
        else:
            @API.subscribe_to_any
            def plugin(*args, **kwargs):
                def func():
                    PluginBuilder.thread_shell_call(
                        self.plugin_data["command"].format(args[1])
                    )

    def build_shutdown_callback(self):
        if "shutdown" in self.plugin_data:
            @API.shutdown
            def plugin():
                PluginBuilder.thread_shell_call(
                    self.plugin_data["shutdown"]
                )
class JsonData:

    def __init__(self, json_data):
        self.plugin_data = json_data

    def is_valid(self):
        validation_schema = Schema({
            "name": unicode,
            "init": unicode,
            "shutdown" : unicode,
            "ents_needed": dict,
            "structure" : dict,
            "questions_needed" : bool,
            "key_words" : list,
            Required("command") : unicode
        })
        try:
            validation_schema(self.plugin_data)
        except MultipleInvalid as e:
            log.exception("Error parsing plugin: {0}".format(e))
            return False
        return True

    def build_plugin(self):
        if self.is_valid():
            PluginBuilder(self.plugin_data).build()
<commit_msg>Json plugins bugfix<commit_after>import os
import json
from will.logger import log
from voluptuous import Schema, Required, MultipleInvalid
import API
import subprocess
import threading

#TODO: get these working

def plugin_loader(file_path):
    json_loader = JsonLoader(file_path)

    try:
        plugin_data = json_loader.load()
        plugin_data.build_plugin()

    except IOError:
        return

class JsonLoader:

    def __init__(self, file_path):
        self.file_path = file_path

    def is_json_file(self, fs_tools=os.path):
        if self.file_path.endswith('.json'):
            if fs_tools.exists(self.file_path) and fs_tools.isfile(self.file_path):  # noqa
                return True
        return False

    def load(self):
        if self.is_json_file():
            log.info("Loading plugin: {0}".format(self.file_path))
            with open(os.path.abspath(self.file_path), 'r') as file_stream:
                return JsonData(json.load(file_stream))
        raise IOError

class PluginBuilder:

    def __init__(self, plugin_data):
        self.plugin_data = plugin_data

    @staticmethod
    def thread_shell_call(command):
        thread = threading.Thread(target=lambda: subprocess.call(
            command, shell=True
        ))
        thread.daemon = True
        thread.start()

    def build(self):
        self.build_init_callback()
        self.build_subscribed_callback()
        self.build_shutdown_callback()

    def build_init_callback(self):
        if "init" in self.plugin_data:
            @API.init
            def plugin():
                PluginBuilder.thread_shell_call(
                    self.plugin_data["init"]
                )

    def build_subscribed_callback(self):
        '''Build the api data for the subscribe decorator'''
        data = False
        api_data = {"name" : '',
                    "ents_needed" : {},
                    "structure" : {"needed" : False},
                    "questions_needed" : False,
                    "key_words" : []
                    }
        if "name" in self.plugin_data:
            api_data["name"] = self.plugin_data["name"]
        if "ents_needed" in self.plugin_data:
            plugin_ents = self.plugin_data["ents_needed"]
            if plugin_ents and plugin_ents != {}:
                api_data['ents_needed'].update(plugin_ents)
                data = True
        if "structure" in self.plugin_data:
            struct_needed = self.plugin_data['structure']['needed']
            if struct_needed:
                api_data['structure']['needed'].append(struct_needed)
                data = True
        if "questions_needed" in self.plugin_data:
            q_needed = self.plugin_data["questions_needed"]
            if q_needed:
                api_data["questions_needed"] = q_needed
                data = True
        if "key_words" in self.plugin_data:
            plugin_key_words = self.plugin_data["key_words"]
            if plugin_key_words:
                for key_word in plugin_key_words:
                    api_data["key_words"].append(key_word)
                data = True
        if data:
            @API.subscribe_to(api_data)
            def plugin(*args):
                leader = args[0]
                full_text = args[1]
                arg_data = args[2][0]
                text = full_text.split(" ")[0]
                plugin_data = {"name": arg_data["name"],
                            "ents": arg_data["ents"],
                            "struct_needed" : arg_data["structure"]["tags"],
                            "full_text" : full_text,
                            "leader" : leader,
                            "headless_text" : text}
                command_str = self.plugin_data["command"]
                for ent in plugin_data["ents"]():
                    ent_place = "{%s}" % plugin_data["ents"][ent]
                    replace_str = ent
                    if isinstance(replace_str, list):
                        replace_str = replace_str[0]
                    if ent_place in command_str:
                        command_str = command_str.replace(ent_place, replace_str)
                for word_type in plugin_data["struct"].keys():
                    word_place = "{%s}" % word_type
                    replace_str = plugin_data["struct"][word_type]
                    if isinstance(replace_str, list):
                        replace_str = replace_str[0]
                    if word_place in command_str:
                        command_str = command_str.replace(word_place, replace_str)
                if "{full_text}" in command_str:
                    command_str = command_str.replace("{full_text}", full_text)
                if "{leader}" in command_str:
                    command_str = command_str.replace("{leader}", leader)
                if "{headless_text}" in command_str:
                    command_str = command_str.replace("{headless_text}", full_text)
                PluginBuilder.thread_shell_call(
                    command_str
                )
                return "Done"
        else:
            @API.subscribe_to_any
            def plugin(*args, **kwargs):
                def func():
                    PluginBuilder.thread_shell_call(
                        self.plugin_data["command"].format(args[1])
                    )

    def build_shutdown_callback(self):
        if "shutdown" in self.plugin_data:
            @API.shutdown
            def plugin():
                PluginBuilder.thread_shell_call(
                    self.plugin_data["shutdown"]
                )
class JsonData:

    def __init__(self, json_data):
        self.plugin_data = json_data

    def is_valid(self):
        validation_schema = Schema({
            "name": unicode,
            "init": unicode,
            "shutdown" : unicode,
            "ents_needed": dict,
            "structure" : dict,
            "questions_needed" : bool,
            "key_words" : list,
            Required("command") : unicode
        })
        try:
            validation_schema(self.plugin_data)
        except MultipleInvalid as e:
            log.exception("Error parsing plugin: {0}".format(e))
            return False
        return True

    def build_plugin(self):
        if self.is_valid():
            PluginBuilder(self.plugin_data).build()
<|endoftext|>


<commit_before>#   Copyright (c) 2006-2007 Open Source Applications Foundation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import webbrowser
import windmill
import exceptions
import os, sys, shutil, subprocess, time
import logging

logger = logging.getLogger(__name__)

os.environ['MOZ_NO_REMOTE'] = str(1)
PROXY_PORT = windmill.settings['SERVER_HTTP_PORT']
DEFAULT_TEST_URL = windmill.settings['TEST_URL']+'/windmill-serv/start.html'
MOZILLA_PROFILE_PATH = windmill.settings['MOZILLA_PROFILE_PATH']
MOZILLA_DEFAULT_PROFILE = windmill.settings['MOZILLA_DEFAULT_PROFILE']

class MozillaProfile(object):

    def __init__(self, path=MOZILLA_PROFILE_PATH, default_profile=MOZILLA_DEFAULT_PROFILE,
                  proxy_port=PROXY_PORT, test_url=DEFAULT_TEST_URL):
        """Create profile dir, and populate it"""

        self.proxy_host = 'localhost'
        self.proxy_port = proxy_port
        self.test_url = test_url

        self.profile_path = path

        if windmill.settings['MOZILLA_CREATE_NEW_PROFILE']:
            if os.path.exists(self.profile_path) is True:
                shutil.rmtree(self.profile_path)

            shutil.copytree(default_profile, self.profile_path)

            self.prefs_js_filename = self.profile_path + '/prefs.js'
            self.prefs_js_f = open( self.prefs_js_filename, 'w')
            self.initial_prefs()

    def initial_prefs(self):
        """Initial prefs population, separated form __init__ for ease of subclassing"""

        # Get rid of default browser check
        self.user_pref('"browser.shell.checkDefaultBrowser", false')
        # Suppress authentication confirmations
        self.user_pref('"network.http.phishy-userpass-length", 255')
        # Disable pop-up blocking
        self.user_pref('"browser.allowpopups", true')
        self.user_pref('"dom.disable_open_during_load", false')
        # Open links in new windows (Firefox 2.0)
        self.user_pref('"browser.link.open_external", 2')
        self.user_pref('"browser.link.open_newwindow", 2')
        # Configure local proxy
        self.user_pref('"network.proxy.http", "%s"' % self.proxy_host)
        self.user_pref('"network.proxy.http_port", %s' % str(self.proxy_port))
        self.user_pref('"network.proxy.no_proxies_on", ""')
        self.user_pref('"network.proxy.type", 1')

        self.user_pref('"network.http.max-connections", 40')
        self.user_pref('"network.http.max-connections-per-server", 16')
        self.user_pref('"network.http.max-persistent-connections-per-proxy", 12')
        self.user_pref('"network.http.max-persistent-connections-per-server", 6')
        self.user_pref('"network.http.pipelining.maxrequests", 6')

        # Turn off favicon requests, no need for even more requests
        self.user_pref('"browser.chrome.favicons", false')

        self.user_pref('"startup.homepage_override_url", "' + self.test_url + '"')
        self.user_pref('"browser.startup.homepage", "' + self.test_url + '"')
        self.user_pref('"startup.homepage_welcome_url", ""')
        # Disable security warnings
        self.user_pref('"security.warn_submit_insecure", false')
        self.user_pref('"security.warn_submit_insecure.show_once", false')
        self.user_pref('"security.warn_entering_secure", false')
        self.user_pref('"security.warn_entering_secure.show_once", false')
        self.user_pref('"security.warn_entering_weak", false')
        self.user_pref('"security.warn_entering_weak.show_once", false')
        self.user_pref('"security.warn_leaving_secure", false')
        self.user_pref('"security.warn_leaving_secure.show_once", false')
        self.user_pref('"security.warn_viewing_mixed", false')
        self.user_pref('"security.warn_viewing_mixed.show_once", false')
        # Disable cache
        self.user_pref('"browser.cache.disk.enable", false')
        self.user_pref('"browser.cache.memory.enable", false')
        # Disable "do you want to remember this password?"
        self.user_pref('"signon.rememberSignons", false')

    def user_pref(self, string):
        self.prefs_js_f.write('user_pref(' + string + ');\n')
        self.prefs_js_f.flush()

    def add_js(self, string):
        self.prefs_js_f.write(string + '\n')
        self.prefs_js_f.flush()

    def clean_up(self):
        shutil.rmtree(self.profile_path)

MOZILLA_BINARY = windmill.settings['MOZILLA_BINARY']

class MozillaBrowser(object):
    """MozillaBrowser class, init requires MozillaProfile instance"""
    def __init__(self, profile, mozilla_bin=MOZILLA_BINARY, start_url=None):

        self.profile = profile
        self.mozilla_bin = mozilla_bin
        self.p_id = None

        self.shell_command = "%s -profile %s %s" % (self.mozilla_bin, self.profile.profile_path, self.profile.test_url)

    def start(self):

        self.p_id = subprocess.Popen(self.shell_command, shell=True).pid

        logger.info(self.shell_command)

    def is_alive(self):

        if self.p_id is None:
            return False

        try:
            os.kill(self.p_id, 0)
            return True
        except exceptions.OSError:
            return False

    def kill(self, signal):

        os.kill(self.p_id, signal)
        try:
            os.kill(self.p_id, 0)
            if sys.platform == 'darwin':
                os.kill(self.p_id+1, signal)
                os.kill(self.p_id, signal)
        except:
            pass

    def stop(self):
        import signal

        self.kill(signal=signal.SIGTERM)

            <commit_msg>More cygwin fixes.<commit_after>#   Copyright (c) 2006-2007 Open Source Applications Foundation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import webbrowser
import windmill
import exceptions
import os, sys, shutil, subprocess, time
import logging

logger = logging.getLogger(__name__)

os.environ['MOZ_NO_REMOTE'] = str(1)
PROXY_PORT = windmill.settings['SERVER_HTTP_PORT']
DEFAULT_TEST_URL = windmill.settings['TEST_URL']+'/windmill-serv/start.html'
MOZILLA_PROFILE_PATH = windmill.settings['MOZILLA_PROFILE_PATH']
MOZILLA_DEFAULT_PROFILE = windmill.settings['MOZILLA_DEFAULT_PROFILE']

class MozillaProfile(object):

    def __init__(self, path=MOZILLA_PROFILE_PATH, default_profile=MOZILLA_DEFAULT_PROFILE,
                  proxy_port=PROXY_PORT, test_url=DEFAULT_TEST_URL):
        """Create profile dir, and populate it"""

        self.proxy_host = 'localhost'
        self.proxy_port = proxy_port
        self.test_url = test_url

        self.profile_path = path

        if windmill.settings['MOZILLA_CREATE_NEW_PROFILE']:
            if os.path.exists(self.profile_path) is True:
                shutil.rmtree(self.profile_path)

            shutil.copytree(default_profile, self.profile_path)

            self.prefs_js_filename = self.profile_path + '/prefs.js'
            self.prefs_js_f = open( self.prefs_js_filename, 'w')
            self.initial_prefs()

        if sys.platform == 'cygwin':
            new_path = windmill.settings['MOZILLA_PROFILE_PATH'].replace('/cygdrive/c/', 'C:\\').replace('/', '\\')
            windmill.settings['MOZILLA_PROFILE_PATH'] = new_path

    def initial_prefs(self):
        """Initial prefs population, separated form __init__ for ease of subclassing"""

        # Get rid of default browser check
        self.user_pref('"browser.shell.checkDefaultBrowser", false')
        # Suppress authentication confirmations
        self.user_pref('"network.http.phishy-userpass-length", 255')
        # Disable pop-up blocking
        self.user_pref('"browser.allowpopups", true')
        self.user_pref('"dom.disable_open_during_load", false')
        # Open links in new windows (Firefox 2.0)
        self.user_pref('"browser.link.open_external", 2')
        self.user_pref('"browser.link.open_newwindow", 2')
        # Configure local proxy
        self.user_pref('"network.proxy.http", "%s"' % self.proxy_host)
        self.user_pref('"network.proxy.http_port", %s' % str(self.proxy_port))
        self.user_pref('"network.proxy.no_proxies_on", ""')
        self.user_pref('"network.proxy.type", 1')

        self.user_pref('"network.http.max-connections", 40')
        self.user_pref('"network.http.max-connections-per-server", 16')
        self.user_pref('"network.http.max-persistent-connections-per-proxy", 12')
        self.user_pref('"network.http.max-persistent-connections-per-server", 6')
        self.user_pref('"network.http.pipelining.maxrequests", 6')

        # Turn off favicon requests, no need for even more requests
        self.user_pref('"browser.chrome.favicons", false')

        self.user_pref('"startup.homepage_override_url", "' + self.test_url + '"')
        self.user_pref('"browser.startup.homepage", "' + self.test_url + '"')
        self.user_pref('"startup.homepage_welcome_url", ""')
        # Disable security warnings
        self.user_pref('"security.warn_submit_insecure", false')
        self.user_pref('"security.warn_submit_insecure.show_once", false')
        self.user_pref('"security.warn_entering_secure", false')
        self.user_pref('"security.warn_entering_secure.show_once", false')
        self.user_pref('"security.warn_entering_weak", false')
        self.user_pref('"security.warn_entering_weak.show_once", false')
        self.user_pref('"security.warn_leaving_secure", false')
        self.user_pref('"security.warn_leaving_secure.show_once", false')
        self.user_pref('"security.warn_viewing_mixed", false')
        self.user_pref('"security.warn_viewing_mixed.show_once", false')
        # Disable cache
        self.user_pref('"browser.cache.disk.enable", false')
        self.user_pref('"browser.cache.memory.enable", false')
        # Disable "do you want to remember this password?"
        self.user_pref('"signon.rememberSignons", false')

    def user_pref(self, string):
        self.prefs_js_f.write('user_pref(' + string + ');\n')
        self.prefs_js_f.flush()

    def add_js(self, string):
        self.prefs_js_f.write(string + '\n')
        self.prefs_js_f.flush()

    def clean_up(self):
        shutil.rmtree(self.profile_path)

MOZILLA_BINARY = windmill.settings['MOZILLA_BINARY']

class MozillaBrowser(object):
    """MozillaBrowser class, init requires MozillaProfile instance"""
    def __init__(self, profile, mozilla_bin=MOZILLA_BINARY, start_url=None):

        self.profile = profile
        self.mozilla_bin = mozilla_bin
        self.p_id = None

        if sys.platform == 'cygwin':
            self.shell_command = '%s -profile "%s" ' % (self.mozilla_bin, self.profile.profile_path, self.profile.test_url)
        else:
            self.shell_command = "%s -profile %s" % (self.mozilla_bin, self.profile.profile_path, self.profile.test_url)

    def start(self):

        self.p_id = subprocess.Popen(self.shell_command, shell=True).pid

        logger.info(self.shell_command)

    def is_alive(self):

        if self.p_id is None:
            return False

        try:
            os.kill(self.p_id, 0)
            return True
        except exceptions.OSError:
            return False

    def kill(self, signal):

        os.kill(self.p_id, signal)
        try:
            os.kill(self.p_id, 0)
            if sys.platform == 'darwin':
                os.kill(self.p_id+1, signal)
                os.kill(self.p_id, signal)
        except:
            pass

    def stop(self):
        import signal

        self.kill(signal=signal.SIGTERM)

            <|endoftext|>


<commit_before>from pprint import pprint
import app, db
import googlemaps
import time

def gen_distance_matrix(start, stores, depart_delay):

    # :param start: Address as a string.
    #   E.g. '1600 Pennsylvania Ave NW, Washington, DC 20006'

    # :param stores: List of Addresses, each element as a string

    # :depart_delay: A nonnegative amount of time (in seconds)
    #       and used to query traffic conditions of when user
    #       leaves relative to when the function is called

    # Generates a distance and time matrix between start locations and stores
    # Matrix is represented as a matrix of (duration_in_traffic, distance)
    # The rows represent starting locations, the columns represent ending
    # locations. The indexing format always follows [start, *stores]

    # Create Google Maps client with Server Key
    gmaps = googlemaps.Client(key=app.config["GOOGLE_MAPS_API_KEY"])

    # Create list of lat/lng for start location and stores
    start_list = list()
    gcode_start = gmaps.geocode(start)
    loc_list.append(gcode_start[0]['geometry']['location'])

    for store in stores:
        gcode_store = gmaps.geocode(store)
        loc_list.append(gcode_store[0]['geometry']['location'])

    pprint(start_list)

    # Setting Departure Time
    depart_time = int(time.time()) + depart_delay

    # Get Google Maps client to make distance matrix
    d_mat = gmaps.distance_matrix(loc_list, loc_list, departure_time=depart_time, mode='driving', units='miles')

    pprint(d_mat)

    # Cleaning distance matrix to be a list of list of (duration_in_traffic, distance)
    distance_matrix = list()
    for src in range(len(d_mat['rows'])):
        tmp_list = list()
        for dst in range(len(d_mat['rows'][src]['elements'])):
            tmp_list.append((d_mat['rows'][src]['elements'][dst]['duration_in_traffic']['value'] / 60., \
                            d_mat['rows'][src]['elements'][dst]['distance']['value'] / 1000.))
        distance_matrix.append(tmp_list)

    return distance_matrix

if __name__ == "__main__":

    print "OK"
<commit_msg>minor change<commit_after>from pprint import pprint
import app, db
import googlemaps
import time

def gen_distance_matrix(start, stores, depart_delay):

    # :param start: Address as a string.
    #   E.g. '1600 Pennsylvania Ave NW, Washington, DC 20006'

    # :param stores: List of Addresses, each element as a string

    # :depart_delay: A nonnegative amount of time (in seconds)
    #       and used to query traffic conditions of when user
    #       leaves relative to when the function is called

    # Generates a distance and time matrix between start locations and stores
    # Matrix is represented as a matrix of (duration_in_traffic, distance)
    # The rows represent starting locations, the columns represent ending
    # locations. The indexing format always follows [start, *stores]

    # Create Google Maps client with Server Key
    gmaps = googlemaps.Client(key=app.config["GOOGLE_MAPS_API_KEY"])

    # Create list of lat/lng for start location and stores
    loc_list = list()
    gcode_start = gmaps.geocode(start)
    loc_list.append(gcode_start[0]['geometry']['location'])

    for store in stores:
        gcode_store = gmaps.geocode(store)
        loc_list.append(gcode_store[0]['geometry']['location'])

    pprint(loc_list)

    # Setting Departure Time
    depart_time = int(time.time()) + depart_delay

    # Get Google Maps client to make distance matrix
    d_mat = gmaps.distance_matrix(loc_list, loc_list, departure_time=depart_time, mode='driving', units='miles')

    pprint(d_mat)

    # Cleaning distance matrix to be a list of list of (duration_in_traffic, distance)
    distance_matrix = list()
    for src in range(len(d_mat['rows'])):
        tmp_list = list()
        for dst in range(len(d_mat['rows'][src]['elements'])):
            tmp_list.append((d_mat['rows'][src]['elements'][dst]['duration_in_traffic']['value'] / 60., \
                            d_mat['rows'][src]['elements'][dst]['distance']['value'] / 1000.))
        distance_matrix.append(tmp_list)

    return distance_matrix

if __name__ == "__main__":

    print "OK"
<|endoftext|>


<commit_before>from six import exec_
import os.path as op
import warnings
from collections import Counter
from .formats import open_ephys, kwik, dat
import logging

DEFAULT_LIMIT = 0.512  # GB

logger = logging.getLogger(__name__)

def get_batch_size(arr, ram_limit_gb=DEFAULT_LIMIT):
    """Return batch size, number of full batches and remainder size for 2D array."""
    batch_size = int(ram_limit_gb * 1e9 / arr.shape[1] / arr.dtype.itemsize)
    return batch_size, arr.shape[0] // batch_size, arr.shape[0] % batch_size

def detect_format(path):
    formats = [f for f in [fmt.detect(path) for fmt in [open_ephys, dat, kwik]] if f is not None]
    if len(formats) == 1:
        fmt = formats[0]
        if 'DAT' in fmt:
            if fmt == 'DAT-File':
                return dat
        else:
            if 'kwik' in fmt:
                return kwik
            else:
                return open_ephys
    logger.info('Detected format(s) {} not valid.'.format(formats))

def run_prb(path):
    """Execute the .prb probe file and import resulting locals return results as dict.
    Args:
        path: file path to probe file with layout as per klusta probe file specs.

    Returns: Dictionary of channel groups with channel list, geometry and connectivity graph.
    """
    if path is None:
        return

    path = op.realpath(op.expanduser(path))
    assert op.exists(path)
    with open(path, 'r') as prb:
        layout = prb.read()

    metadata = {}
    exec_(layout, {}, metadata)
    metadata = {k.lower(): v for (k, v) in metadata.items()}
    return metadata

def flat_channel_list(prb):
    """Flat lists of all channels and bad channels given a probe dictionary.

    Args:
        prb: Path to probe file.

    Returns:
        Tuple of lists of channels.
    """
    channels = sum([prb['channel_groups'][cg]['channels'] for cg in sorted(prb['channel_groups'])], [])
    dead_channels = prb['dead_channels']

    return channels, dead_channels

def monotonic_prb(prb):
    """Return probe file dict with monotonically increasing channel group and channel numbers."""
    # FIXME: Should otherwise copy any other fields over (unknown fields warning)
    # FIXME: Correct ref like dc to indices
    chan_n = 0
    groups = prb['channel_groups']
    monotonic = {}
    for n, chg in enumerate(groups.keys()):
        monotonic[n] = list(range(chan_n, chan_n + len(groups[chg]['channels'])))
        chan_n += len(groups[chg]['channels'])
    ret_dict = {'channel_groups': monotonic}

    # correct bad channel indices
    if 'dead_channels' in prb.keys():
        fcl, fbc = flat_channel_list(prb)
        ret_dict['dead_channels'] = sorted([fcl.index(bc) for bc in fbc])
    return ret_dict

def make_prb():
    raise NotImplementedError

def make_prm():
    raise NotImplementedError

def channel_ranges(channel_list):
    """List of channels in to ranges of consecutive channels.

    Args:
        channel_list: list of channel numbers (ints)

    Returns:
        List of list of channels grouped into consecutive sequences.

    Example: channel_ranges([1, 3, 4, 5]) -> [[1], [3, 4, 5]]
    """
    duplicates = [c for c, n in Counter(channel_list).items() if n > 1]
    if len(duplicates):
        warnings.warn("Warning: Channel(s) {} listed more than once".format(duplicates))

    ranges = [[]]
    for channel in channel_list:
        if len(ranges[-1]) and abs(channel - ranges[-1][-1]) != 1:
            ranges.append([])
        ranges[-1].append(channel)
    return ranges

def fmt_channel_ranges(channels, shorten_seq=5, rs="tm", c_sep="_", zp=2):
    """String of channel numbers separated with delimiters with consecutive channels
    are shortened when sequence length above threshold.

    Args:
        channels: list of channels
        shorten_seq: number of consecutive channels to be shortened. (default: 5)
        rs: range delimiter (default: 'tm')
        c_sep: channel delimiter (default: '_')
        zp: zero pad channel numbers (default: 2)

    Returns: String of channels in order they appeared in the list of channels.

    Example: fmt_channel_ranges([[1], [3], [5, 6, 7, 8, 9, 10]]) -> 01_03_05tm10
    """
    c_ranges = channel_ranges(channels)
    range_strings = [c_sep.join(["{c:0{zp}}".format(c=c, zp=zp) for c in c_seq])
                     if len(c_seq) < shorten_seq
                     else "{start:0{zp}}{rs}{end:0{zp}}".format(start=c_seq[0], end=c_seq[-1], rs=rs, zp=zp)
                     for c_seq in c_ranges]
    return c_sep.join(range_strings)

def fmt_time(s, minimal=True):
    """
    Args:
        s: time in seconds (float for fractional)
        minimal: Flag, if true, only return strings for times > 0, leave rest outs
    Returns: String formatted 99h 59min 59.9s, where elements < 1 are left out optionally.

    """
    ms = s - int(s)
    s = int(s)
    if s < 60 and minimal:
        return "{s:02.3f}s".format(s=s + ms)

    m, s = divmod(s, 60)
    if m < 60 and minimal:
        return "{m:02d}min {s:02.3f}s".format(m=m, s=s + ms)

    h, m = divmod(m, 60)
    return "{h:02d}h {m:02d}min {s:02.3f}s".format(h=h, m=m, s=s + ms)

def get_needed_channels(cli_args=None):
    """Gets a list of channels that are needed in order to process a given channel.
    """
    if cli_args is None:
        import argparse
        parser = argparse.ArgumentParser(description=get_needed_channels.__doc__)
        parser.add_argument('probe_file', nargs=1,
                            help="""Phe probe file to be used""")
        parser.add_argument("groups", nargs='+',
                            help="""A list of groups""")
        parser.add_argument("-f", "--filenames", action='store_true',
                            help="""Returns a list of open-ephys continuous filenames, instead of a list of channel
                                    numbers""")
        parser.add_argument("-n", "--node", type=int,
                            help="""A node number for the filenames (default 100)""")
        parser.add_argument("--zerobased", action='store_true',
                            help="Use klusta zero-based convention instead of open-ephys 1-based one")
        cli_args = parser.parse_args()

    probe_file = cli_args.probe_file[0]
    groups = [int(g) for g in cli_args.groups]

    do_filenames = False
    if cli_args.filenames:
        do_filenames = True

    if cli_args.node:
        node = cli_args.node
        do_filenames = True
    else:
        node = 100

    zero_based = False
    if cli_args.zerobased:
        zero_based = True

    layout = run_prb(probe_file)

    channels = []
    for g in groups:
        channels.extend(layout['channel_groups'][g]['channels'])
        if 'reference' in layout['channel_groups']:
            channels.extend(layout['channel_groups'][g]['reference'])

    if 'ref_a' in layout:
        channels.extend(layout['ref_a'])

    if 'ref_b' in layout:
        channels.extend(layout['ref_b'])

    if not zero_based:
        channels = [c + 1 for c in channels]
    channels = set(channels)

    if do_filenames:
        fnames = [str(node) + '_CH' + str(c) + '.continuous' for c in channels]
        print('\n'.join(fnames))
    else:
        print(' '.join(map(str, channels)))<commit_msg>Fix .prb file formatting<commit_after>from six import exec_
import os.path as op
import warnings
from collections import Counter
from .formats import open_ephys, kwik, dat
import logging

DEFAULT_LIMIT = 0.512  # GB

logger = logging.getLogger(__name__)

def get_batch_size(arr, ram_limit_gb=DEFAULT_LIMIT):
    """Return batch size, number of full batches and remainder size for 2D array."""
    batch_size = int(ram_limit_gb * 1e9 / arr.shape[1] / arr.dtype.itemsize)
    return batch_size, arr.shape[0] // batch_size, arr.shape[0] % batch_size

def detect_format(path):
    formats = [f for f in [fmt.detect(path) for fmt in [open_ephys, dat, kwik]] if f is not None]
    if len(formats) == 1:
        fmt = formats[0]
        if 'DAT' in fmt:
            if fmt == 'DAT-File':
                return dat
        else:
            if 'kwik' in fmt:
                return kwik
            else:
                return open_ephys
    logger.info('Detected format(s) {} not valid.'.format(formats))

def run_prb(path):
    """Execute the .prb probe file and import resulting locals return results as dict.
    Args:
        path: file path to probe file with layout as per klusta probe file specs.

    Returns: Dictionary of channel groups with channel list, geometry and connectivity graph.
    """
    if path is None:
        return

    path = op.realpath(op.expanduser(path))
    assert op.exists(path)
    with open(path, 'r') as prb:
        layout = prb.read()

    metadata = {}
    exec_(layout, {}, metadata)
    metadata = {k.lower(): v for (k, v) in metadata.items()}
    return metadata

def flat_channel_list(prb):
    """Flat lists of all channels and bad channels given a probe dictionary.

    Args:
        prb: Path to probe file.

    Returns:
        Tuple of lists of channels.
    """
    channels = sum([prb['channel_groups'][cg]['channels'] for cg in sorted(prb['channel_groups'])], [])
    dead_channels = prb['dead_channels']

    return channels, dead_channels

def monotonic_prb(prb):
    """Return probe file dict with monotonically increasing channel group and channel numbers."""
    # FIXME: Should otherwise copy any other fields over (unknown fields warning)
    # FIXME: Correct ref like dc to indices
    chan_n = 0
    groups = prb['channel_groups']
    monotonic = {}
    for n, chg in enumerate(groups.keys()):
        monotonic[n] = {'channels': list(range(chan_n, chan_n + len(groups[chg]['channels'])))}
        chan_n += len(groups[chg]['channels'])

    # correct bad channel indices
    if 'dead_channels' in prb.keys():
        fcl, fbc = flat_channel_list(prb)
        dead_channels = sorted([fcl.index(bc) for bc in fbc])
    else:
        dead_channels = []
    return monotonic, dead_channels

def make_prb():
    raise NotImplementedError

def make_prm():
    raise NotImplementedError

def channel_ranges(channel_list):
    """List of channels in to ranges of consecutive channels.

    Args:
        channel_list: list of channel numbers (ints)

    Returns:
        List of list of channels grouped into consecutive sequences.

    Example: channel_ranges([1, 3, 4, 5]) -> [[1], [3, 4, 5]]
    """
    duplicates = [c for c, n in Counter(channel_list).items() if n > 1]
    if len(duplicates):
        warnings.warn("Warning: Channel(s) {} listed more than once".format(duplicates))

    ranges = [[]]
    for channel in channel_list:
        if len(ranges[-1]) and abs(channel - ranges[-1][-1]) != 1:
            ranges.append([])
        ranges[-1].append(channel)
    return ranges

def fmt_channel_ranges(channels, shorten_seq=5, rs="tm", c_sep="_", zp=2):
    """String of channel numbers separated with delimiters with consecutive channels
    are shortened when sequence length above threshold.

    Args:
        channels: list of channels
        shorten_seq: number of consecutive channels to be shortened. (default: 5)
        rs: range delimiter (default: 'tm')
        c_sep: channel delimiter (default: '_')
        zp: zero pad channel numbers (default: 2)

    Returns: String of channels in order they appeared in the list of channels.

    Example: fmt_channel_ranges([[1], [3], [5, 6, 7, 8, 9, 10]]) -> 01_03_05tm10
    """
    c_ranges = channel_ranges(channels)
    range_strings = [c_sep.join(["{c:0{zp}}".format(c=c, zp=zp) for c in c_seq])
                     if len(c_seq) < shorten_seq
                     else "{start:0{zp}}{rs}{end:0{zp}}".format(start=c_seq[0], end=c_seq[-1], rs=rs, zp=zp)
                     for c_seq in c_ranges]
    return c_sep.join(range_strings)

def fmt_time(s, minimal=True):
    """
    Args:
        s: time in seconds (float for fractional)
        minimal: Flag, if true, only return strings for times > 0, leave rest outs
    Returns: String formatted 99h 59min 59.9s, where elements < 1 are left out optionally.

    """
    ms = s - int(s)
    s = int(s)
    if s < 60 and minimal:
        return "{s:02.3f}s".format(s=s + ms)

    m, s = divmod(s, 60)
    if m < 60 and minimal:
        return "{m:02d}min {s:02.3f}s".format(m=m, s=s + ms)

    h, m = divmod(m, 60)
    return "{h:02d}h {m:02d}min {s:02.3f}s".format(h=h, m=m, s=s + ms)

def get_needed_channels(cli_args=None):
    """Gets a list of channels that are needed in order to process a given channel.
    """
    if cli_args is None:
        import argparse
        parser = argparse.ArgumentParser(description=get_needed_channels.__doc__)
        parser.add_argument('probe_file', nargs=1,
                            help="""Phe probe file to be used""")
        parser.add_argument("groups", nargs='+',
                            help="""A list of groups""")
        parser.add_argument("-f", "--filenames", action='store_true',
                            help="""Returns a list of open-ephys continuous filenames, instead of a list of channel
                                    numbers""")
        parser.add_argument("-n", "--node", type=int,
                            help="""A node number for the filenames (default 100)""")
        parser.add_argument("--zerobased", action='store_true',
                            help="Use klusta zero-based convention instead of open-ephys 1-based one")
        cli_args = parser.parse_args()

    probe_file = cli_args.probe_file[0]
    groups = [int(g) for g in cli_args.groups]

    do_filenames = False
    if cli_args.filenames:
        do_filenames = True

    if cli_args.node:
        node = cli_args.node
        do_filenames = True
    else:
        node = 100

    zero_based = False
    if cli_args.zerobased:
        zero_based = True

    layout = run_prb(probe_file)

    channels = []
    for g in groups:
        channels.extend(layout['channel_groups'][g]['channels'])
        if 'reference' in layout['channel_groups']:
            channels.extend(layout['channel_groups'][g]['reference'])

    if 'ref_a' in layout:
        channels.extend(layout['ref_a'])

    if 'ref_b' in layout:
        channels.extend(layout['ref_b'])

    if not zero_based:
        channels = [c + 1 for c in channels]
    channels = set(channels)

    if do_filenames:
        fnames = [str(node) + '_CH' + str(c) + '.continuous' for c in channels]
        print('\n'.join(fnames))
    else:
        print(' '.join(map(str, channels)))<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright 2011 Aaron Steele and John Wieczorek
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""This script implements the MoL workflow process for layers."""

import copy
import csv
import glob
import logging
from unicodewriterhelper import UnicodeWriter
from optparse import OptionParser
import os
import simplejson
import shlex
import subprocess
import sys
import yaml

class Config(object):
    """Wraps the YAML object for a MoL config.yaml object."""

    @classmethod
    def lower_keys(cls, x):
        """Lower cases all nested dictionary keys."""
        if isinstance(x, list):
            return [cls.lower_keys(v) for v in x]
        if isinstance(x, dict):
            return dict((k.lower(), cls.lower_keys(v)) for k, v in x.iteritems())
        return x

    class Collection(object):
        def __init__(self, collection):
            self.collection = collection

        def get_row(self):
            row = {}
            for k,v in self.collection['required'].iteritems():
                row[k] = v
            for k,v in self.collection['optional'].iteritems():
                if k in ['taxonomy', 'basemaps']:
                    v = simplejson.dumps(v)
                row[k] = v
            return row

        def get_columns(self):
            cols = []
            cols.extend(self.collection['required'].keys())
            cols.extend(self.collection['optional'].keys())
            cols.extend(self.collection['dbfmapping']['required'].keys())
            cols.extend(self.collection['dbfmapping']['optional'].keys())
            cols.extend(['layer_source', 'layer_collection', 'layer_filename', 'layer_polygons'])
            return cols

        def get_mapping(self, required=True):
            """Returns a reverse dbfmapping for convienience."""
            if required:
                mapping = self.collection['dbfmapping']['required']
            else:
                mapping = self.collection['dbfmapping']['optional']
            return dict((source, mol) for mol,source in mapping.iteritems())

        def getdir(self):
            return self.collection['directoryname']

        def get(self, key, default=None):
            return self.collection.get(key, default)

    def __init__(self, filename):
        self.config = Config.lower_keys(yaml.load(open(filename, 'r').read()))

    def collection_names(self):
        return [x.get('directoryname') for x in self.collections()]

    def collections(self):
        for collection in self.config['collections']:
            yield Config.Collection(collection)

def source2csv(source_dir, options):
    ''' Loads the collections in the given source directory.

        Arguments:
            source_dir - the relative path to the directory in which the config.yaml file is located.
    '''
    config = Config(os.path.join(source_dir, 'config.yaml'))
    logging.info('Collections in %s: %s' % (source_dir, config.collection_names()))

    for collection in config.collections(): # For each collection dir in the source dir
        coll_dir = collection.getdir()

        original_dir = os.getcwd()           # We'll need this to restore us to this dir at the end of processing this collection.
        os.chdir(os.path.join(source_dir, coll_dir))

        # Create collection.csv writer
        coll_file = open('collection.csv.txt', 'w')
        coll_cols = collection.get_columns()
        coll_cols.sort()
#        coll_csv = UnicodeWriter(coll_file, coll_cols)
        coll_csv = csv.DictWriter(coll_file, coll_cols)
        coll_csv.writer.writerow(coll_csv.fieldnames)
        coll_row = collection.get_row()
        coll_row['layer_source'] = source_dir
        coll_row['layer_collection'] = coll_dir

        # Create polygons.csv writer
        poly_file = open('collection.polygons.csv.txt', 'w')
#        poly_dw = UnicodeWriter(poly_file, ['shapefilename', 'json'])
        poly_dw = csv.DictWriter(poly_file, ['shapefilename', 'json'])
        poly_dw.writer.writerow(poly_dw.fieldnames)

        # Convert DBF to CSV and add to collection.csv
        shpfiles = glob.glob('*.shp')
        logging.info('Processing %d layers in the %s/%s' % (len(shpfiles), source_dir, coll_dir))
        for sf in shpfiles:
            logging.info('Extracting DBF fields from %s' % sf)
            csvfile = '%s.csv' % sf
            if os.path.exists(csvfile): # ogr2ogr barfs if there are *any* csv files in the dir
                os.remove(csvfile)
            # Following line only for convenience for running on Mac in eclipse
#            command = '/Library/Frameworks/GDAL.framework/Programs/ogr2ogr -f CSV "%s" "%s"' % (csvfile, sf)
            command = 'ogr2ogr -f CSV "%s" "%s"' % (csvfile, sf)
            args = shlex.split(command)
            try:
                subprocess.call(args)
            except OSError as errmsg:
                print """Error occurred while executing command line '{0}': {2}
    Please ensure that {1} is executable and available on your path.
                """.format(command, args[0], errmsg)
                raise

            # Copy and update coll_row with DBF fields
            row = copy.copy(coll_row)
            row['layer_filename'] = os.path.splitext(sf)[0]
            dr = csv.DictReader(open(csvfile, 'r'), skipinitialspace=True)

            layer_polygons = []

            for dbf in dr: # For each row in the DBF CSV file (1 row per polygon)

                polygon = {}

                for source, mol in collection.get_mapping().iteritems(): # Required DBF fields
                    sourceval = dbf.get(source)
                    if not sourceval:
                        logging.error('Missing required DBF field %s' % mol)
                        sys.exit(1)
                    row[mol] = sourceval
                    polygon[mol] = sourceval
                    logging.info('REQUIRED %s=%s' % (mol, sourceval))

                for source, mol in collection.get_mapping(required=False).iteritems(): #Optional DBF fields
                    sourceval = dbf.get(source)
                    if not sourceval:
                        continue
                    row[mol] = sourceval
                    polygon[mol] = sourceval

                # Write coll_row to collection.csv
                coll_csv.writerow(row)
                layer_polygons.append(polygon)

            # Create JSON representation of dbfjson
            polygons_json = simplejson.dumps(layer_polygons) # TODO: Showing up as string instead of JSON in API
            d=dict(shapefilename=row['layer_filename'], json=polygons_json)
            poly_dw.writerow(dict(shapefilename=row['layer_filename'], json=polygons_json))
        poly_file.flush()
        poly_file.close()

        # Important: Close the DictWriter file before trying to bulkload it
        logging.info('All collection metadata saved to %s' % coll_file.name)
        coll_file.flush()
        coll_file.close()

        # Bulkload...

        if options.dry_run:
            logging.info('Dry run complete!')
            sys.exit(1)

        # os.chdir(current_dir)
        os.chdir('../../')
        filename = os.path.abspath('%s/%s/collection.csv.txt' % (source_dir, coll_dir))
        config_file = os.path.abspath(options.config_file)

        if options.localhost:
            options.url = 'http://localhost:8080/_ah/remote_api'

        # Bulkload Layer entities to App Engine for entire collection
        cmd = "appcfg.py upload_data --config_file=%s --filename=%s --kind=%s --url=%s"
        cmdline = cmd % (config_file, filename, 'Layer', options.url)
        args = shlex.split(cmdline)
        #logging.info(cmdline)
        subprocess.call(args)

        # Bulkload LayerIndex entities to App Engine for entire collection
        cmd = "appcfg.py upload_data --config_file=%s --filename=%s --kind=%s --url=%s"
        cmdline = cmd % (config_file, filename, 'LayerIndex', options.url)
        args = shlex.split(cmdline)
        #logging.info(cmdline)
        subprocess.call(args)

        # Go back to the original directory for the next collection.
        os.chdir(original_dir)

def _getoptions():
    ''' Parses command line options and returns them.'''
    parser = OptionParser()
    parser.add_option('--config_file',
                      type='string',
                      dest='config_file',
                      metavar='FILE',
                      help='Bulkload YAML config file.')
    parser.add_option('-d', '--dry_run',
                      action="store_true",
                      dest='dry_run',
                      help='Creates CSV file but does not bulkload it')
    parser.add_option('-l', '--localhost',
                      action="store_true",
                      dest='localhost',
                      help='Shortcut for bulkloading to http://localhost:8080/_ah/remote_api')
    parser.add_option('-s', '--source_dir',
                      type='string',
                      dest='source_dir',
                      help='Directory containing source to load.')
    parser.add_option('--url',
                      type='string',
                      dest='url',
                      help='URL endpoint to /remote_api to bulkload to.')
    return parser.parse_args()[0]

def main():
    logging.basicConfig(level=logging.DEBUG)
    options = _getoptions()
    current_dir = os.path.curdir
    if options.dry_run:
            logging.info('Performing a dry run...')

    if options.source_dir is not None:
        if os.path.isdir(options.source_dir):
            logging.info('Processing source directory: %s' % options.source_dir)
            source2csv(options.source_dir, options)
            sys.exit(0)
        else:
            logging.info('Unable to locate source directory %s.' % options.source_dir)
            sys.exit(1)
    else:
        source_dirs = [x for x in os.listdir('.') if os.path.isdir(x)]
        logging.info('Processing source directories: %s' % source_dirs)
        for sd in source_dirs: # For each source dir (e.g., jetz, iucn)
            source2csv(sd, options)

    logging.info('Loading finished!')

if __name__ == '__main__':
    main()
<commit_msg>Getting closer on #106<commit_after>#!/usr/bin/env python
#
# Copyright 2011 Aaron Steele and John Wieczorek
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""This script implements the MoL workflow process for layers."""

import copy
import csv
import glob
import logging
from unicodewriterhelper import UnicodeWriter
from optparse import OptionParser
import os
import simplejson
import shlex
import subprocess
import sys
import yaml

class Config(object):
    """Wraps the YAML object for a MoL config.yaml object."""

    @classmethod
    def lower_keys(cls, x):
        """Lower cases all nested dictionary keys."""
        if isinstance(x, list):
            return [cls.lower_keys(v) for v in x]
        if isinstance(x, dict):
            return dict((k.lower(), cls.lower_keys(v)) for k, v in x.iteritems())
        return x

    class Collection(object):
        def __init__(self, collection):
            self.collection = collection

        def __repr__(self):
            return str(self.__dict__)

        def get_row(self):
            row = {}
            for k,v in self.collection['required'].iteritems():
                row[k] = v
            for k,v in self.collection['optional'].iteritems():
                if k in ['taxonomy', 'basemaps']:
                    v = simplejson.dumps(v)
                row[k] = v
            return row

        def get_columns(self):
            cols = []
            cols.extend(self.collection['required'].keys())
            cols.extend(self.collection['optional'].keys())
            cols.extend(self.collection['dbfmapping']['required'].keys())
            cols.extend(self.collection['dbfmapping']['optional'].keys())
            cols.extend(['layer_source', 'layer_collection', 'layer_filename', 'layer_polygons'])
            return cols

        def get_mapping(self, required=True):
            """Returns a reverse dbfmapping for convienience."""
            if required:
                mapping = self.collection['dbfmapping']['required']
            else:
                mapping = self.collection['dbfmapping']['optional']
            return dict((source, mol) for mol,source in mapping.iteritems())

        def getdir(self):
            return self.collection['directoryname']

        def get(self, key, default=None):
            return self.collection.get(key, default)

    def __init__(self, filename):
        self.config = Config.lower_keys(yaml.load(open(filename, 'r').read()))

    def collection_names(self):
        return [x.get('directoryname') for x in self.collections()]

    def collections(self):
        return [Config.Collection(collection) for collection in self.config['collections']]

def source2csv(source_dir, options):
    ''' Loads the collections in the given source directory.

        Arguments:
            source_dir - the relative path to the directory in which the config.yaml file is located.
    '''
    config = Config(os.path.join(source_dir, 'config.yaml'))
    logging.info('Collections in %s: %s' % (source_dir, config.collection_names()))

    collections = config.collections()
    logging.info(len(collections))

    for collection in config.collections(): # For each collection dir in the source dir
        coll_dir = collection.getdir()

        original_dir = os.getcwd()           # We'll need this to restore us to this dir at the end of processing this collection.
        os.chdir(os.path.join(source_dir, coll_dir))

        # Create collection.csv writer
        coll_file = open('collection.csv.txt', 'w')
        coll_cols = collection.get_columns()
        coll_cols.sort()
#        coll_csv = UnicodeWriter(coll_file, coll_cols)
        coll_csv = csv.DictWriter(coll_file, coll_cols)
        coll_csv.writer.writerow(coll_csv.fieldnames)
        coll_row = collection.get_row()
        coll_row['layer_source'] = source_dir
        coll_row['layer_collection'] = coll_dir

        # Create polygons.csv writer
        poly_file = open('collection.polygons.csv.txt', 'w')
#        poly_dw = UnicodeWriter(poly_file, ['shapefilename', 'json'])
        poly_dw = csv.DictWriter(poly_file, ['shapefilename', 'json'])
        poly_dw.writer.writerow(poly_dw.fieldnames)

        # Convert DBF to CSV and add to collection.csv
        shpfiles = glob.glob('*.shp')
        logging.info('Processing %d layers in the %s/%s' % (len(shpfiles), source_dir, coll_dir))
        for sf in shpfiles:
            logging.info('Extracting DBF fields from %s' % sf)
            csvfile = '%s.csv' % sf
            if os.path.exists(csvfile): # ogr2ogr barfs if there are *any* csv files in the dir
                os.remove(csvfile)
            # Following line only for convenience for running on Mac in eclipse
#            command = '/Library/Frameworks/GDAL.framework/Programs/ogr2ogr -f CSV "%s" "%s"' % (csvfile, sf)
            command = 'ogr2ogr -f CSV "%s" "%s"' % (csvfile, sf)
            args = shlex.split(command)
            try:
                subprocess.call(args)
            except OSError as errmsg:
                print """Error occurred while executing command line '{0}': {2}
    Please ensure that {1} is executable and available on your path.
                """.format(command, args[0], errmsg)
                raise

            # Copy and update coll_row with DBF fields
            row = copy.copy(coll_row)
            row['layer_filename'] = os.path.splitext(sf)[0]
            dr = csv.DictReader(open(csvfile, 'r'), skipinitialspace=True)

            layer_polygons = []

            for dbf in dr: # For each row in the DBF CSV file (1 row per polygon)

                polygon = {}

                for source, mol in collection.get_mapping().iteritems(): # Required DBF fields
                    sourceval = dbf.get(source)
                    if not sourceval:
                        logging.error('Missing required DBF field %s' % mol)
                        sys.exit(1)
                    row[mol] = sourceval
                    polygon[mol] = sourceval
                    logging.info('REQUIRED %s=%s' % (mol, sourceval))

                for source, mol in collection.get_mapping(required=False).iteritems(): #Optional DBF fields
                    sourceval = dbf.get(source)
                    if not sourceval:
                        continue
                    row[mol] = sourceval
                    polygon[mol] = sourceval

                # Write coll_row to collection.csv
                coll_csv.writerow(row)
                layer_polygons.append(polygon)

            # Create JSON representation of dbfjson
            polygons_json = simplejson.dumps(layer_polygons) # TODO: Showing up as string instead of JSON in API
            d=dict(shapefilename=row['layer_filename'], json=polygons_json)
            poly_dw.writerow(dict(shapefilename=row['layer_filename'], json=polygons_json))
        poly_file.flush()
        poly_file.close()

        # Important: Close the DictWriter file before trying to bulkload it
        logging.info('All collection metadata saved to %s' % coll_file.name)
        coll_file.flush()
        coll_file.close()

        # Bulkload...

        # os.chdir(current_dir)
        if not options.dry_run:
            os.chdir('../../')
            filename = os.path.abspath('%s/%s/collection.csv.txt' % (source_dir, coll_dir))
            config_file = os.path.abspath(options.config_file)

            if options.localhost:
                options.url = 'http://localhost:8080/_ah/remote_api'

            # Bulkload Layer entities to App Engine for entire collection
            cmd = "appcfg.py upload_data --config_file=%s --filename=%s --kind=%s --url=%s"
            cmdline = cmd % (config_file, filename, 'Layer', options.url)
            args = shlex.split(cmdline)
            subprocess.call(args)

            # Bulkload LayerIndex entities to App Engine for entire collection
            cmd = "appcfg.py upload_data --config_file=%s --filename=%s --kind=%s --url=%s"
            cmdline = cmd % (config_file, filename, 'LayerIndex', options.url)
            args = shlex.split(cmdline)
            subprocess.call(args)

        # Go back to the original directory for the next collection.
        os.chdir(original_dir)

def _getoptions():
    ''' Parses command line options and returns them.'''
    parser = OptionParser()
    parser.add_option('--config_file',
                      type='string',
                      dest='config_file',
                      metavar='FILE',
                      help='Bulkload YAML config file.')
    parser.add_option('-d', '--dry_run',
                      action="store_true",
                      dest='dry_run',
                      help='Creates CSV file but does not bulkload it')
    parser.add_option('-l', '--localhost',
                      action="store_true",
                      dest='localhost',
                      help='Shortcut for bulkloading to http://localhost:8080/_ah/remote_api')
    parser.add_option('-s', '--source_dir',
                      type='string',
                      dest='source_dir',
                      help='Directory containing source to load.')
    parser.add_option('--url',
                      type='string',
                      dest='url',
                      help='URL endpoint to /remote_api to bulkload to.')
    return parser.parse_args()[0]

def main():
    logging.basicConfig(level=logging.DEBUG)
    options = _getoptions()
    current_dir = os.path.curdir
    if options.dry_run:
            logging.info('Performing a dry run...')

    if options.source_dir is not None:
        if os.path.isdir(options.source_dir):
            logging.info('Processing source directory: %s' % options.source_dir)
            source2csv(options.source_dir, options)
            sys.exit(0)
        else:
            logging.info('Unable to locate source directory %s.' % options.source_dir)
            sys.exit(1)
    else:
        source_dirs = [x for x in os.listdir('.') if os.path.isdir(x)]
        logging.info('Processing source directories: %s' % source_dirs)
        for sd in source_dirs: # For each source dir (e.g., jetz, iucn)
            source2csv(sd, options)

    logging.info('Loading finished!')

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># Import the Python Imaging Library if it is available. On error, ignore
# the problem and continue. PIL being absent should only affect the
# graphic lump loading/saving methods and the user may not be interested
# in installing PIL just to pass this line if not interested in using the
# graphics functionality at all.
try:
    from PIL import Image, ImageDraw, ImageOps
except:
    pass

import os
import omg.palette
from omg.util import *
from omg      import six

class Lump(object):
    """Basic lump class. Instances of Lump (and its subclasses)
    always have the following:

        .data       -- a bytes object holding the lump's data
        .from_file  -- load the data to a file
        .to_file    -- save the data to a file

    The default Lump class merely copies the raw data when
    loading/saving to files, but subclasses may convert data
    appropriately (for example, Graphic supports various image
    formats)."""

    def __init__(self, data=None, from_file=None):
        """Create a new instance. The `data` parameter may be a string
        representing data for the lump. The `source` parameter may be
        a path to a file or a file-like object to load from."""
        self.data = bytes()
        if issubclass(type(data), Lump):
            self.data = data.data
        elif data is not None:
            self.data = data or bytes()
        if from_file:
            self.from_file(from_file)

    def from_file(self, source):
        """Load data from a file. Source may be a path name string
        or a file-like object (with a `write` method)."""
        self.data = readfile(source)

    def to_file(self, target):
        """Write data to a file. Target may be a path name string
        or a file-like object (with a `write` method)."""
        writefile(target, self.data)

    def copy(self):
        return deepcopy(self)

class Music(Lump):
    """Subclass of Lump, for music lumps. Not yet implemented."""
    pass

class Sound(Lump):
    """Subclass of Lump, for sound lumps. Not yet implemented."""
    pass

class Graphic(Lump):
    """Subclass of Lump, for Doom format graphics. Supports
    conversion from/to RAWs (sequences of bytes) and PIL
    Image objects, as well as saving to/loading from various
    file formats (via PIL).

    Useful attributes:
        .dimensions     -- (width, height)
        .width          -- width of the image
        .height         -- height of the image
        .x_offset       -- x offset
        .y_offset       -- y offset
    """

    def __init__(self, data=None, from_file=None, palette=None):
        self.palette = palette or omg.palette.default
        Lump.__init__(self, data, from_file)

    def get_offsets(self):
        """Retrieve the (x, y) offsets of the graphic."""
        return unpack('<hh', self.data[4:8])

    def set_offsets(self, xy):
        """Set the (x, y) offsets of the graphic."""
        self.data = self.data[0:4] + pack('<hh', *xy) + self.data[8:]

    def get_dimensions(self):
        """Retrieve the (width, height) dimensions of the graphic."""
        return unpack('<hh', self.data[0:4])

    offsets = property(get_offsets, set_offsets)
    x_offset = property(lambda self: self.offsets[0],
        lambda self, x: self.set_offsets((x, self.y_offset)))
    y_offset = property(lambda self: self.offsets[1],
        lambda self, y: self.set_offsets((self.x_offset, y)))

    dimensions = property(get_dimensions)
    width  = property(lambda self: self.dimensions[0])
    height = property(lambda self: self.dimensions[1])

    def from_pixels(self, data, width, height, x_offset=0, y_offset=0):
        """Load a list of 8bpp pixels.
        Pixels with a negative value are transparent."""
        # First pass: extract pixel data in column+post format
        columns_in = [data[n:width*height:width] for n in range(width)]
        columns_out = []
        for column in columns_in:
            # Find the y position where each chunk starts
            start_rows = []
            postdata = []
            in_trans = True
            tall = False
            offset = 0
            for y in range(height):
                # split at 128 for vanilla-compatible images without premature tiling
                if height < 256:
                    if y == 128:
                        in_trans = True

                # for tall patch support
                elif offset == 254:
                    in_trans = True
                    tall = True
                    # dummy post
                    start_rows.append(254)
                    postdata.append(bytearray())
                    # start relative offsets
                    offset = 0

                if column[y] is None:
                    in_trans = True
                else:
                    if in_trans:
                        # start a new post
                        start_rows.append(offset)
                        postdata.append(bytearray())
                        in_trans = False
                        if tall:
                            # reset relative offset for tall patches
                            offset = 0
                    postdata[-1].append(column[y])

                offset += 1
            columns_out.append(zip(start_rows, postdata))
        # Second pass: compile column+post data, adding pointers
        data = []
        columnptrs = []
        pointer = 4*width + 8
        for column in columns_out:
            columnptrs.append(pack('<i', pointer))
            for row, pixels in column:
                data.append(b"%c%c\x00%s\x00" % (row, len(pixels), pixels))
                pointer += 4 + len(pixels)
            data.append(b'\xff')
            pointer += 1
        # Merge everything together
        self.data = bytes().join([pack('4h', width, height, x_offset, y_offset),
                    bytes().join(columnptrs), bytes().join(data)])

    def from_raw(self, data, width, height, x_offset=0, y_offset=0, pal=None):
        """Load a raw 8-bpp image, converting to the Doom picture format
        (used by all graphics except flats)"""
        pal = pal or omg.palette.default
        pixels = [i if i != pal.tran_index else None for i in six.iterbytes(data)]
        self.from_pixels(pixels, width, height, x_offset, y_offset)

    def to_pixels(self):
        """Returns self converted to a list of 8bpp pixels.
        Pixels with value None are transparent."""
        data = self.data
        width, height = self.dimensions
        output = [None] * (width*height)
        pointers = unpack('<%il'%width, data[8 : 8 + width*4])
        for x in range(width):
            y = -1
            pointer = pointers[x]
            if pointer >= len(data):
                continue

            while six.indexbytes(data, pointer) != 0xff:
                offset = six.indexbytes(data, pointer)
                if offset <= y:
                    y += offset # for tall patches
                else:
                    y = offset
                post_length = six.indexbytes(data, pointer+1)
                op = y*width + x
                for p in range(pointer + 3, pointer + post_length + 3):
                    if op >= len(output) or p >= len(data):
                        break
                    output[op] = six.indexbytes(data, p)
                    op += width
                pointer += post_length + 4
        return output

    def to_raw(self, tran_index=None):
        """Returns self converted to a raw (8-bpp) image.

        `tran_index` specifies the palette index to use for
        transparent pixels. The value defaults to that of the
        Graphic object's palette instance."""
        tran_index = tran_index or self.palette.tran_index
        output = [i or tran_index for i in self.to_pixels()]
        return bytes(bytearray(output))

    def to_Image(self, mode='P'):
        """Convert to a PIL Image instance"""
        if mode != 'RGBA' or isinstance(self, Flat):
            # target image has no alpha,
            # or source image is a flat (which has no transparent pixels)
            im = Image.new('P', self.dimensions, None)
            if isinstance(self, Flat):
                im.frombytes(self.data)
            else:
                im.frombytes(self.to_raw())
            im.putpalette(self.palette.save_bytes)
            return im.convert(mode)
        else:
            # target image is RGBA and source image is not a flat
            im = Image.new('RGBA', self.dimensions, None)
            data = bytes().join([self.palette.bytes[i*3:i*3+3] + b'\xff' if i is not None \
                                 else b'\0\0\0\0' for i in self.to_pixels()])
            im.frombytes(data)
            return im

    def from_Image(self, im, translate=False):
        """Load from a PIL Image instance

        If the input image is 24-bit or 32-bit, the colors will be
        looked up in the current palette.

        If the input image is 8-bit, indices will simply be copied
        from the input image. To properly translate colors between
        palettes, set the `translate` parameter."""

        pixels = im.tobytes()
        width, height = im.size
        xoff, yoff = (width // 2)-1, height-5
        if im.mode == "RGB":
            pixels = join([chr(self.palette.match(unpack('BBB', \
                pixels[i*3:(i+1)*3]))) for i in range(width*height)])

            self.from_raw(pixels, width, height, xoff, yoff, self.palette)

        elif im.mode == "RGBA":
            pixels = [unpack('BBBB', pixels[i*4:(i+1)*4]) for i in range(width*height)]
            pixels = [self.palette.match(i[0:3]) if i[3] > 0 else None for i in pixels]

            self.from_pixels(pixels, width, height, xoff, yoff)

        elif im.mode == 'P':
            srcpal = im.palette.tobytes()
            if im.palette.mode == "RGB":
                palsize = 3
            elif im.palette.mode == "RGBA":
                palsize = 4
            else:
                raise TypeError("palette mode must be 'RGB' or 'RGBA'")

            if translate:
                R = [c for c in six.iterbytes(srcpal[0::palsize])]
                G = [c for c in six.iterbytes(srcpal[1::palsize])]
                B = [c for c in six.iterbytes(srcpal[2::palsize])]

                srcpal = zip(R, G, B)
                lexicon = [six.int2byte(self.palette.match(c)) for c in srcpal]
                pixels = join([lexicon[b] for b in six.iterbytes(pixels)])
            else:
                # Simply copy pixels. However, make sure to translate
                # all colors matching the transparency color to the
                # right index. This is necessary because programs
                # aren't consistent in choice of position for the
                # transparent entry.
                packed_color = pack("BBB", *self.palette.tran_color)
                ri = 0
                while ri != -1:
                    ri = srcpal.find(packed_color, ri+palsize)
                    if not ri % palsize and ri//palsize != self.palette.tran_index:
                        pixels = pixels.replace(six.int2byte(ri//palsize),
                            six.int2byte(self.palette.tran_index))

            self.from_raw(pixels, width, height, xoff, yoff, self.palette)
        else:
            raise TypeError("image mode must be 'P', 'RGB', or 'RGBA'")

    def from_file(self, filename, translate=False):
        """Load graphic from an image file."""
        if filename[-4:].lower() == '.lmp':
            self.data = readfile(filename)
        else:
            im = Image.open(filename)
            self.from_Image(im, translate)

    def to_file(self, filename, mode='P'):
        """Save the graphic to an image file.

        The output format is selected based on the filename extension.
        For example, "file.jpg" saves to JPEG format. If the file has
        no extension, PNG format is used.

        Special cases: ".lmp" saves the raw lump data, and ".raw" saves
        the raw pixel data.

        `mode` may be be 'P', 'RGB', or 'RGBA' for palette or 24/32 bit
        output, respectively. However, .raw ignores this parameter and
        always writes in palette mode."""

        format = os.path.splitext(filename)[1][1:].upper()
        if   format == 'LMP': writefile(filename, self.data)
        elif format == 'RAW': writefile(filename, self.to_raw())
        else:
            im = self.to_Image(mode)
            if format:
                im.save(filename)
            else:
                im.save(filename, "PNG")

    def translate(self, pal):
        """Translate (in-place) the graphic to another palette."""
        lexicon = [chr(pal.match(self.palette.colors[i])) for i in range(256)]
        lexicon[self.palette.tran_index] = chr(pal.tran_index)
        if isinstance(self, Flat):
            self.data = join([lexicon[b] for b in self.data])
        else:
            raw = self.to_raw()
            #raw = raw.replace(chr(self.palette.tran_index), chr(pal.tran_index))
            self.load_raw(join([lexicon[b] for b in raw]),
                self.width, self.height,
                self.x_offset, self.y_offset)

class Flat(Graphic):
    """Subclass of Graphic, for flat graphics"""

    def get_dimensions(self):
        sz = len(self.data)
        if sz == 4096: return (64, 64)
        if sz == 4160: return (64, 65)
        root = int(sz**0.5)
        if root**2 != sz:
            raise TypeError("unable to determine size: not a square number")
        return (root, root)

    dimensions = property(get_dimensions)
    width  = property(lambda self: self.dimensions[0])
    height = property(lambda self: self.dimensions[1])

    def load_raw(self, data, *unused):
        self.data = data

    def to_raw(self):
        return self.data
<commit_msg>fix some leftover usage of chr() for py3<commit_after># Import the Python Imaging Library if it is available. On error, ignore
# the problem and continue. PIL being absent should only affect the
# graphic lump loading/saving methods and the user may not be interested
# in installing PIL just to pass this line if not interested in using the
# graphics functionality at all.
try:
    from PIL import Image, ImageDraw, ImageOps
except:
    pass

import os
import omg.palette
from omg.util import *
from omg      import six

class Lump(object):
    """Basic lump class. Instances of Lump (and its subclasses)
    always have the following:

        .data       -- a bytes object holding the lump's data
        .from_file  -- load the data to a file
        .to_file    -- save the data to a file

    The default Lump class merely copies the raw data when
    loading/saving to files, but subclasses may convert data
    appropriately (for example, Graphic supports various image
    formats)."""

    def __init__(self, data=None, from_file=None):
        """Create a new instance. The `data` parameter may be a string
        representing data for the lump. The `source` parameter may be
        a path to a file or a file-like object to load from."""
        self.data = bytes()
        if issubclass(type(data), Lump):
            self.data = data.data
        elif data is not None:
            self.data = data or bytes()
        if from_file:
            self.from_file(from_file)

    def from_file(self, source):
        """Load data from a file. Source may be a path name string
        or a file-like object (with a `write` method)."""
        self.data = readfile(source)

    def to_file(self, target):
        """Write data to a file. Target may be a path name string
        or a file-like object (with a `write` method)."""
        writefile(target, self.data)

    def copy(self):
        return deepcopy(self)

class Music(Lump):
    """Subclass of Lump, for music lumps. Not yet implemented."""
    pass

class Sound(Lump):
    """Subclass of Lump, for sound lumps. Not yet implemented."""
    pass

class Graphic(Lump):
    """Subclass of Lump, for Doom format graphics. Supports
    conversion from/to RAWs (sequences of bytes) and PIL
    Image objects, as well as saving to/loading from various
    file formats (via PIL).

    Useful attributes:
        .dimensions     -- (width, height)
        .width          -- width of the image
        .height         -- height of the image
        .x_offset       -- x offset
        .y_offset       -- y offset
    """

    def __init__(self, data=None, from_file=None, palette=None):
        self.palette = palette or omg.palette.default
        Lump.__init__(self, data, from_file)

    def get_offsets(self):
        """Retrieve the (x, y) offsets of the graphic."""
        return unpack('<hh', self.data[4:8])

    def set_offsets(self, xy):
        """Set the (x, y) offsets of the graphic."""
        self.data = self.data[0:4] + pack('<hh', *xy) + self.data[8:]

    def get_dimensions(self):
        """Retrieve the (width, height) dimensions of the graphic."""
        return unpack('<hh', self.data[0:4])

    offsets = property(get_offsets, set_offsets)
    x_offset = property(lambda self: self.offsets[0],
        lambda self, x: self.set_offsets((x, self.y_offset)))
    y_offset = property(lambda self: self.offsets[1],
        lambda self, y: self.set_offsets((self.x_offset, y)))

    dimensions = property(get_dimensions)
    width  = property(lambda self: self.dimensions[0])
    height = property(lambda self: self.dimensions[1])

    def from_pixels(self, data, width, height, x_offset=0, y_offset=0):
        """Load a list of 8bpp pixels.
        Pixels with a negative value are transparent."""
        # First pass: extract pixel data in column+post format
        columns_in = [data[n:width*height:width] for n in range(width)]
        columns_out = []
        for column in columns_in:
            # Find the y position where each chunk starts
            start_rows = []
            postdata = []
            in_trans = True
            tall = False
            offset = 0
            for y in range(height):
                # split at 128 for vanilla-compatible images without premature tiling
                if height < 256:
                    if y == 128:
                        in_trans = True

                # for tall patch support
                elif offset == 254:
                    in_trans = True
                    tall = True
                    # dummy post
                    start_rows.append(254)
                    postdata.append(bytearray())
                    # start relative offsets
                    offset = 0

                if column[y] is None:
                    in_trans = True
                else:
                    if in_trans:
                        # start a new post
                        start_rows.append(offset)
                        postdata.append(bytearray())
                        in_trans = False
                        if tall:
                            # reset relative offset for tall patches
                            offset = 0
                    postdata[-1].append(column[y])

                offset += 1
            columns_out.append(zip(start_rows, postdata))
        # Second pass: compile column+post data, adding pointers
        data = []
        columnptrs = []
        pointer = 4*width + 8
        for column in columns_out:
            columnptrs.append(pack('<i', pointer))
            for row, pixels in column:
                data.append(b"%c%c\x00%s\x00" % (row, len(pixels), pixels))
                pointer += 4 + len(pixels)
            data.append(b'\xff')
            pointer += 1
        # Merge everything together
        self.data = bytes().join([pack('4h', width, height, x_offset, y_offset),
                    bytes().join(columnptrs), bytes().join(data)])

    def from_raw(self, data, width, height, x_offset=0, y_offset=0, pal=None):
        """Load a raw 8-bpp image, converting to the Doom picture format
        (used by all graphics except flats)"""
        pal = pal or omg.palette.default
        pixels = [i if i != pal.tran_index else None for i in six.iterbytes(data)]
        self.from_pixels(pixels, width, height, x_offset, y_offset)

    def to_pixels(self):
        """Returns self converted to a list of 8bpp pixels.
        Pixels with value None are transparent."""
        data = self.data
        width, height = self.dimensions
        output = [None] * (width*height)
        pointers = unpack('<%il'%width, data[8 : 8 + width*4])
        for x in range(width):
            y = -1
            pointer = pointers[x]
            if pointer >= len(data):
                continue

            while six.indexbytes(data, pointer) != 0xff:
                offset = six.indexbytes(data, pointer)
                if offset <= y:
                    y += offset # for tall patches
                else:
                    y = offset
                post_length = six.indexbytes(data, pointer+1)
                op = y*width + x
                for p in range(pointer + 3, pointer + post_length + 3):
                    if op >= len(output) or p >= len(data):
                        break
                    output[op] = six.indexbytes(data, p)
                    op += width
                pointer += post_length + 4
        return output

    def to_raw(self, tran_index=None):
        """Returns self converted to a raw (8-bpp) image.

        `tran_index` specifies the palette index to use for
        transparent pixels. The value defaults to that of the
        Graphic object's palette instance."""
        tran_index = tran_index or self.palette.tran_index
        output = [i or tran_index for i in self.to_pixels()]
        return bytes(bytearray(output))

    def to_Image(self, mode='P'):
        """Convert to a PIL Image instance"""
        if mode != 'RGBA' or isinstance(self, Flat):
            # target image has no alpha,
            # or source image is a flat (which has no transparent pixels)
            im = Image.new('P', self.dimensions, None)
            if isinstance(self, Flat):
                im.frombytes(self.data)
            else:
                im.frombytes(self.to_raw())
            im.putpalette(self.palette.save_bytes)
            return im.convert(mode)
        else:
            # target image is RGBA and source image is not a flat
            im = Image.new('RGBA', self.dimensions, None)
            data = bytes().join([self.palette.bytes[i*3:i*3+3] + b'\xff' if i is not None \
                                 else b'\0\0\0\0' for i in self.to_pixels()])
            im.frombytes(data)
            return im

    def from_Image(self, im, translate=False):
        """Load from a PIL Image instance

        If the input image is 24-bit or 32-bit, the colors will be
        looked up in the current palette.

        If the input image is 8-bit, indices will simply be copied
        from the input image. To properly translate colors between
        palettes, set the `translate` parameter."""

        pixels = im.tobytes()
        width, height = im.size
        xoff, yoff = (width // 2)-1, height-5
        if im.mode == "RGB":
            pixels = join([six.int2byte(self.palette.match(unpack('BBB', \
                pixels[i*3:(i+1)*3]))) for i in range(width*height)])

            self.from_raw(pixels, width, height, xoff, yoff, self.palette)

        elif im.mode == "RGBA":
            pixels = [unpack('BBBB', pixels[i*4:(i+1)*4]) for i in range(width*height)]
            pixels = [self.palette.match(i[0:3]) if i[3] > 0 else None for i in pixels]

            self.from_pixels(pixels, width, height, xoff, yoff)

        elif im.mode == 'P':
            srcpal = im.palette.tobytes()
            if im.palette.mode == "RGB":
                palsize = 3
            elif im.palette.mode == "RGBA":
                palsize = 4
            else:
                raise TypeError("palette mode must be 'RGB' or 'RGBA'")

            if translate:
                R = [c for c in six.iterbytes(srcpal[0::palsize])]
                G = [c for c in six.iterbytes(srcpal[1::palsize])]
                B = [c for c in six.iterbytes(srcpal[2::palsize])]

                srcpal = zip(R, G, B)
                lexicon = [six.int2byte(self.palette.match(c)) for c in srcpal]
                pixels = join([lexicon[b] for b in six.iterbytes(pixels)])
            else:
                # Simply copy pixels. However, make sure to translate
                # all colors matching the transparency color to the
                # right index. This is necessary because programs
                # aren't consistent in choice of position for the
                # transparent entry.
                packed_color = pack("BBB", *self.palette.tran_color)
                ri = 0
                while ri != -1:
                    ri = srcpal.find(packed_color, ri+palsize)
                    if not ri % palsize and ri//palsize != self.palette.tran_index:
                        pixels = pixels.replace(six.int2byte(ri//palsize),
                            six.int2byte(self.palette.tran_index))

            self.from_raw(pixels, width, height, xoff, yoff, self.palette)
        else:
            raise TypeError("image mode must be 'P', 'RGB', or 'RGBA'")

    def from_file(self, filename, translate=False):
        """Load graphic from an image file."""
        if filename[-4:].lower() == '.lmp':
            self.data = readfile(filename)
        else:
            im = Image.open(filename)
            self.from_Image(im, translate)

    def to_file(self, filename, mode='P'):
        """Save the graphic to an image file.

        The output format is selected based on the filename extension.
        For example, "file.jpg" saves to JPEG format. If the file has
        no extension, PNG format is used.

        Special cases: ".lmp" saves the raw lump data, and ".raw" saves
        the raw pixel data.

        `mode` may be be 'P', 'RGB', or 'RGBA' for palette or 24/32 bit
        output, respectively. However, .raw ignores this parameter and
        always writes in palette mode."""

        format = os.path.splitext(filename)[1][1:].upper()
        if   format == 'LMP': writefile(filename, self.data)
        elif format == 'RAW': writefile(filename, self.to_raw())
        else:
            im = self.to_Image(mode)
            if format:
                im.save(filename)
            else:
                im.save(filename, "PNG")

    def translate(self, pal):
        """Translate (in-place) the graphic to another palette."""
        lexicon = [six.int2byte(pal.match(self.palette.colors[i])) for i in range(256)]
        lexicon[self.palette.tran_index] = six.int2byte(pal.tran_index)
        if isinstance(self, Flat):
            self.data = join([lexicon[b] for b in self.data])
        else:
            raw = self.to_raw()
            #raw = raw.replace(six.int2byte(self.palette.tran_index), six.int2byte(pal.tran_index))
            self.load_raw(join([lexicon[b] for b in raw]),
                self.width, self.height,
                self.x_offset, self.y_offset)

class Flat(Graphic):
    """Subclass of Graphic, for flat graphics"""

    def get_dimensions(self):
        sz = len(self.data)
        if sz == 4096: return (64, 64)
        if sz == 4160: return (64, 65)
        root = int(sz**0.5)
        if root**2 != sz:
            raise TypeError("unable to determine size: not a square number")
        return (root, root)

    dimensions = property(get_dimensions)
    width  = property(lambda self: self.dimensions[0])
    height = property(lambda self: self.dimensions[1])

    def load_raw(self, data, *unused):
        self.data = data

    def to_raw(self):
        return self.data
<|endoftext|>


<commit_before>#!/usr/bin/env python2.7

from __future__ import absolute_import
import redis
import json

#POOL = redis.ConnectionPool(max_connections=4, host='localhost', db=5, port=6379)
#REDIS_CLIENT = redis.Redis(connection_pool=POOL)
data = json.load(open('config.json'))
REDIS_CLIENT = redis.Redis(host=data['oohost'], db=data['oodb'], port=data['ooport'])

def only_one(function=None, key="", timeout=None):
    """Enforce only one celery task at a time."""

    def _dec(run_func):
        """Decorator."""

        def _caller(*args, **kwargs):
            """Caller."""
            ret_value = None
            have_lock = False
            lock = REDIS_CLIENT.lock(key, timeout=timeout)
            try:
                have_lock = lock.acquire(blocking=False)
                if have_lock:
                    ret_value = run_func(*args, **kwargs)
            finally:
                if have_lock:
                    lock.release()

            return ret_value

        return _caller

    return _dec(function) if function is not None else _dec
<commit_msg>Set full static path for config.json file.  Necessary for launching manually<commit_after>#!/usr/bin/env python2.7

from __future__ import absolute_import
import redis
import json
import os

#POOL = redis.ConnectionPool(max_connections=4, host='localhost', db=5, port=6379)
#REDIS_CLIENT = redis.Redis(connection_pool=POOL)
configfile = "%s/config.json" % (os.path.dirname(os.path.realpath(__file__)))
data = json.load(open(configfile))
REDIS_CLIENT = redis.Redis(host=data['oohost'], db=data['oodb'], port=data['ooport'])

def only_one(function=None, key="", timeout=None):
    """Enforce only one celery task at a time."""

    def _dec(run_func):
        """Decorator."""

        def _caller(*args, **kwargs):
            """Caller."""
            ret_value = None
            have_lock = False
            lock = REDIS_CLIENT.lock(key, timeout=timeout)
            try:
                have_lock = lock.acquire(blocking=False)
                if have_lock:
                    ret_value = run_func(*args, **kwargs)
            finally:
                if have_lock:
                    lock.release()

            return ret_value

        return _caller

    return _dec(function) if function is not None else _dec
<|endoftext|>