<commit_before>#!/usr/bin/env python

import pywren
import boto3
import click
import shutil
import os
import json
import zipfile
from glob2 import glob
import io
import time
import botocore
from multiprocessing import Process
from pywren import wrenhandler
import logging
import watchtower
import subprocess
import math
import platform
try:
    # For Python 3.0 and later
    from urllib.request import urlopen
except ImportError:
    # Fall back to Python 2's urllib2
    from urllib2 import urlopen

logger = logging.getLogger(__name__)

SQS_VISIBILITY_SEC = 10
PROCESS_SLEEP_DUR_SEC=2
AWS_REGION_DEBUG='us-west-2'
QUEUE_SLEEP_DUR_SEC=2
IDLE_TERMINATE_THRESHOLD = 0.95

INSTANCE_ID_URL = "http://169.254.169.254/latest/meta-data/instance-id"
def get_my_ec2_instance(aws_region):

    ec2 = boto3.resource('ec2', region_name=aws_region)

    instance_id =  urlopen(INSTANCE_ID_URL).read()
    instances = ec2.instances.filter(InstanceIds=[instance_id])

    for instance in instances:
        return instance

# def get_my_ec2_uptime():
#     instance = get_my_ec2_instance()

#     launch_time = instance.launch_time
#     time_delta =  datetime.datetime.now(launch_time.tzinfo) - launch_time
#     print launch_time, time_delta
#     #hour_frac = (time_delta.total_seconds() % 3600) / 3600

#     return time_delta.total_seconds()

def tags_to_dict(d):
    if d is None:
        return {}
    return {a['Key'] : a['Value'] for a in d}

def get_my_ec2_meta(instance):

    tags = tags_to_dict(instance.tags)

    r = {'public_dns_name' : instance.public_dns_name,
         'public_ip_address' : instance.public_ip_address,
         'instance_id': instance.id}
    r.update(tags)
    return r

def get_my_uptime():
    with open('/proc/uptime', 'r') as f:
        uptime_seconds = float(f.readline().split()[0])
        return uptime_seconds

def check_is_ec2():
    """
    last-minute check to make sure we are on EC2.

    """
    try:
        instance_id =  urlopen(INSTANCE_ID_URL, timeout=3).read()
        return True
    except:
        return False

def ec2_self_terminate(idle_time, uptime, message_count):
    if check_is_ec2():
        logger.info("self-terminating after idle for {:.0f} sec ({:.0f} s uptime), processed {:d} messages".format(idle_time, uptime, message_count))
        for h in logger.handlers:
            h.flush()

        subprocess.call("sudo shutdown -h now", shell=True)
    else:
        logger.warn("attempted to self-terminate on non-EC2 instance. Check config")

def idle_granularity_valid(idle_terminate_granularity,
                           queue_receive_message_timeout):
    return (1.0 - IDLE_TERMINATE_THRESHOLD)*idle_terminate_granularity >  (queue_receive_message_timeout)*1.1

def server_runner(aws_region, sqs_queue_name,
                  max_run_time, run_dir, server_name, log_stream_prefix,
                  max_idle_time=None,
                  idle_terminate_granularity = None,
                  queue_receive_message_timeout=10):
    """
    Extract messages from queue and pass them off
    """

    sqs = boto3.resource('sqs', region_name=aws_region)

    # Get the queue
    queue = sqs.get_queue_by_name(QueueName=sqs_queue_name)
    local_message_i = 0
    last_processed_timestamp = time.time()

    terminate_thold_sec = (IDLE_TERMINATE_THRESHOLD * idle_terminate_granularity)
    terminate_window_sec = idle_terminate_granularity - terminate_thold_sec
    queue_receive_message_timeout = min(math.floor(terminate_window_sec/1.2), queue_receive_message_timeout)
    queue_receive_message_timeout = int(max(queue_receive_message_timeout, 1))
    if not idle_granularity_valid(idle_terminate_granularity,
                              queue_receive_message_timeout):
        raise Exception("Idle time granularity window smaller than queue receive message timeout with headroom, instance will not self-terminate")
    message_count = 0
    idle_time = 0
    while(True):
        logger.debug("reading queue" )
        response = queue.receive_messages(WaitTimeSeconds=queue_receive_message_timeout)
        if len(response) > 0:
            m = response[0]
            logger.info("Dispatching message_id={}".format(m.message_id))

            process_message(m, local_message_i, max_run_time, run_dir,
                            aws_region, server_name, log_stream_prefix)
            message_count += 1
            last_processed_timestamp = time.time()
            idle_time = 0
        else:
            idle_time = time.time() - last_processed_timestamp

            logger.debug("no message, idle for {:3.0f} sec".format(idle_time))

        # this is EC2_only
        if max_idle_time is not None and \
           idle_terminate_granularity is not None:
            if idle_time > max_idle_time:
                my_uptime = get_my_uptime()
                time_frac = (my_uptime % idle_terminate_granularity)

                logger.debug("Instance has been up for {:.0f} and inactive for {:.0f} time_frac={:.0f} terminate_thold={:.0f}".format(my_uptime,
                                                                                idle_time,
                                                                                time_frac,
                                                                                terminate_thold_sec))

                if time_frac > terminate_thold_sec:
                    logger.info("Instance has been up for {:.0f} and inactive for {:.0f}, terminating".format(my_uptime,
                                                                                idle_time))
                    for h in logger.handlers:
                        h.flush()
                    ec2_self_terminate(idle_time, my_uptime, message_count)

def process_message(m, local_message_i, max_run_time, run_dir,
                    aws_region,
                    server_name, log_stream_prefix):

    logger.info("processing message_id={}".format(m.message_id))

    event = json.loads(m.body)
    call_id = event['call_id']
    callset_id = event['callset_id']

    extra_env_debug = event.get('extra_env', {})
    unique_job_id = "{}:{}:{}".format(m.message_id, call_id, callset_id)
    logger.info("processing message_id={} callset_id={} call_id={}".format(m.message_id, callset_id, call_id))

    # FIXME this is all for debugging
    if 'DEBUG_THROW_EXCEPTION' in extra_env_debug:
        m.delete()
        raise Exception("Debug exception")
    message_id = m.message_id

    # run this in a thread: pywren.wrenhandler.generic_handler(event)
    p =  Process(target=job_handler, args=(event, local_message_i,
                                           run_dir, aws_region, server_name,
                                           log_stream_prefix))
    # is thread done
    p.start()
    pid = p.pid
    logger.info("processing message_id={} callset_id={} call_id={} process_pid={}".format(m.message_id, callset_id, call_id, pid))
    start_time = time.time()

    run_time = time.time() - start_time
    last_visibility_update_time = time.time()
    while run_time < max_run_time:
        time_since_visibility_update = time.time() - last_visibility_update_time
        est_visibility_left = SQS_VISIBILITY_SEC - time_since_visibility_update
        if est_visibility_left < (PROCESS_SLEEP_DUR_SEC*1.5):
            logger.debug("{} - {:3.1f}s since last visibility update, setting to {:3.1f} sec".format(message_id, time_since_visibility_update,
                                                                                SQS_VISIBILITY_SEC))
            last_visibility_update_time = time.time()
            response = m.change_visibility(VisibilityTimeout=SQS_VISIBILITY_SEC)

        if not p.is_alive():
            logger.debug("{} - attempting to join process".format(message_id))
            # FIXME will this join ever hang?
            p.join()
            break
        else:
            logger.debug("{} - {:3.1f}s since visibility update, sleeping".format(message_id, time_since_visibility_update))
            time.sleep(PROCESS_SLEEP_DUR_SEC)

        run_time = time.time() - start_time

    if p.exitcode is None:
        p.terminate()  # FIXME PRINT LOTS OF ERRORS HERE
        logger.warn("{} - Had to manually terminate process ".format(message_id))
    logger.info("deleting message_id={} callset_id={} call_id={}".format(m.message_id, callset_id, call_id))

    m.delete()

def copy_runtime(tgt_dir):
    files = glob(os.path.join(pywren.SOURCE_DIR, "./*.py"))
    for f in files:
        shutil.copy(f, os.path.join(tgt_dir, os.path.basename(f)))

def job_handler(job, job_i, run_dir, aws_region,
                server_name, log_stream_prefix,
                extra_context = None,
                delete_taskdir=True):
    """
    Run a deserialized job in run_dir

    Just for debugging
    """
    debug_pid = open("/tmp/pywren.scripts.standalone.{}.{}.log".format(os.getpid(),
                                                                       time.time()), 'w')
    print "subprocess job_handler job i=", job_i, "pid=", os.getpid()
    session = boto3.session.Session(region_name=aws_region)
    # we do this here instead of in the global context
    # because of how multiprocessing works
    handler = watchtower.CloudWatchLogHandler(send_interval=20,
                                              log_group="pywren.standalone",
                                              stream_name=log_stream_prefix + "-{logger_name}",
                                              boto3_session=session,
                                              max_batch_count=10)
    log_format_str ='{} %(asctime)s - %(name)s - %(levelname)s - %(message)s'.format(server_name)

    formatter = logging.Formatter(log_format_str, "%Y-%m-%d %H:%M:%S")
    handler.setFormatter(formatter)

    wren_log = pywren.wrenhandler.logger # logging.getLogger('pywren.wrenhandler')
    wren_log.setLevel(logging.DEBUG)
    wren_log.propagate = 0
    wren_log.addHandler(handler)

    original_dir = os.getcwd()

    task_run_dir = os.path.join(run_dir, str(job_i))
    shutil.rmtree(task_run_dir, True) # delete old modules
    os.makedirs(task_run_dir)
    copy_runtime(task_run_dir)

    context = {'jobnum' : job_i}
    if extra_context is not None:
        context.update(extra_context)

    os.chdir(task_run_dir)
    try:
        wrenhandler.generic_handler(job, context)
    finally:
        if delete_taskdir:
            shutil.rmtree(task_run_dir)
        os.chdir(original_dir)
    handler.flush()

@click.command()
@click.option('--max_run_time', default=3600,
              help='max run time for a job', type=int)
@click.option('--run_dir', default="/tmp/pywren.rundir",
              help='directory to hold intermediate output')
@click.option('--aws_region', default="us-west-2",
              help='aws region')
@click.option('--sqs_queue_name', default="pywren-queue",
              help='queue')
@click.option('--max_idle_time', default=None, type=int,
              help='maximum time for queue to remine idle before we try to self-terminate (sec)')
@click.option('--idle_terminate_granularity', default=None, type=int,
              help="only terminate if we have been up for an integral number of this")
@click.option('--queue_receive_message_timeout', default=10, type=int,
              help="longpoll timeout for getting sqs messages")
def server(aws_region, max_run_time, run_dir, sqs_queue_name, max_idle_time,
           idle_terminate_granularity, queue_receive_message_timeout):

    session = boto3.session.Session(region_name=aws_region)

    # make boto quiet locally FIXME is there a better way of doing this?
    logging.getLogger('boto').setLevel(logging.CRITICAL)
    logging.getLogger('boto3').setLevel(logging.CRITICAL)
    logging.getLogger('botocore').setLevel(logging.CRITICAL)

    if platform.node() != 'c65':

        instance = get_my_ec2_instance(aws_region)
        ec2_metadata = get_my_ec2_meta(instance)
        server_name = ec2_metadata['Name']
        log_stream_prefix = ec2_metadata['instance_id']
    else:
        server_name='c65'
        log_stream_prefix='millennium-c65'

    log_format_str ='{} %(asctime)s - %(name)s - %(levelname)s - %(message)s'.format(server_name)

    formatter = logging.Formatter(log_format_str, "%Y-%m-%d %H:%M:%S")

    handler = watchtower.CloudWatchLogHandler(send_interval=20,
                                              log_group="pywren.standalone",
                                              stream_name=log_stream_prefix + "-{logger_name}",
                                              boto3_session=session,
                                              max_batch_count=10)

    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)

    #config = pywren.wrenconfig.default()
    server_runner(aws_region, sqs_queue_name,
                  max_run_time, os.path.abspath(run_dir),
                  server_name, log_stream_prefix,
                  max_idle_time,
                  idle_terminate_granularity,
                  queue_receive_message_timeout)

<commit_msg>updates to file writing logging<commit_after>#!/usr/bin/env python

import pywren
import boto3
import click
import shutil
import os
import json
import zipfile
from glob2 import glob
import io
import time
import botocore
from multiprocessing import Process
from pywren import wrenhandler
import logging
import watchtower
import subprocess
import math
import platform
try:
    # For Python 3.0 and later
    from urllib.request import urlopen
except ImportError:
    # Fall back to Python 2's urllib2
    from urllib2 import urlopen

logger = logging.getLogger(__name__)

SQS_VISIBILITY_SEC = 10
PROCESS_SLEEP_DUR_SEC=2
AWS_REGION_DEBUG='us-west-2'
QUEUE_SLEEP_DUR_SEC=2
IDLE_TERMINATE_THRESHOLD = 0.95

INSTANCE_ID_URL = "http://169.254.169.254/latest/meta-data/instance-id"
def get_my_ec2_instance(aws_region):

    ec2 = boto3.resource('ec2', region_name=aws_region)

    instance_id =  urlopen(INSTANCE_ID_URL).read()
    instances = ec2.instances.filter(InstanceIds=[instance_id])

    for instance in instances:
        return instance

# def get_my_ec2_uptime():
#     instance = get_my_ec2_instance()

#     launch_time = instance.launch_time
#     time_delta =  datetime.datetime.now(launch_time.tzinfo) - launch_time
#     print launch_time, time_delta
#     #hour_frac = (time_delta.total_seconds() % 3600) / 3600

#     return time_delta.total_seconds()

def tags_to_dict(d):
    if d is None:
        return {}
    return {a['Key'] : a['Value'] for a in d}

def get_my_ec2_meta(instance):

    tags = tags_to_dict(instance.tags)

    r = {'public_dns_name' : instance.public_dns_name,
         'public_ip_address' : instance.public_ip_address,
         'instance_id': instance.id}
    r.update(tags)
    return r

def get_my_uptime():
    with open('/proc/uptime', 'r') as f:
        uptime_seconds = float(f.readline().split()[0])
        return uptime_seconds

def check_is_ec2():
    """
    last-minute check to make sure we are on EC2.

    """
    try:
        instance_id =  urlopen(INSTANCE_ID_URL, timeout=3).read()
        return True
    except:
        return False

def ec2_self_terminate(idle_time, uptime, message_count):
    if check_is_ec2():
        logger.info("self-terminating after idle for {:.0f} sec ({:.0f} s uptime), processed {:d} messages".format(idle_time, uptime, message_count))
        for h in logger.handlers:
            h.flush()

        subprocess.call("sudo shutdown -h now", shell=True)
    else:
        logger.warn("attempted to self-terminate on non-EC2 instance. Check config")

def idle_granularity_valid(idle_terminate_granularity,
                           queue_receive_message_timeout):
    return (1.0 - IDLE_TERMINATE_THRESHOLD)*idle_terminate_granularity >  (queue_receive_message_timeout)*1.1

def server_runner(aws_region, sqs_queue_name,
                  max_run_time, run_dir, server_name, log_stream_prefix,
                  max_idle_time=None,
                  idle_terminate_granularity = None,
                  queue_receive_message_timeout=10):
    """
    Extract messages from queue and pass them off
    """

    sqs = boto3.resource('sqs', region_name=aws_region)

    # Get the queue
    queue = sqs.get_queue_by_name(QueueName=sqs_queue_name)
    local_message_i = 0
    last_processed_timestamp = time.time()

    terminate_thold_sec = (IDLE_TERMINATE_THRESHOLD * idle_terminate_granularity)
    terminate_window_sec = idle_terminate_granularity - terminate_thold_sec
    queue_receive_message_timeout = min(math.floor(terminate_window_sec/1.2), queue_receive_message_timeout)
    queue_receive_message_timeout = int(max(queue_receive_message_timeout, 1))
    if not idle_granularity_valid(idle_terminate_granularity,
                              queue_receive_message_timeout):
        raise Exception("Idle time granularity window smaller than queue receive message timeout with headroom, instance will not self-terminate")
    message_count = 0
    idle_time = 0
    while(True):
        logger.debug("reading queue" )
        response = queue.receive_messages(WaitTimeSeconds=queue_receive_message_timeout)
        if len(response) > 0:
            m = response[0]
            logger.info("Dispatching message_id={}".format(m.message_id))

            process_message(m, local_message_i, max_run_time, run_dir,
                            aws_region, server_name, log_stream_prefix)
            message_count += 1
            last_processed_timestamp = time.time()
            idle_time = 0
        else:
            idle_time = time.time() - last_processed_timestamp

            logger.debug("no message, idle for {:3.0f} sec".format(idle_time))

        # this is EC2_only
        if max_idle_time is not None and \
           idle_terminate_granularity is not None:
            if idle_time > max_idle_time:
                my_uptime = get_my_uptime()
                time_frac = (my_uptime % idle_terminate_granularity)

                logger.debug("Instance has been up for {:.0f} and inactive for {:.0f} time_frac={:.0f} terminate_thold={:.0f}".format(my_uptime,
                                                                                idle_time,
                                                                                time_frac,
                                                                                terminate_thold_sec))

                if time_frac > terminate_thold_sec:
                    logger.info("Instance has been up for {:.0f} and inactive for {:.0f}, terminating".format(my_uptime,
                                                                                idle_time))
                    for h in logger.handlers:
                        h.flush()
                    ec2_self_terminate(idle_time, my_uptime, message_count)

def process_message(m, local_message_i, max_run_time, run_dir,
                    aws_region,
                    server_name, log_stream_prefix):

    logger.info("processing message_id={}".format(m.message_id))

    event = json.loads(m.body)
    call_id = event['call_id']
    callset_id = event['callset_id']

    extra_env_debug = event.get('extra_env', {})
    unique_job_id = "{}:{}:{}".format(m.message_id, call_id, callset_id)
    logger.info("processing message_id={} callset_id={} call_id={}".format(m.message_id, callset_id, call_id))

    # FIXME this is all for debugging
    if 'DEBUG_THROW_EXCEPTION' in extra_env_debug:
        m.delete()
        raise Exception("Debug exception")
    message_id = m.message_id

    # run this in a thread: pywren.wrenhandler.generic_handler(event)
    p =  Process(target=job_handler, args=(event, local_message_i,
                                           run_dir, aws_region, server_name,
                                           log_stream_prefix))
    # is thread done
    p.start()
    pid = p.pid
    logger.info("processing message_id={} callset_id={} call_id={} process_pid={}".format(m.message_id, callset_id, call_id, pid))
    start_time = time.time()

    run_time = time.time() - start_time
    last_visibility_update_time = time.time()
    while run_time < max_run_time:
        time_since_visibility_update = time.time() - last_visibility_update_time
        est_visibility_left = SQS_VISIBILITY_SEC - time_since_visibility_update
        if est_visibility_left < (PROCESS_SLEEP_DUR_SEC*1.5):
            logger.debug("{} - {:3.1f}s since last visibility update, setting to {:3.1f} sec".format(message_id, time_since_visibility_update,
                                                                                SQS_VISIBILITY_SEC))
            last_visibility_update_time = time.time()
            response = m.change_visibility(VisibilityTimeout=SQS_VISIBILITY_SEC)

        if not p.is_alive():
            logger.debug("{} - attempting to join process".format(message_id))
            # FIXME will this join ever hang?
            p.join()
            break
        else:
            logger.debug("{} - {:3.1f}s since visibility update, sleeping".format(message_id, time_since_visibility_update))
            time.sleep(PROCESS_SLEEP_DUR_SEC)

        run_time = time.time() - start_time

    if p.exitcode is None:
        p.terminate()  # FIXME PRINT LOTS OF ERRORS HERE
        logger.warn("{} - Had to manually terminate process ".format(message_id))
    logger.info("deleting message_id={} callset_id={} call_id={}".format(m.message_id, callset_id, call_id))

    m.delete()

def copy_runtime(tgt_dir):
    files = glob(os.path.join(pywren.SOURCE_DIR, "./*.py"))
    for f in files:
        shutil.copy(f, os.path.join(tgt_dir, os.path.basename(f)))

def job_handler(job, job_i, run_dir, aws_region,
                server_name, log_stream_prefix,
                extra_context = None,
                delete_taskdir=True):
    """
    Run a deserialized job in run_dir

    Just for debugging
    """
    debug_pid = open("/tmp/pywren.scripts.standalone.{}.{}.log".format(os.getpid(),
                                                                       time.time()), 'w')
    print "subprocess job_handler job i=", job_i, "pid=", os.getpid()
    session = boto3.session.Session(region_name=aws_region)
    # we do this here instead of in the global context
    # because of how multiprocessing works
    handler = watchtower.CloudWatchLogHandler(send_interval=20,
                                              log_group="pywren.standalone",
                                              stream_name=log_stream_prefix + "-{logger_name}",
                                              boto3_session=session,
                                              max_batch_count=10)
    log_format_str ='{} %(asctime)s - %(name)s - %(levelname)s - %(message)s'.format(server_name)

    formatter = logging.Formatter(log_format_str, "%Y-%m-%d %H:%M:%S")
    handler.setFormatter(formatter)

    wren_log = pywren.wrenhandler.logger # logging.getLogger('pywren.wrenhandler')
    wren_log.setLevel(logging.DEBUG)
    wren_log.propagate = 0
    wren_log.addHandler(handler)

    original_dir = os.getcwd()
    debug_pid.write("added log handler\n")

    task_run_dir = os.path.join(run_dir, str(job_i))
    shutil.rmtree(task_run_dir, True) # delete old modules
    os.makedirs(task_run_dir)
    copy_runtime(task_run_dir)

    context = {'jobnum' : job_i}
    if extra_context is not None:
        context.update(extra_context)

    os.chdir(task_run_dir)
    try:
        debug_pid.write("invoking generic_handler\n")

        wrenhandler.generic_handler(job, context)
    finally:
        debug_pid.write("generic handler finally\n")

        if delete_taskdir:
            shutil.rmtree(task_run_dir)
        os.chdir(original_dir)
    handler.flush()
    debug_pid.write("done and returning\n")

@click.command()
@click.option('--max_run_time', default=3600,
              help='max run time for a job', type=int)
@click.option('--run_dir', default="/tmp/pywren.rundir",
              help='directory to hold intermediate output')
@click.option('--aws_region', default="us-west-2",
              help='aws region')
@click.option('--sqs_queue_name', default="pywren-queue",
              help='queue')
@click.option('--max_idle_time', default=None, type=int,
              help='maximum time for queue to remine idle before we try to self-terminate (sec)')
@click.option('--idle_terminate_granularity', default=None, type=int,
              help="only terminate if we have been up for an integral number of this")
@click.option('--queue_receive_message_timeout', default=10, type=int,
              help="longpoll timeout for getting sqs messages")
def server(aws_region, max_run_time, run_dir, sqs_queue_name, max_idle_time,
           idle_terminate_granularity, queue_receive_message_timeout):

    session = boto3.session.Session(region_name=aws_region)

    # make boto quiet locally FIXME is there a better way of doing this?
    logging.getLogger('boto').setLevel(logging.CRITICAL)
    logging.getLogger('boto3').setLevel(logging.CRITICAL)
    logging.getLogger('botocore').setLevel(logging.CRITICAL)

    if platform.node() != 'c65':

        instance = get_my_ec2_instance(aws_region)
        ec2_metadata = get_my_ec2_meta(instance)
        server_name = ec2_metadata['Name']
        log_stream_prefix = ec2_metadata['instance_id']
    else:
        server_name='c65'
        log_stream_prefix='millennium-c65'

    log_format_str ='{} %(asctime)s - %(name)s - %(levelname)s - %(message)s'.format(server_name)

    formatter = logging.Formatter(log_format_str, "%Y-%m-%d %H:%M:%S")

    handler = watchtower.CloudWatchLogHandler(send_interval=20,
                                              log_group="pywren.standalone",
                                              stream_name=log_stream_prefix + "-{logger_name}",
                                              boto3_session=session,
                                              max_batch_count=10)

    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)

    #config = pywren.wrenconfig.default()
    server_runner(aws_region, sqs_queue_name,
                  max_run_time, os.path.abspath(run_dir),
                  server_name, log_stream_prefix,
                  max_idle_time,
                  idle_terminate_granularity,
                  queue_receive_message_timeout)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
XLS-to-dict and csv-to-dict are essentially backends for xls2json.
"""
import datetime
import re
from collections import OrderedDict
from functools import reduce
from io import BytesIO

import unicodecsv as csv
import xlrd
from xlrd import XLRDError
from xlrd.xldate import XLDateAmbiguous

from pyxform import constants
from pyxform.errors import PyXFormError
from pyxform.utils import basestring, unichr, unicode

XL_DATE_AMBIGOUS_MSG = (
    "The xls file provided has an invalid date on the %s sheet, under"
    " the %s column on row number %s"
)

def _list_to_dict_list(list_items):
    """
    Takes a list and creates a dict with the list values as keys.
    Returns a list of the created dict or an empty list
    """
    if list_items:
        k = OrderedDict()
        for item in list_items:
            k["%s" % item] = ""
        return [k]
    return []

def xls_to_dict(path_or_file):
    """
    Return a Python dictionary with a key for each worksheet
    name. For each sheet there is a list of dictionaries, each
    dictionary corresponds to a single row in the worksheet. A
    dictionary has keys taken from the column headers and values
    equal to the cell value for that row and column.
    All the keys and leaf elements are unicode text.
    """
    try:
        if isinstance(path_or_file, basestring):
            workbook = xlrd.open_workbook(filename=path_or_file)
        else:
            workbook = xlrd.open_workbook(file_contents=path_or_file.read())
    except XLRDError as error:
        raise PyXFormError("Error reading .xls file: %s" % error)

    def xls_to_dict_normal_sheet(sheet):
        def iswhitespace(string):
            return isinstance(string, basestring) and len(string.strip()) == 0

        # Check for duplicate column headers
        column_header_list = list()
        for column in range(0, sheet.ncols):
            column_header = sheet.cell_value(0, column)
            if column_header in column_header_list:
                raise PyXFormError("Duplicate column header: %s" % column_header)
            # xls file with 3 columns mostly have a 3 more columns that are
            # blank by default or something, skip during check
            if column_header is not None:
                if not iswhitespace(column_header):
                    # strip whitespaces from the header
                    clean_header = re.sub(r"( )+", " ", column_header.strip())
                    column_header_list.append(clean_header)

        result = []
        for row in range(1, sheet.nrows):
            row_dict = OrderedDict()
            for column in range(0, sheet.ncols):
                # Changing to cell_value function
                # convert to string, in case it is not string
                key = "%s" % sheet.cell_value(0, column)
                key = key.strip()
                value = sheet.cell_value(row, column)
                # remove whitespace at the beginning and end of value
                if isinstance(value, basestring):
                    value = value.strip()
                value_type = sheet.cell_type(row, column)
                if value is not None:
                    if not iswhitespace(value):
                        try:
                            row_dict[key] = xls_value_to_unicode(
                                value, value_type, workbook.datemode
                            )
                        except XLDateAmbiguous:
                            raise PyXFormError(
                                XL_DATE_AMBIGOUS_MSG % (sheet.name, column_header, row)
                            )
                # Taking this condition out so I can get accurate row numbers.
                # TODO: Do the same for csvs
                # if row_dict != {}:
            result.append(row_dict)
        return result, _list_to_dict_list(column_header_list)

    def xls_value_from_sheet(sheet, row, column):
        value = sheet.cell_value(row, column)
        value_type = sheet.cell_type(row, column)
        if value is not None and value != "":
            try:
                return xls_value_to_unicode(value, value_type, workbook.datemode)
            except XLDateAmbiguous:
                raise PyXFormError(XL_DATE_AMBIGOUS_MSG % (sheet.name, column, row))
        else:
            raise PyXFormError("Empty Value")

    def _xls_to_dict_cascade_sheet(sheet):
        result = []
        rs_dict = OrderedDict()  # tmp dict to hold entire structure

        def slugify(s):
            return re.sub(r"\W+", "_", s.strip().lower())

        prefix = "$PREFIX$"
        # get col headers and position first, ignore first column
        for column in range(1, sheet.ncols):
            col_name = sheet.cell_value(0, column)
            rs_dict[col_name] = {
                "pos": column,
                "data": [],
                "itemset": col_name,
                "type": constants.SELECT_ONE,
                "name": prefix
                if (column == sheet.ncols - 1)
                else "".join([prefix, "_", col_name]),
                "label": sheet.cell_value(1, column),
            }
            if column > 1:
                rs_dict[col_name]["parent"] = sheet.cell_value(0, column - 1)
            else:
                rs_dict[col_name]["choices"] = []
            choice_filter = ""
            for a in range(1, column):
                prev_col_name = sheet.cell_value(0, a)
                if choice_filter != "":
                    choice_filter += " and %s=${%s_%s}" % (
                        prev_col_name,
                        prefix,
                        prev_col_name,
                    )
                else:
                    choice_filter += "%s=${%s_%s}" % (
                        prev_col_name,
                        prefix,
                        prev_col_name,
                    )
            rs_dict[col_name]["choice_filter"] = choice_filter
        # get data, use new cascade dict structure, data starts on 3 row
        for row in range(2, sheet.nrows):
            # go through each header aka column
            for col_name in rs_dict:
                column = rs_dict[col_name]["pos"]
                cell_data = xls_value_from_sheet(sheet, row, column)
                try:
                    rs_dict[col_name]["data"].index(slugify(cell_data))
                except ValueError:
                    rs_dict[col_name]["data"].append(slugify(cell_data))
                    if "choices" in rs_dict[col_name]:
                        rs_dict[col_name]["choices"].append(
                            {"name": slugify(cell_data), "label": cell_data}
                        )
                data = {
                    "name": slugify(cell_data),
                    "label": cell_data.strip(),
                    constants.LIST_NAME: col_name,
                }
                for prev_column in range(1, column):
                    prev_col_name = sheet.cell_value(0, prev_column)
                    data[prev_col_name] = slugify(
                        xls_value_from_sheet(sheet, row, prev_column)
                    )
                result.append(data)
        # order
        kl = []
        for column in range(1, sheet.ncols):
            col_name = sheet.cell_value(0, column)
            if "parent" in rs_dict[col_name]:
                rs_dict[col_name].pop("parent")
            if "pos" in rs_dict[col_name]:
                rs_dict[col_name].pop("pos")
            if "data" in rs_dict[col_name]:
                rs_dict[col_name].pop("data")
            kl.append(rs_dict[col_name])

            # create list with no duplicates
        choices = []
        for rec in result:
            c = 0
            for check in result:
                if rec == check:
                    c += 1
            if c == 1:
                choices.append(rec)
            else:
                try:
                    choices.index(rec)
                except ValueError:
                    choices.append(rec)
        return [{"choices": choices, "questions": kl}]

    result = OrderedDict()
    for sheet in workbook.sheets():
        # Do not process sheets that have nothing to do with XLSForm.
        if sheet.name not in constants.SUPPORTED_SHEET_NAMES:
            if len(workbook.sheets()) == 1:
                result[constants.SURVEY], result[
                    "%s_header" % sheet.name
                ] = xls_to_dict_normal_sheet(sheet)
            else:
                continue
        if sheet.name == constants.CASCADING_CHOICES:
            result[sheet.name] = _xls_to_dict_cascade_sheet(sheet)
        else:
            result[sheet.name], result[
                "%s_header" % sheet.name
            ] = xls_to_dict_normal_sheet(sheet)

    return result

def xls_value_to_unicode(value, value_type, datemode):
    """
    Take a xls formatted value and try to make a unicode string
    representation.
    """
    if value_type == xlrd.XL_CELL_BOOLEAN:
        return "TRUE" if value else "FALSE"
    elif value_type == xlrd.XL_CELL_NUMBER:
        # Try to display as an int if possible.
        int_value = int(value)
        if int_value == value:
            return unicode(int_value)
        else:
            return unicode(value)
    elif value_type is xlrd.XL_CELL_DATE:
        # Warn that it is better to single quote as a string.
        # error_location = cellFormatString % (ss_row_idx, ss_col_idx)
        # raise Exception(
        #   "Cannot handle excel formatted date at " + error_location)
        datetime_or_time_only = xlrd.xldate_as_tuple(value, datemode)
        if datetime_or_time_only[:3] == (0, 0, 0):
            # must be time only
            return unicode(datetime.time(*datetime_or_time_only[3:]))
        return unicode(datetime.datetime(*datetime_or_time_only))
    else:
        # ensure unicode and replace nbsp spaces with normal ones
        # to avoid this issue:
        # https://github.com/modilabs/pyxform/issues/83
        return unicode(value).replace(unichr(160), " ")

def get_cascading_json(sheet_list, prefix, level):
    return_list = []
    for row in sheet_list:
        if "stopper" in row:
            if row["stopper"] == level:
                # last element's name IS the prefix; doesn't need level
                return_list[-1]["name"] = prefix
                return return_list
            else:
                continue
        elif "lambda" in row:

            def replace_prefix(d, prefix):
                for k, v in d.items():
                    if isinstance(v, basestring):
                        d[k] = v.replace("$PREFIX$", prefix)
                    elif isinstance(v, dict):
                        d[k] = replace_prefix(v, prefix)
                    elif isinstance(v, list):
                        d[k] = map(lambda x: replace_prefix(x, prefix), v)
                return d

            return_list.append(replace_prefix(row["lambda"], prefix))
    raise PyXFormError(
        "Found a cascading_select "
        + level
        + ", but could not find "
        + level
        + "in cascades sheet."
    )

def csv_to_dict(path_or_file):
    if isinstance(path_or_file, basestring):
        csv_data = open(path_or_file, "rb")
    else:
        csv_data = path_or_file

    _dict = OrderedDict()

    def first_column_as_sheet_name(row):
        if len(row) == 0:
            return None, None
        elif len(row) == 1:
            return row[0], None
        else:
            s_or_c = row[0]
            content = row[1:]
            if s_or_c == "":
                s_or_c = None
            # concatenate all the strings in content
            if reduce(lambda x, y: x + y, content) == "":
                # content is a list of empty strings
                content = None
            return s_or_c, content

    reader = csv.reader(csv_data, encoding="utf-8")
    sheet_name = None
    current_headers = None
    for row in reader:
        survey_or_choices, content = first_column_as_sheet_name(row)
        if survey_or_choices is not None:
            sheet_name = survey_or_choices
            if sheet_name not in _dict:
                _dict[unicode(sheet_name)] = []
            current_headers = None
        if content is not None:
            if current_headers is None:
                current_headers = content
                _dict["%s_header" % sheet_name] = _list_to_dict_list(current_headers)
            else:
                _d = OrderedDict()
                for key, val in zip(current_headers, content):
                    if val != "":
                        # Slight modification so values are striped
                        # this is because csvs often spaces following commas
                        # (but the csv reader might already handle that.)
                        _d[unicode(key)] = unicode(val.strip())
                _dict[sheet_name].append(_d)
    csv_data.close()
    return _dict

"""
I want the ability to go:

    xls => csv
so that we can go:
    xls => csv => survey

and some day:
    csv => xls

"""

def convert_file_to_csv_string(path):
    """
    This will open a csv or xls file and return a CSV in the format:
        sheet_name1
        ,col1,col2
        ,r1c1,r1c2
        ,r2c1,r2c2
        sheet_name2
        ,col1,col2
        ,r1c1,r1c2
        ,r2c1,r2c2

    Currently, it processes csv files and xls files to ensure consistent
    csv delimiters, etc. for tests.
    """
    if path.endswith(".csv"):
        imported_sheets = csv_to_dict(path)
    else:
        imported_sheets = xls_to_dict(path)
    foo = BytesIO()
    writer = csv.writer(foo, delimiter=",", quotechar='"', quoting=csv.QUOTE_MINIMAL)
    for sheet_name, rows in imported_sheets.items():
        writer.writerow([sheet_name])
        out_keys = []
        out_rows = []
        for row in rows:
            out_row = []
            for key in row.keys():
                if key not in out_keys:
                    out_keys.append(key)
            for out_key in out_keys:
                out_row.append(row.get(out_key, None))
            out_rows.append(out_row)
        writer.writerow([None] + out_keys)
        for out_row in out_rows:
            writer.writerow([None] + out_row)
    return foo.getvalue().decode("utf-8")
<commit_msg>Using the correct name for the header section.<commit_after># -*- coding: utf-8 -*-
"""
XLS-to-dict and csv-to-dict are essentially backends for xls2json.
"""
import datetime
import re
from collections import OrderedDict
from functools import reduce
from io import BytesIO

import unicodecsv as csv
import xlrd
from xlrd import XLRDError
from xlrd.xldate import XLDateAmbiguous

from pyxform import constants
from pyxform.errors import PyXFormError
from pyxform.utils import basestring, unichr, unicode

XL_DATE_AMBIGOUS_MSG = (
    "The xls file provided has an invalid date on the %s sheet, under"
    " the %s column on row number %s"
)

def _list_to_dict_list(list_items):
    """
    Takes a list and creates a dict with the list values as keys.
    Returns a list of the created dict or an empty list
    """
    if list_items:
        k = OrderedDict()
        for item in list_items:
            k["%s" % item] = ""
        return [k]
    return []

def xls_to_dict(path_or_file):
    """
    Return a Python dictionary with a key for each worksheet
    name. For each sheet there is a list of dictionaries, each
    dictionary corresponds to a single row in the worksheet. A
    dictionary has keys taken from the column headers and values
    equal to the cell value for that row and column.
    All the keys and leaf elements are unicode text.
    """
    try:
        if isinstance(path_or_file, basestring):
            workbook = xlrd.open_workbook(filename=path_or_file)
        else:
            workbook = xlrd.open_workbook(file_contents=path_or_file.read())
    except XLRDError as error:
        raise PyXFormError("Error reading .xls file: %s" % error)

    def xls_to_dict_normal_sheet(sheet):
        def iswhitespace(string):
            return isinstance(string, basestring) and len(string.strip()) == 0

        # Check for duplicate column headers
        column_header_list = list()
        for column in range(0, sheet.ncols):
            column_header = sheet.cell_value(0, column)
            if column_header in column_header_list:
                raise PyXFormError("Duplicate column header: %s" % column_header)
            # xls file with 3 columns mostly have a 3 more columns that are
            # blank by default or something, skip during check
            if column_header is not None:
                if not iswhitespace(column_header):
                    # strip whitespaces from the header
                    clean_header = re.sub(r"( )+", " ", column_header.strip())
                    column_header_list.append(clean_header)

        result = []
        for row in range(1, sheet.nrows):
            row_dict = OrderedDict()
            for column in range(0, sheet.ncols):
                # Changing to cell_value function
                # convert to string, in case it is not string
                key = "%s" % sheet.cell_value(0, column)
                key = key.strip()
                value = sheet.cell_value(row, column)
                # remove whitespace at the beginning and end of value
                if isinstance(value, basestring):
                    value = value.strip()
                value_type = sheet.cell_type(row, column)
                if value is not None:
                    if not iswhitespace(value):
                        try:
                            row_dict[key] = xls_value_to_unicode(
                                value, value_type, workbook.datemode
                            )
                        except XLDateAmbiguous:
                            raise PyXFormError(
                                XL_DATE_AMBIGOUS_MSG % (sheet.name, column_header, row)
                            )
                # Taking this condition out so I can get accurate row numbers.
                # TODO: Do the same for csvs
                # if row_dict != {}:
            result.append(row_dict)
        return result, _list_to_dict_list(column_header_list)

    def xls_value_from_sheet(sheet, row, column):
        value = sheet.cell_value(row, column)
        value_type = sheet.cell_type(row, column)
        if value is not None and value != "":
            try:
                return xls_value_to_unicode(value, value_type, workbook.datemode)
            except XLDateAmbiguous:
                raise PyXFormError(XL_DATE_AMBIGOUS_MSG % (sheet.name, column, row))
        else:
            raise PyXFormError("Empty Value")

    def _xls_to_dict_cascade_sheet(sheet):
        result = []
        rs_dict = OrderedDict()  # tmp dict to hold entire structure

        def slugify(s):
            return re.sub(r"\W+", "_", s.strip().lower())

        prefix = "$PREFIX$"
        # get col headers and position first, ignore first column
        for column in range(1, sheet.ncols):
            col_name = sheet.cell_value(0, column)
            rs_dict[col_name] = {
                "pos": column,
                "data": [],
                "itemset": col_name,
                "type": constants.SELECT_ONE,
                "name": prefix
                if (column == sheet.ncols - 1)
                else "".join([prefix, "_", col_name]),
                "label": sheet.cell_value(1, column),
            }
            if column > 1:
                rs_dict[col_name]["parent"] = sheet.cell_value(0, column - 1)
            else:
                rs_dict[col_name]["choices"] = []
            choice_filter = ""
            for a in range(1, column):
                prev_col_name = sheet.cell_value(0, a)
                if choice_filter != "":
                    choice_filter += " and %s=${%s_%s}" % (
                        prev_col_name,
                        prefix,
                        prev_col_name,
                    )
                else:
                    choice_filter += "%s=${%s_%s}" % (
                        prev_col_name,
                        prefix,
                        prev_col_name,
                    )
            rs_dict[col_name]["choice_filter"] = choice_filter
        # get data, use new cascade dict structure, data starts on 3 row
        for row in range(2, sheet.nrows):
            # go through each header aka column
            for col_name in rs_dict:
                column = rs_dict[col_name]["pos"]
                cell_data = xls_value_from_sheet(sheet, row, column)
                try:
                    rs_dict[col_name]["data"].index(slugify(cell_data))
                except ValueError:
                    rs_dict[col_name]["data"].append(slugify(cell_data))
                    if "choices" in rs_dict[col_name]:
                        rs_dict[col_name]["choices"].append(
                            {"name": slugify(cell_data), "label": cell_data}
                        )
                data = {
                    "name": slugify(cell_data),
                    "label": cell_data.strip(),
                    constants.LIST_NAME: col_name,
                }
                for prev_column in range(1, column):
                    prev_col_name = sheet.cell_value(0, prev_column)
                    data[prev_col_name] = slugify(
                        xls_value_from_sheet(sheet, row, prev_column)
                    )
                result.append(data)
        # order
        kl = []
        for column in range(1, sheet.ncols):
            col_name = sheet.cell_value(0, column)
            if "parent" in rs_dict[col_name]:
                rs_dict[col_name].pop("parent")
            if "pos" in rs_dict[col_name]:
                rs_dict[col_name].pop("pos")
            if "data" in rs_dict[col_name]:
                rs_dict[col_name].pop("data")
            kl.append(rs_dict[col_name])

            # create list with no duplicates
        choices = []
        for rec in result:
            c = 0
            for check in result:
                if rec == check:
                    c += 1
            if c == 1:
                choices.append(rec)
            else:
                try:
                    choices.index(rec)
                except ValueError:
                    choices.append(rec)
        return [{"choices": choices, "questions": kl}]

    result = OrderedDict()
    for sheet in workbook.sheets():
        # Do not process sheets that have nothing to do with XLSForm.
        if sheet.name not in constants.SUPPORTED_SHEET_NAMES:
            if len(workbook.sheets()) == 1:
                result[constants.SURVEY], result[
                    "%s_header" % constants.SURVEY
                ] = xls_to_dict_normal_sheet(sheet)
            else:
                continue
        if sheet.name == constants.CASCADING_CHOICES:
            result[sheet.name] = _xls_to_dict_cascade_sheet(sheet)
        else:
            result[sheet.name], result[
                "%s_header" % sheet.name
            ] = xls_to_dict_normal_sheet(sheet)

    return result

def xls_value_to_unicode(value, value_type, datemode):
    """
    Take a xls formatted value and try to make a unicode string
    representation.
    """
    if value_type == xlrd.XL_CELL_BOOLEAN:
        return "TRUE" if value else "FALSE"
    elif value_type == xlrd.XL_CELL_NUMBER:
        # Try to display as an int if possible.
        int_value = int(value)
        if int_value == value:
            return unicode(int_value)
        else:
            return unicode(value)
    elif value_type is xlrd.XL_CELL_DATE:
        # Warn that it is better to single quote as a string.
        # error_location = cellFormatString % (ss_row_idx, ss_col_idx)
        # raise Exception(
        #   "Cannot handle excel formatted date at " + error_location)
        datetime_or_time_only = xlrd.xldate_as_tuple(value, datemode)
        if datetime_or_time_only[:3] == (0, 0, 0):
            # must be time only
            return unicode(datetime.time(*datetime_or_time_only[3:]))
        return unicode(datetime.datetime(*datetime_or_time_only))
    else:
        # ensure unicode and replace nbsp spaces with normal ones
        # to avoid this issue:
        # https://github.com/modilabs/pyxform/issues/83
        return unicode(value).replace(unichr(160), " ")

def get_cascading_json(sheet_list, prefix, level):
    return_list = []
    for row in sheet_list:
        if "stopper" in row:
            if row["stopper"] == level:
                # last element's name IS the prefix; doesn't need level
                return_list[-1]["name"] = prefix
                return return_list
            else:
                continue
        elif "lambda" in row:

            def replace_prefix(d, prefix):
                for k, v in d.items():
                    if isinstance(v, basestring):
                        d[k] = v.replace("$PREFIX$", prefix)
                    elif isinstance(v, dict):
                        d[k] = replace_prefix(v, prefix)
                    elif isinstance(v, list):
                        d[k] = map(lambda x: replace_prefix(x, prefix), v)
                return d

            return_list.append(replace_prefix(row["lambda"], prefix))
    raise PyXFormError(
        "Found a cascading_select "
        + level
        + ", but could not find "
        + level
        + "in cascades sheet."
    )

def csv_to_dict(path_or_file):
    if isinstance(path_or_file, basestring):
        csv_data = open(path_or_file, "rb")
    else:
        csv_data = path_or_file

    _dict = OrderedDict()

    def first_column_as_sheet_name(row):
        if len(row) == 0:
            return None, None
        elif len(row) == 1:
            return row[0], None
        else:
            s_or_c = row[0]
            content = row[1:]
            if s_or_c == "":
                s_or_c = None
            # concatenate all the strings in content
            if reduce(lambda x, y: x + y, content) == "":
                # content is a list of empty strings
                content = None
            return s_or_c, content

    reader = csv.reader(csv_data, encoding="utf-8")
    sheet_name = None
    current_headers = None
    for row in reader:
        survey_or_choices, content = first_column_as_sheet_name(row)
        if survey_or_choices is not None:
            sheet_name = survey_or_choices
            if sheet_name not in _dict:
                _dict[unicode(sheet_name)] = []
            current_headers = None
        if content is not None:
            if current_headers is None:
                current_headers = content
                _dict["%s_header" % sheet_name] = _list_to_dict_list(current_headers)
            else:
                _d = OrderedDict()
                for key, val in zip(current_headers, content):
                    if val != "":
                        # Slight modification so values are striped
                        # this is because csvs often spaces following commas
                        # (but the csv reader might already handle that.)
                        _d[unicode(key)] = unicode(val.strip())
                _dict[sheet_name].append(_d)
    csv_data.close()
    return _dict

"""
I want the ability to go:

    xls => csv
so that we can go:
    xls => csv => survey

and some day:
    csv => xls

"""

def convert_file_to_csv_string(path):
    """
    This will open a csv or xls file and return a CSV in the format:
        sheet_name1
        ,col1,col2
        ,r1c1,r1c2
        ,r2c1,r2c2
        sheet_name2
        ,col1,col2
        ,r1c1,r1c2
        ,r2c1,r2c2

    Currently, it processes csv files and xls files to ensure consistent
    csv delimiters, etc. for tests.
    """
    if path.endswith(".csv"):
        imported_sheets = csv_to_dict(path)
    else:
        imported_sheets = xls_to_dict(path)
    foo = BytesIO()
    writer = csv.writer(foo, delimiter=",", quotechar='"', quoting=csv.QUOTE_MINIMAL)
    for sheet_name, rows in imported_sheets.items():
        writer.writerow([sheet_name])
        out_keys = []
        out_rows = []
        for row in rows:
            out_row = []
            for key in row.keys():
                if key not in out_keys:
                    out_keys.append(key)
            for out_key in out_keys:
                out_row.append(row.get(out_key, None))
            out_rows.append(out_row)
        writer.writerow([None] + out_keys)
        for out_row in out_rows:
            writer.writerow([None] + out_row)
    return foo.getvalue().decode("utf-8")
<|endoftext|>


<commit_before>from tornado.web import authenticated, HTTPError

from os.path import isdir, join, exists
from os import makedirs, listdir

from shutil import copyfileobj, rmtree

from .util import check_access
from .base_handlers import BaseHandler

from qiita_core.qiita_settings import qiita_config

from qiita_db.util import (get_files_from_uploads_folders,
                           get_mountpoint)
from qiita_db.study import Study
from qiita_db.user import User
from qiita_db.exceptions import QiitaDBUnknownIDError

class StudyUploadFileHandler(BaseHandler):
    @authenticated
    def display_template(self, study_id, msg):
        """Simple function to avoid duplication of code"""
        study_id = int(study_id)
        study = Study(study_id)
        check_access(User(self.current_user), study, no_public=True,
                     raise_error=True)

        # getting the ontologies
        self.render('upload.html', user=self.current_user,
                    study_title=study.title, study_info=study.info,
                    study_id=study_id,
                    extensions=','.join(qiita_config.valid_upload_extension),
                    max_upload_size=qiita_config.max_upload_size,
                    files=get_files_from_uploads_folders(str(study_id)))

    @authenticated
    def get(self, study_id):
        try:
            study = Study(int(study_id))
        except QiitaDBUnknownIDError:
            raise HTTPError(404, "Study %s does not exist" % study_id)
        check_access(User(self.current_user), study, no_public=True,
                     raise_error=True)
        self.display_template(study_id, "")

class UploadFileHandler(BaseHandler):
    # """ main upload class
    # based on
    # https://github.com/23/resumable.js/blob/master/samples/Backend%20on%20PHP.md
    # """
    def validate_file_extension(self, filename):
        """simple method to avoid duplication of code

        This validation is server side in case they can go around the client
        side validation
        """
        if not filename.endswith(tuple(qiita_config.valid_upload_extension)):
            self.set_status(415)
            raise HTTPError(415, "User %s is trying to upload %d" %
                                 (self.current_user, filename))

    @authenticated
    def post(self):
        resumable_identifier = self.get_argument('resumableIdentifier')
        resumable_filename = self.get_argument('resumableFilename')
        resumable_chunk_number = int(self.get_argument('resumableChunkNumber'))
        resumable_total_chunks = int(self.get_argument('resumableTotalChunks'))
        study_id = self.get_argument('study_id')
        data = self.request.files['file'][0]['body']

        check_access(User(self.current_user), Study(int(study_id)),
                     no_public=True, raise_error=True)

        self.validate_file_extension(resumable_filename)

        _, base_fp = get_mountpoint("uploads")[0]
        fp = join(base_fp, study_id, resumable_identifier)
        # creating temporal folder for upload
        if not isdir(fp):
            makedirs(fp)
        dfp = join(fp, '%s.part.%d' % (resumable_filename,
                                       resumable_chunk_number))

        # writting the output file
        with open(dfp, 'wb') as f:
            f.write(bytes(data))

        # validating if all files have been uploaded
        num_files = len([n for n in listdir(fp)])
        if resumable_total_chunks == num_files:
            # creating final destination
            ffp = join(base_fp, study_id, resumable_filename)
            with open(ffp, 'wb') as f:
                for c in range(1, resumable_total_chunks+1):
                    chunk = join(fp, '%s.part.%d' % (resumable_filename, c))
                    copyfileobj(open(chunk, 'rb'), f)

                # deleting the tmp folder with contents and finish file
                rmtree(fp)
                self.set_status(200)

    @authenticated
    def get(self):
        """ this is the first point of entry into the upload service

        this should either set the status as 400 (error) so the file/chunk is
        sent via post or 200 (valid) to not send the file
        """
        study_id = self.get_argument('study_id')
        resumable_filename = self.get_argument('resumableFilename')
        resumable_chunk_number = self.get_argument('resumableChunkNumber')

        check_access(User(self.current_user), Study(int(study_id)),
                     no_public=True, raise_error=True)

        self.validate_file_extension(resumable_filename)

        # temporaly filename or chunck
        _, fp = get_mountpoint("uploads")[0]
        tfp = join(fp, study_id,
                   resumable_filename + '.part.' + resumable_chunk_number)

        if exists(tfp):
            self.set_status(200)
        else:
            self.set_status(400)
<commit_msg>fix another %d error<commit_after>from tornado.web import authenticated, HTTPError

from os.path import isdir, join, exists
from os import makedirs, listdir

from shutil import copyfileobj, rmtree

from .util import check_access
from .base_handlers import BaseHandler

from qiita_core.qiita_settings import qiita_config

from qiita_db.util import (get_files_from_uploads_folders,
                           get_mountpoint)
from qiita_db.study import Study
from qiita_db.user import User
from qiita_db.exceptions import QiitaDBUnknownIDError

class StudyUploadFileHandler(BaseHandler):
    @authenticated
    def display_template(self, study_id, msg):
        """Simple function to avoid duplication of code"""
        study_id = int(study_id)
        study = Study(study_id)
        check_access(User(self.current_user), study, no_public=True,
                     raise_error=True)

        # getting the ontologies
        self.render('upload.html', user=self.current_user,
                    study_title=study.title, study_info=study.info,
                    study_id=study_id,
                    extensions=','.join(qiita_config.valid_upload_extension),
                    max_upload_size=qiita_config.max_upload_size,
                    files=get_files_from_uploads_folders(str(study_id)))

    @authenticated
    def get(self, study_id):
        try:
            study = Study(int(study_id))
        except QiitaDBUnknownIDError:
            raise HTTPError(404, "Study %s does not exist" % study_id)
        check_access(User(self.current_user), study, no_public=True,
                     raise_error=True)
        self.display_template(study_id, "")

class UploadFileHandler(BaseHandler):
    # """ main upload class
    # based on
    # https://github.com/23/resumable.js/blob/master/samples/Backend%20on%20PHP.md
    # """
    def validate_file_extension(self, filename):
        """simple method to avoid duplication of code

        This validation is server side in case they can go around the client
        side validation
        """
        if not filename.endswith(tuple(qiita_config.valid_upload_extension)):
            self.set_status(415)
            raise HTTPError(415, "User %s is trying to upload %s" %
                                 (self.current_user, str(filename)))

    @authenticated
    def post(self):
        resumable_identifier = self.get_argument('resumableIdentifier')
        resumable_filename = self.get_argument('resumableFilename')
        resumable_chunk_number = int(self.get_argument('resumableChunkNumber'))
        resumable_total_chunks = int(self.get_argument('resumableTotalChunks'))
        study_id = self.get_argument('study_id')
        data = self.request.files['file'][0]['body']

        check_access(User(self.current_user), Study(int(study_id)),
                     no_public=True, raise_error=True)

        self.validate_file_extension(resumable_filename)

        _, base_fp = get_mountpoint("uploads")[0]
        fp = join(base_fp, study_id, resumable_identifier)
        # creating temporal folder for upload
        if not isdir(fp):
            makedirs(fp)
        dfp = join(fp, '%s.part.%d' % (resumable_filename,
                                       resumable_chunk_number))

        # writting the output file
        with open(dfp, 'wb') as f:
            f.write(bytes(data))

        # validating if all files have been uploaded
        num_files = len([n for n in listdir(fp)])
        if resumable_total_chunks == num_files:
            # creating final destination
            ffp = join(base_fp, study_id, resumable_filename)
            with open(ffp, 'wb') as f:
                for c in range(1, resumable_total_chunks+1):
                    chunk = join(fp, '%s.part.%d' % (resumable_filename, c))
                    copyfileobj(open(chunk, 'rb'), f)

                # deleting the tmp folder with contents and finish file
                rmtree(fp)
                self.set_status(200)

    @authenticated
    def get(self):
        """ this is the first point of entry into the upload service

        this should either set the status as 400 (error) so the file/chunk is
        sent via post or 200 (valid) to not send the file
        """
        study_id = self.get_argument('study_id')
        resumable_filename = self.get_argument('resumableFilename')
        resumable_chunk_number = self.get_argument('resumableChunkNumber')

        check_access(User(self.current_user), Study(int(study_id)),
                     no_public=True, raise_error=True)

        self.validate_file_extension(resumable_filename)

        # temporaly filename or chunck
        _, fp = get_mountpoint("uploads")[0]
        tfp = join(fp, study_id,
                   resumable_filename + '.part.' + resumable_chunk_number)

        if exists(tfp):
            self.set_status(200)
        else:
            self.set_status(400)
<|endoftext|>


<commit_before>import time
import traceback
import ConfigParser

from lib.pipeline import *
from lib.utils import *
from lib.cache import *
from lib.event import *

SYSTEM_CONF_FILE = "conf/system.conf"
PIPELINE_CONF_FILE = "conf/pipeline.conf"
BOTS_CONF_FILE = "conf/bots.conf"
LOGS_PATH = "logs/"

class Bot(object):

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

        self.logger = self.load_logger()
        self.logger.info('Bot is starting')

        self.load_configurations()

        src_queue, dest_queues = self.load_queues()
        self.pipeline = self.load_pipeline(src_queue, dest_queues)

        self.init()

    def init(self):
        pass

    def start(self):
         self.logger.info('Bot start processing')

         count = 0
         while True:
             try:
                self.process()
                time.sleep(int(self.parameters.processing_interval))
                count = 0
             except:
                # self.close() # add this method
                self.logger.info('Bot is restarting')
                self.connect()
                self.logger.error(traceback.format_exc())
                count += 1
                if count > 5:
                    self.stop()
                time.sleep(5)

    def stop(self):
        self.logger.error("Bot found an error. Exiting")
        exit(-1)

    def load_configurations(self):
        self.parameters = Parameters()
        config = ConfigParser.ConfigParser()
        config.read(BOTS_CONF_FILE)

        default_section = "default"
        self.logger.debug("Loading configuration in default section from '%s' file" % BOTS_CONF_FILE)

        if config.has_section(default_section):
            for option in config.options(default_section):
                setattr(self.parameters, option, config.get(default_section, option))
                self.logger.debug("Parameter '%s' loaded with the value '%s'" % (option, config.get(default_section, option)))

        self.logger.debug("Loading configuration in %s section from '%s' file" % (self.bot_id, BOTS_CONF_FILE))

        if config.has_section(self.bot_id):
            for option in config.options(self.bot_id):
                setattr(self.parameters, option, config.get(self.bot_id, option))
                self.logger.debug("Parameter '%s' loaded with the value '%s'" % (option, config.get(self.bot_id, option)))

    def load_logger(self):
        config = ConfigParser.ConfigParser()
        config.read(SYSTEM_CONF_FILE)
        loglevel = config.get('Logging','level')
        return log(LOGS_PATH, self.bot_id, loglevel)

    def load_pipeline(self, src_queue, dest_queues):
        self.logger.debug("Connecting to pipeline queues")
        return Pipeline(src_queue, dest_queues)

    def load_queues(self):
        config = ConfigParser.ConfigParser()
        config.read(PIPELINE_CONF_FILE)
        self.logger.debug("Loading pipeline queues from '%s' file" % PIPELINE_CONF_FILE)

        for option in config.options("Pipeline"):
            if option == self.bot_id:
                queues = config.get("Pipeline", self.bot_id)
                queues = queues.split('|')

                if len(queues) == 2:
                    src_queue = queues[0].strip()
                    dest_queues = queues[1].strip()

                    if src_queue == "None":
                        src_queue = None
                    if dest_queues == "None":
                        dest_queue = None

                    dest_queues_list = list()
                    for queue in dest_queues.split(','):
                        dest_queues_list.append(queue.strip())

                    self.logger.info("Source queue '%s'" % src_queue)
                    self.logger.info("Destination queue(s) '%s'" % dest_queues_list)

                    return [src_queue, dest_queues_list]

        self.logger.error("Failed to load queues")
        self.stop()

    def send_message(self, message):
        try:
            message = unicode(message)
        except:
            message = force_decode(message)
        self.pipeline.send(message)

    def receive_message(self):
        raw_message = self.pipeline.receive()
        try:
            message = Event.from_unicode(raw_message)
        except:
            message = raw_message
        return message

    def acknowledge_message(self):
        self.pipeline.acknowledge()

class Parameters(object):
    pass

<commit_msg>update bot<commit_after>import time
import traceback
import ConfigParser

from lib.pipeline import *
from lib.utils import *
from lib.cache import *
from lib.event import *

SYSTEM_CONF_FILE = "conf/system.conf"
PIPELINE_CONF_FILE = "conf/pipeline.conf"
BOTS_CONF_FILE = "conf/bots.conf"
LOGS_PATH = "logs/"

class Bot(object):

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

        self.logger = self.load_logger()
        self.logger.info('Bot is starting')

        self.load_configurations()

        src_queue, dest_queues = self.load_queues()
        self.pipeline = self.load_pipeline(src_queue, dest_queues)

        self.init()

    def init(self):
        pass

    def start(self):
         self.logger.info('Bot start processing')

         while True:
             try:
                self.process()
                time.sleep(int(self.parameters.processing_interval))
             except:
                # self.close() # add this method
                self.logger.info('Bot is restarting')
                self.connect()
                self.logger.error(traceback.format_exc())
                time.sleep(5)

    def stop(self):
        self.logger.error("Bot found an error. Exiting")
        exit(-1)

    def load_configurations(self):
        self.parameters = Parameters()
        config = ConfigParser.ConfigParser()
        config.read(BOTS_CONF_FILE)

        default_section = "default"
        self.logger.debug("Loading configuration in default section from '%s' file" % BOTS_CONF_FILE)

        if config.has_section(default_section):
            for option in config.options(default_section):
                setattr(self.parameters, option, config.get(default_section, option))
                self.logger.debug("Parameter '%s' loaded with the value '%s'" % (option, config.get(default_section, option)))

        self.logger.debug("Loading configuration in %s section from '%s' file" % (self.bot_id, BOTS_CONF_FILE))

        if config.has_section(self.bot_id):
            for option in config.options(self.bot_id):
                setattr(self.parameters, option, config.get(self.bot_id, option))
                self.logger.debug("Parameter '%s' loaded with the value '%s'" % (option, config.get(self.bot_id, option)))

    def load_logger(self):
        config = ConfigParser.ConfigParser()
        config.read(SYSTEM_CONF_FILE)
        loglevel = config.get('Logging','level')
        return log(LOGS_PATH, self.bot_id, loglevel)

    def load_pipeline(self, src_queue, dest_queues):
        self.logger.debug("Connecting to pipeline queues")
        return Pipeline(src_queue, dest_queues)

    def load_queues(self):
        config = ConfigParser.ConfigParser()
        config.read(PIPELINE_CONF_FILE)
        self.logger.debug("Loading pipeline queues from '%s' file" % PIPELINE_CONF_FILE)

        for option in config.options("Pipeline"):
            if option == self.bot_id:
                queues = config.get("Pipeline", self.bot_id)

                src_queue, dest_queues = self.parse_queues(queues)
                if src_queue or dest_queues:
                    self.logger.info("Source queue '%s'" % src_queue)
                    self.logger.info("Destination queue(s) '%s'" % dest_queues)
                    return [src_queue, dest_queues]

        self.logger.error("Failed to load queues")
        self.stop()

    def parse_queues(self, queues):
        queues = queues.split('|')

        if len(queues) == 2:
            src_queue = queues[0].strip()
            if len(src_queue) == 0:
                src_queue = None

            dest_queues = queues[1].strip()
            if len(dest_queues) == 0:
                dest_queues_list = None
            else:
                dest_queues_list = list()
                for queue in dest_queues.split(','):
                    dest_queues_list.append(queue.strip())

            return [src_queue, dest_queues_list]
        else:
            return [None, None]

    def send_message(self, message):
        try:
            message = unicode(message)
        except:
            message = force_decode(message)
        self.pipeline.send(message)

    def receive_message(self):
        raw_message = self.pipeline.receive()
        try:
            message = Event.from_unicode(raw_message)
        except:
            message = raw_message
        return message

    def acknowledge_message(self):
        self.pipeline.acknowledge()

class Parameters(object):
    pass

<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

import Tkinter as tk
import os
from collections import deque

import pygame

from colors import BLACK, ORANGE, DARK_GRAY

FPS = 30

CELL_SIZE = 10
ROWS = 64
COLS = 128
DISPLAY_WIDTH = COLS * CELL_SIZE
DISPLAY_HEIGHT = ROWS * CELL_SIZE

LEFT_MOUSE_BUTTON = 1
RIGHT_MOUSE_BUTTON = 3

def make_line_of_tiles(point1, point2):
    tile_lst = []
    if point1[0] == point2[0]:
        startpoint = min(point1[1], point2[1])
        endpoint = max(point1[1], point2[1])
        for ypos in range(startpoint, endpoint+1, CELL_SIZE):
            tile_lst.append((point1[0], ypos))
    else:
        startpoint = min(point1[0], point2[0])
        endpoint = max(point1[0], point2[0])
        for xpos in range(startpoint, endpoint+1, CELL_SIZE):
            tile_lst.append((xpos, point1[1]))

    return tile_lst

class ChangeTilesCommand(object):
    def __init__(self, tiles, tilemap, remove=False):
        self.tile_lst = tiles
        self.tilemap = tilemap
        self.tiles_changed = []
        self.remove = remove

    def do(self):
        for tile in self.tile_lst:
            if self.remove:
                if tile in self.tilemap:
                    self.tilemap.remove(tile)
                    self.tiles_changed.append(tile)
            else:
                if tile not in self.tilemap:
                    self.tilemap.append(tile)
                    self.tiles_changed.append(tile)

    def undo(self):
        for tile in self.tiles_changed:
            if self.remove:
                self.tilemap.append(tile)
            else:
                self.tilemap.remove(tile)

class MapEditor(object):
    def __init__(self):
        self.root = tk.Tk()
        self.root.bind('<Motion>', self.motion)
        self.root.bind('<ButtonPress>', self.on_button_press)
        self.root.bind('<ButtonRelease>', self.on_button_release)
        self.root.bind('<Control-g>', self.toggle_grid)
        self.root.bind('<KeyPress>', self.on_key_press)
        self.root.bind('<Shift-KeyRelease>', self.shift_release)
        self.root.bind('<Control-z>', self.on_undo)
        self.root.bind('<Control-y>', self.on_redo)
        self.root.rowconfigure(0, weight=1)
        self.root.columnconfigure(0, weight=1)

        top = self.root.winfo_toplevel()
        self.menu_bar = tk.Menu(top)
        top['menu'] = self.menu_bar
        self.file_menu = tk.Menu(self.menu_bar)
        self.menu_bar.add_cascade(label='File', menu=self.file_menu)
        self.file_menu.add_command(label='new')
        self.file_menu.add_command(label='open')
        self.file_menu.add_command(label='save')
        self.file_menu.add_command(label='save as')
        self.file_menu.add_command(label='exit',
            command=self.exit_cmd)

        self.edit_menu = tk.Menu(self.menu_bar)
        self.menu_bar.add_cascade(label='Edit', menu=self.edit_menu)
        self.edit_menu.add_command(label='Undo', command=self.undo)
        self.edit_menu.add_command(label='Redo', command=self.redo)

        self.embed = tk.Frame(self.root, width=1400, height=700)
        self.embed.grid(row=0, column=0)

        os.environ['SDL_WINDOWID'] = str(self.embed.winfo_id())

        self.root.update()
        self.fps_clock = pygame.time.Clock()
        pygame.init()
        self.screen = pygame.display.set_mode((DISPLAY_WIDTH,
                DISPLAY_HEIGHT))

        self.quit = False

        self.selected = None
        self.mouse_x = 0
        self.mouse_y = 0

        self.wall_tex = pygame.image.load('../gfx/wall.png').convert()
        self.tiles = []

        self.show_grid = True

        self.shift_pressed = False
        self.point1 = None

        self.tile_line = None

        self.mouse_state = {LEFT_MOUSE_BUTTON:0, RIGHT_MOUSE_BUTTON:0}

        self.undo_stack = deque()
        self.redo_stack = deque()

    def exec_cmd(self, cmd):
        cmd.do()

        if hasattr(cmd, 'undo'):
            self.undo_stack.append(cmd)

    def on_undo(self, _):
        self.undo()

    def undo(self):
        if len(self.undo_stack) > 0:
            cmd = self.undo_stack.pop()
            self.redo_stack.append(cmd)
            cmd.undo()

    def on_redo(self, _):
        self.redo()

    def redo(self):
        if len(self.redo_stack) > 0:
            cmd = self.redo_stack.pop()
            self.undo_stack.append(cmd)
            cmd.do()

    def on_key_press(self, event):
        if event.keysym == 'Shift_L':
            self.shift_pressed = True

    def shift_release(self, _):
        self.shift_pressed = False
        self.point1 = None

    def motion(self, event):
        self.mouse_x = event.x
        self.mouse_y = event.y

    def on_button_press(self, event):
        self.mouse_state[event.num] = True

    def on_button_release(self, event):
        self.mouse_state[event.num] = False

        if event.num == LEFT_MOUSE_BUTTON and self.shift_pressed:
            if self.point1 is not None:
                if (self.point1[0] == self.selected[0] or
                    self.point1[1] == self.selected[1]):

                    tile_lst = make_line_of_tiles(self.point1,
                    self.selected)

                    self.exec_cmd(ChangeTilesCommand(tile_lst,
                        self.tiles))

                    self.point1 = self.selected
            else:
                self.point1 = self.selected

    def toggle_grid(self, _):
        self.show_grid = not self.show_grid

    def update(self):
        self.selected = ((self.mouse_x / CELL_SIZE) * CELL_SIZE,
            (self.mouse_y / CELL_SIZE) * CELL_SIZE)

        if (self.mouse_state[LEFT_MOUSE_BUTTON] and
        self.selected not in self.tiles) and not self.shift_pressed:
            self.exec_cmd(ChangeTilesCommand([self.selected],
                self.tiles))

        if (self.mouse_state[RIGHT_MOUSE_BUTTON] and
        self.selected in self.tiles):
            self.exec_cmd(ChangeTilesCommand([self.selected],
                self.tiles, remove=True))

        if self.point1 is not None:
            if self.point1[0] == self.selected[0]:
                self.tile_line = make_line_of_tiles(self.point1,
                self.selected)
            elif self.point1[1] == self.selected[1]:
                self.tile_line = make_line_of_tiles(self.point1,
                self.selected)
            else:
                self.tile_line = None
        else:
            self.tile_line = None

    def draw_grid(self):
        for xpos in range(0, DISPLAY_WIDTH, CELL_SIZE):
            pygame.draw.line(self.screen, DARK_GRAY, (xpos, 0),
                (xpos, DISPLAY_HEIGHT))

        for ypos in range(0, DISPLAY_HEIGHT, CELL_SIZE):
            pygame.draw.line(self.screen, DARK_GRAY, (0, ypos),
                (DISPLAY_WIDTH, ypos))

    def draw(self):
        if self.show_grid:
            self.draw_grid()

        for tile in self.tiles:
            self.screen.blit(self.wall_tex, tile)

        pygame.draw.rect(self.screen, ORANGE,
            pygame.Rect(self.selected,
            (CELL_SIZE, CELL_SIZE)))

        if self.tile_line is not None:
            for tile in self.tile_line:
                self.screen.blit(self.wall_tex, tile)

    def run(self):
        while not self.quit:
            for _ in pygame.event.get():
                pass

            self.screen.fill(BLACK)
            self.fps_clock.tick(FPS)
            self.update()
            self.draw()
            pygame.display.flip()
            self.root.update()

    def exit_cmd(self):
        self.quit = True

def main():
    mapedit = MapEditor()
    mapedit.run()

if __name__ == '__main__':
    main()
<commit_msg>Minor improvement to mapedit<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

import Tkinter as tk
import os
from collections import deque

import pygame

from colors import BLACK, ORANGE, DARK_GRAY

FPS = 30

CELL_SIZE = 10
ROWS = 64
COLS = 128
DISPLAY_WIDTH = COLS * CELL_SIZE
DISPLAY_HEIGHT = ROWS * CELL_SIZE

LEFT_MOUSE_BUTTON = 1
RIGHT_MOUSE_BUTTON = 3

def make_line_of_tiles(point1, point2):
    tile_lst = []
    if point1[0] == point2[0]:
        startpoint = min(point1[1], point2[1])
        endpoint = max(point1[1], point2[1])
        for ypos in range(startpoint, endpoint+1, CELL_SIZE):
            tile_lst.append((point1[0], ypos))
    else:
        startpoint = min(point1[0], point2[0])
        endpoint = max(point1[0], point2[0])
        for xpos in range(startpoint, endpoint+1, CELL_SIZE):
            tile_lst.append((xpos, point1[1]))

    return tile_lst

class ChangeTilesCommand(object):
    def __init__(self, tiles, tilemap, remove=False):
        self.tile_lst = tiles
        self.tilemap = tilemap
        self.tiles_changed = []
        self.remove = remove

    def do(self):
        for tile in self.tile_lst:
            if self.remove:
                if tile in self.tilemap:
                    self.tilemap.remove(tile)
                    self.tiles_changed.append(tile)
            else:
                if tile not in self.tilemap:
                    self.tilemap.append(tile)
                    self.tiles_changed.append(tile)

    def undo(self):
        for tile in self.tiles_changed:
            if self.remove:
                self.tilemap.append(tile)
            else:
                self.tilemap.remove(tile)

class MapEditor(object):
    def __init__(self):
        self.root = tk.Tk()
        self.root.bind('<Motion>', self.motion)
        self.root.bind('<ButtonPress>', self.on_button_press)
        self.root.bind('<ButtonRelease>', self.on_button_release)
        self.root.bind('<Control-g>', self.toggle_grid)
        self.root.bind('<KeyPress>', self.on_key_press)
        self.root.bind('<Shift-KeyRelease>', self.shift_release)
        self.root.bind('<Control-z>', self.on_undo)
        self.root.bind('<Control-y>', self.on_redo)
        self.root.rowconfigure(0, weight=1)
        self.root.columnconfigure(0, weight=1)

        top = self.root.winfo_toplevel()
        self.menu_bar = tk.Menu(top)
        top['menu'] = self.menu_bar
        self.file_menu = tk.Menu(self.menu_bar)
        self.menu_bar.add_cascade(label='File', menu=self.file_menu)
        self.file_menu.add_command(label='new')
        self.file_menu.add_command(label='open')
        self.file_menu.add_command(label='save')
        self.file_menu.add_command(label='save as')
        self.file_menu.add_command(label='exit',
            command=self.exit_cmd)

        self.edit_menu = tk.Menu(self.menu_bar)
        self.menu_bar.add_cascade(label='Edit', menu=self.edit_menu)
        self.edit_menu.add_command(label='Undo', command=self.undo)
        self.edit_menu.add_command(label='Redo', command=self.redo)

        self.embed = tk.Frame(self.root, width=1400, height=700)
        self.embed.grid(row=0, column=0)

        os.environ['SDL_WINDOWID'] = str(self.embed.winfo_id())

        self.root.update()
        self.fps_clock = pygame.time.Clock()
        pygame.init()
        self.screen = pygame.display.set_mode((DISPLAY_WIDTH,
                DISPLAY_HEIGHT))

        self.quit = False

        self.selected = None
        self.mouse_x = 0
        self.mouse_y = 0

        self.wall_tex = pygame.image.load('../gfx/wall.png').convert()
        self.tiles = []

        self.show_grid = True

        self.shift_pressed = False
        self.point1 = None

        self.tile_line = None

        self.mouse_state = {LEFT_MOUSE_BUTTON:0, RIGHT_MOUSE_BUTTON:0}

        self.undo_stack = deque()
        self.redo_stack = deque()

    def exec_cmd(self, cmd):
        cmd.do()

        if hasattr(cmd, 'undo'):
            self.undo_stack.append(cmd)

    def on_undo(self, _):
        self.undo()

    def undo(self):
        if len(self.undo_stack) > 0:
            cmd = self.undo_stack.pop()
            self.redo_stack.append(cmd)
            cmd.undo()

    def on_redo(self, _):
        self.redo()

    def redo(self):
        if len(self.redo_stack) > 0:
            cmd = self.redo_stack.pop()
            self.undo_stack.append(cmd)
            cmd.do()

    def on_key_press(self, event):
        if event.keysym == 'Shift_L':
            self.shift_pressed = True

    def shift_release(self, _):
        self.shift_pressed = False
        self.point1 = None

    def motion(self, event):
        self.mouse_x = event.x
        self.mouse_y = event.y

    def on_button_press(self, event):
        self.mouse_state[event.num] = True

    def on_button_release(self, event):
        self.mouse_state[event.num] = False

        if event.num == LEFT_MOUSE_BUTTON and self.shift_pressed:
            if self.point1 is not None:
                if (self.point1[0] == self.selected[0] or
                    self.point1[1] == self.selected[1]):

                    tile_lst = make_line_of_tiles(self.point1,
                    self.selected)

                    self.exec_cmd(ChangeTilesCommand(tile_lst,
                        self.tiles))

                    self.point1 = self.selected
            else:
                self.point1 = self.selected

    def toggle_grid(self, _):
        self.show_grid = not self.show_grid

    def update(self):
        self.selected = ((self.mouse_x / CELL_SIZE) * CELL_SIZE,
            (self.mouse_y / CELL_SIZE) * CELL_SIZE)

        if (self.mouse_state[LEFT_MOUSE_BUTTON] and
        self.selected not in self.tiles) and not self.shift_pressed:
            self.exec_cmd(ChangeTilesCommand([self.selected],
                self.tiles))

        if (self.mouse_state[RIGHT_MOUSE_BUTTON] and
        self.selected in self.tiles):
            self.exec_cmd(ChangeTilesCommand([self.selected],
                self.tiles, remove=True))

        if self.point1 is not None:
            if self.point1[0] == self.selected[0]:
                self.tile_line = make_line_of_tiles(self.point1,
                self.selected)
            elif self.point1[1] == self.selected[1]:
                self.tile_line = make_line_of_tiles(self.point1,
                self.selected)
            else:
                self.tile_line = None
        else:
            self.tile_line = None

    def draw_grid(self):
        for xpos in range(0, DISPLAY_WIDTH, CELL_SIZE):
            pygame.draw.line(self.screen, DARK_GRAY, (xpos, 0),
                (xpos, DISPLAY_HEIGHT))

        for ypos in range(0, DISPLAY_HEIGHT, CELL_SIZE):
            pygame.draw.line(self.screen, DARK_GRAY, (0, ypos),
                (DISPLAY_WIDTH, ypos))

    def draw(self):
        if self.show_grid:
            self.draw_grid()

        for tile in self.tiles:
            self.screen.blit(self.wall_tex, tile)

        pygame.draw.rect(self.screen, ORANGE,
            pygame.Rect(self.selected,
            (CELL_SIZE, CELL_SIZE)))

        if self.tile_line is not None:
            for tile in self.tile_line:
                self.screen.blit(self.wall_tex, tile)

        if self.point1 is not None:
            self.screen.blit(self.wall_tex, self.point1)

    def run(self):
        while not self.quit:
            for _ in pygame.event.get():
                pass

            self.screen.fill(BLACK)
            self.fps_clock.tick(FPS)
            self.update()
            self.draw()
            pygame.display.flip()
            self.root.update()

    def exit_cmd(self):
        self.quit = True

def main():
    mapedit = MapEditor()
    mapedit.run()

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