<commit_before>import logging; _L = logging.getLogger('openaddr.ci')

from ..compat import standard_library
from ..jobs import JOB_TIMEOUT

from os.path import relpath, splitext
from urllib.parse import urljoin
from base64 import b64decode
from uuid import uuid4
import json, os

from flask import Flask, request, Response, current_app, jsonify
from uritemplate import expand
from requests import get, post
from psycopg2 import connect
from pq import PQ

def load_config():
    return dict(GITHUB_AUTH=(os.environ['GITHUB_TOKEN'], 'x-oauth-basic'),
                DATABASE_URL=os.environ['DATABASE_URL'])

app = Flask(__name__)
app.config.update(load_config())

MAGIC_OK_MESSAGE = 'Everything is fine'
TASK_QUEUE, DONE_QUEUE, DUE_QUEUE = 'tasks', 'finished', 'due'

@app.route('/')
def index():
    return 'Yo.'

@app.route('/hook', methods=['POST'])
def hook():
    github_auth = current_app.config['GITHUB_AUTH']
    webhook_payload = json.loads(request.data.decode('utf8'))
    files = process_payload_files(webhook_payload, github_auth)
    status_url = get_status_url(webhook_payload)

    if not files:
        update_empty_status(status_url, github_auth)
        return jsonify({'url': None, 'files': []})

    filenames = list(files.keys())
    job_url_template = urljoin(request.url, '/jobs/{id}')

    with db_connect(current_app.config['DATABASE_URL']) as conn:
        queue = db_queue(conn, TASK_QUEUE)
        try:
            job_id = create_queued_job(queue, files, job_url_template, status_url)
            job_url = expand(job_url_template, dict(id=job_id))
        except Exception as e:
            # Oops, tell Github something went wrong.
            update_error_status(status_url, str(e), filenames, github_auth)
            return Response(json.dumps({'error': str(e), 'files': files}),
                            500, content_type='application/json')
        else:
            # That worked, tell Github we're working on it.
            update_pending_status(status_url, job_url, filenames, github_auth)
            return jsonify({'id': job_id, 'url': job_url, 'files': files})

@app.route('/jobs/<job_id>', methods=['GET'])
def get_job(job_id):
    '''
    '''
    with db_connect(current_app.config['DATABASE_URL']) as conn:
        with db_cursor(conn) as db:
            try:
                status, task_files, file_states, file_results, github_status_url = read_job(db, job_id)
            except TypeError:
                return Response('Job {} not found'.format(job_id), 404)

    job = dict(status=status, task_files=task_files, file_states=file_states,
               file_results=file_results, github_status_url=github_status_url)

    return jsonify(job)

def get_touched_payload_files(payload):
    ''' Return a set of files modified in payload commits.
    '''
    touched = set()

    # Iterate over commits in chronological order.
    for commit in payload['commits']:
        for filelist in (commit['added'], commit['modified']):
            # Include any potentially-new files.
            touched.update(filelist)

        for filename in commit['removed']:
            # Skip files that no longer exist.
            touched.remove(filename)

    current_app.logger.debug('Touched files {}'.format(', '.join(touched)))

    return touched

def get_touched_branch_files(payload, github_auth):
    ''' Return a set of files modified between master and payload head.
    '''
    branch_sha = payload['head_commit']['id']

    compare1_url = payload['repository']['compare_url']
    compare1_url = expand(compare1_url, dict(base='master', head=branch_sha))
    current_app.logger.debug('Compare URL 1 {}'.format(compare1_url))

    compare1 = get(compare1_url, auth=github_auth).json()
    merge_base_sha = compare1['merge_base_commit']['sha']

    # That's no branch.
    if merge_base_sha == branch_sha:
        return set()

    compare2_url = payload['repository']['compare_url']
    compare2_url = expand(compare2_url, dict(base=merge_base_sha, head=branch_sha))
    current_app.logger.debug('Compare URL 2 {}'.format(compare2_url))

    compare2 = get(compare2_url, auth=github_auth).json()
    touched = set([file['filename'] for file in compare2['files']])
    current_app.logger.debug('Touched files {}'.format(', '.join(touched)))

    return touched

def process_payload_files(payload, github_auth):
    ''' Return a dictionary of file paths to raw JSON contents and file IDs.
    '''
    files = dict()

    touched = get_touched_payload_files(payload)
    touched |= get_touched_branch_files(payload, github_auth)

    commit_sha = payload['head_commit']['id']

    for filename in touched:
        if relpath(filename, 'sources').startswith('..'):
            # Skip things outside of sources directory.
            continue

        if splitext(filename)[1] != '.json':
            # Skip non-JSON files.
            continue

        contents_url = payload['repository']['contents_url']
        contents_url = expand(contents_url, dict(path=filename))
        contents_url = '{contents_url}?ref={commit_sha}'.format(**locals())

        current_app.logger.debug('Contents URL {}'.format(contents_url))

        got = get(contents_url, auth=github_auth)
        contents = got.json()

        if got.status_code not in range(200, 299):
            current_app.logger.warning('Skipping {} - {}'.format(filename, got.status_code))
            continue

        content, encoding = contents['content'], contents['encoding']

        current_app.logger.debug('Contents SHA {}'.format(contents['sha']))

        if encoding == 'base64':
            files[filename] = b64decode(content).decode('utf8'), contents['sha']
        else:
            raise ValueError('Unrecognized encoding "{}"'.format(encoding))

    return files

def get_status_url(payload):
    ''' Get Github status API URL from webhook payload.
    '''
    commit_sha = payload['head_commit']['id']
    status_url = payload['repository']['statuses_url']
    status_url = expand(status_url, dict(sha=commit_sha))

    current_app.logger.debug('Status URL {}'.format(status_url))

    return status_url

def post_github_status(status_url, status_json, github_auth):
    ''' POST status JSON to Github status API.
    '''
    # Github only wants 140 chars of description.
    status_json['description'] = status_json['description'][:140]

    posted = post(status_url, data=json.dumps(status_json), auth=github_auth,
                  headers={'Content-Type': 'application/json'})

    if posted.status_code not in range(200, 299):
        raise ValueError('Failed status post to {}'.format(status_url))

    if posted.json()['state'] != status_json['state']:
        raise ValueError('Mismatched status post to {}'.format(status_url))

def update_pending_status(status_url, job_url, filenames, github_auth):
    ''' Push pending status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='pending',
                  description='Checking {}'.format(', '.join(filenames)),
                  target_url=job_url)

    return post_github_status(status_url, status, github_auth)

def update_error_status(status_url, message, filenames, github_auth):
    ''' Push error status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='error',
                  description='Errored on {}: {}'.format(', '.join(filenames), message))

    return post_github_status(status_url, status, github_auth)

def update_failing_status(status_url, job_url, bad_files, filenames, github_auth):
    ''' Push failing status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='failure',
                  description='Failed on {} from {}'.format(', '.join(bad_files), ', '.join(filenames)),
                  target_url=job_url)

    return post_github_status(status_url, status, github_auth)

def update_empty_status(status_url, github_auth):
    ''' Push success status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='success',
                  description='Nothing to check')

    return post_github_status(status_url, status, github_auth)

def update_success_status(status_url, job_url, filenames, github_auth):
    ''' Push success status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='success',
                  description='Succeeded on {}'.format(', '.join(filenames)),
                  target_url=job_url)

    return post_github_status(status_url, status, github_auth)

def calculate_job_id(files):
    '''
    '''
    return str(uuid4())

    #
    # Previously, we created a deterministic hash of
    # the files, but for now that might be too optimistic.
    #
    blob = json.dumps(files, ensure_ascii=True, sort_keys=True)
    job_id = sha1(blob).hexdigest()

    return job_id

def create_queued_job(queue, files, job_url_template, status_url):
    ''' Create a new job, and add its files to the queue.
    '''
    filenames = list(files.keys())
    file_states = {name: None for name in filenames}
    file_results = {name: None for name in filenames}

    job_id = calculate_job_id(files)
    job_url = job_url_template and expand(job_url_template, dict(id=job_id))
    job_status = None

    with queue as db:
        task_files = add_files_to_queue(queue, job_id, job_url, files)
        add_job(db, job_id, None, task_files, file_states, file_results, status_url)

    return job_id

def add_files_to_queue(queue, job_id, job_url, files):
    ''' Make a new task for each file, return dict of file IDs to file names.
    '''
    tasks = {}

    for (file_name, (content, file_id)) in files.items():
        queue.put(dict(id=job_id, url=job_url, name=file_name,
                       content=content, file_id=file_id))

        tasks[file_id] = file_name

    return tasks

def add_job(db, job_id, status, task_files, file_states, file_results, status_url):
    ''' Save information about a job to the database.

        Throws an IntegrityError exception if the job ID exists.
    '''
    db.execute('''INSERT INTO jobs
                  (task_files, file_states, file_results, github_status_url, status, id)
                  VALUES (%s::json, %s::json, %s::json, %s, %s, %s)''',
               (json.dumps(task_files), json.dumps(file_states),
                json.dumps(file_results), status_url, status, job_id))

def write_job(db, job_id, status, task_files, file_states, file_results, status_url):
    ''' Save information about a job to the database.
    '''
    db.execute('''UPDATE jobs
                  SET task_files=%s::json, file_states=%s::json,
                      file_results=%s::json, github_status_url=%s, status=%s
                  WHERE id = %s''',
               (json.dumps(task_files), json.dumps(file_states),
                json.dumps(file_results), status_url, status, job_id))

def read_job(db, job_id):
    ''' Read information about a job from the database.

        Returns (status, task_files, file_states, file_results, github_status_url) or None.
    '''
    db.execute('''SELECT status, task_files, file_states, file_results, github_status_url
                  FROM jobs WHERE id = %s''', (job_id, ))

    try:
        status, task_files, states, file_results, github_status_url = db.fetchone()
    except TypeError:
        return None
    else:
        return status, task_files, states, file_results, github_status_url

def pop_task_from_taskqueue(task_queue, done_queue, due_queue, output_dir):
    '''
    '''
    with task_queue as db:
        task = task_queue.get()

        # PQ will return NULL after 1 second timeout if not ask
        if task is None:
            return

    _L.info("Got job {}".format(task.data))

    # Send a Due task, possibly for later.
    due_task_data = dict(task_data=task.data, file_id=task.data['file_id'])
    due_queue.put(due_task_data, JOB_TIMEOUT)

    # Run the task.
    from .worker import do_work # <-- TODO: un-suck this.
    result = do_work(task.data['id'], task.data['content'], output_dir)

    # Send a Done task
    done_task_data = {k: task.data[k] for k in ('id', 'url', 'name')}
    done_task_data['result'] = result
    done_queue.put(done_task_data)

def pop_task_from_donequeue(queue, github_auth):
    ''' Look for a completed job in the "done" task queue, update Github status.
    '''
    with queue as db:
        task = queue.get()

        if task is None:
            return

        results = task.data['result']
        message = results['message']
        run_state = results.get('output', None)
        job_url = task.data['url']
        filename = task.data['name']
        file_id = task.data['file_id']
        job_id = task.data['id']

        #
        # Add to the runs table.
        #
        db.execute('''INSERT INTO runs
                      (source_path, source_id, state, datetime)
                      VALUES (%s, %s, %s::json, NOW())''',
                   (filename, file_id, json.dumps(run_state)))

        try:
            _, task_files, file_states, file_results, status_url = read_job(db, job_id)
        except TypeError:
            raise Exception('Job {} not found'.format(job_id))

        if filename not in file_states:
            raise Exception('Unknown file from job {}: "{}"'.format(job_id, filename))

        filenames = list(task_files.values())
        file_states[filename] = bool(message == MAGIC_OK_MESSAGE)
        file_results[filename] = results

        if False in file_states.values():
            # Any task failure means the whole job has failed.
            job_status = False
        elif None in file_states.values():
            job_status = None
        else:
            job_status = True

        write_job(db, job_id, job_status, task_files, file_states, file_results, status_url)

        if not status_url:
            return

        if job_status is False:
            bad_files = [name for (name, state) in file_states.items() if state is False]
            update_failing_status(status_url, job_url, bad_files, filenames, github_auth)

        elif job_status is None:
            update_pending_status(status_url, job_url, filenames, github_auth)

        elif job_status is True:
            update_success_status(status_url, job_url, filenames, github_auth)

def db_connect(dsn):
    ''' Connect to database using DSN string.
    '''
    return connect(dsn)

def db_queue(conn, name):
    return PQ(conn, table='queue')[name]

def db_cursor(conn):
    return conn.cursor()

if __name__ == '__main__':
    app.run(debug=True)
<commit_msg>Started using queue.put() expected_at argument for priority<commit_after>import logging; _L = logging.getLogger('openaddr.ci')

from ..compat import standard_library
from ..jobs import JOB_TIMEOUT

from os.path import relpath, splitext
from urllib.parse import urljoin
from datetime import timedelta
from base64 import b64decode
from uuid import uuid4
import json, os

from flask import Flask, request, Response, current_app, jsonify
from uritemplate import expand
from requests import get, post
from psycopg2 import connect
from pq import PQ

def load_config():
    return dict(GITHUB_AUTH=(os.environ['GITHUB_TOKEN'], 'x-oauth-basic'),
                DATABASE_URL=os.environ['DATABASE_URL'])

app = Flask(__name__)
app.config.update(load_config())

MAGIC_OK_MESSAGE = 'Everything is fine'
TASK_QUEUE, DONE_QUEUE, DUE_QUEUE = 'tasks', 'finished', 'due'

@app.route('/')
def index():
    return 'Yo.'

@app.route('/hook', methods=['POST'])
def hook():
    github_auth = current_app.config['GITHUB_AUTH']
    webhook_payload = json.loads(request.data.decode('utf8'))
    files = process_payload_files(webhook_payload, github_auth)
    status_url = get_status_url(webhook_payload)

    if not files:
        update_empty_status(status_url, github_auth)
        return jsonify({'url': None, 'files': []})

    filenames = list(files.keys())
    job_url_template = urljoin(request.url, '/jobs/{id}')

    with db_connect(current_app.config['DATABASE_URL']) as conn:
        queue = db_queue(conn, TASK_QUEUE)
        try:
            job_id = create_queued_job(queue, files, job_url_template, status_url)
            job_url = expand(job_url_template, dict(id=job_id))
        except Exception as e:
            # Oops, tell Github something went wrong.
            update_error_status(status_url, str(e), filenames, github_auth)
            return Response(json.dumps({'error': str(e), 'files': files}),
                            500, content_type='application/json')
        else:
            # That worked, tell Github we're working on it.
            update_pending_status(status_url, job_url, filenames, github_auth)
            return jsonify({'id': job_id, 'url': job_url, 'files': files})

@app.route('/jobs/<job_id>', methods=['GET'])
def get_job(job_id):
    '''
    '''
    with db_connect(current_app.config['DATABASE_URL']) as conn:
        with db_cursor(conn) as db:
            try:
                status, task_files, file_states, file_results, github_status_url = read_job(db, job_id)
            except TypeError:
                return Response('Job {} not found'.format(job_id), 404)

    job = dict(status=status, task_files=task_files, file_states=file_states,
               file_results=file_results, github_status_url=github_status_url)

    return jsonify(job)

def get_touched_payload_files(payload):
    ''' Return a set of files modified in payload commits.
    '''
    touched = set()

    # Iterate over commits in chronological order.
    for commit in payload['commits']:
        for filelist in (commit['added'], commit['modified']):
            # Include any potentially-new files.
            touched.update(filelist)

        for filename in commit['removed']:
            # Skip files that no longer exist.
            touched.remove(filename)

    current_app.logger.debug('Touched files {}'.format(', '.join(touched)))

    return touched

def get_touched_branch_files(payload, github_auth):
    ''' Return a set of files modified between master and payload head.
    '''
    branch_sha = payload['head_commit']['id']

    compare1_url = payload['repository']['compare_url']
    compare1_url = expand(compare1_url, dict(base='master', head=branch_sha))
    current_app.logger.debug('Compare URL 1 {}'.format(compare1_url))

    compare1 = get(compare1_url, auth=github_auth).json()
    merge_base_sha = compare1['merge_base_commit']['sha']

    # That's no branch.
    if merge_base_sha == branch_sha:
        return set()

    compare2_url = payload['repository']['compare_url']
    compare2_url = expand(compare2_url, dict(base=merge_base_sha, head=branch_sha))
    current_app.logger.debug('Compare URL 2 {}'.format(compare2_url))

    compare2 = get(compare2_url, auth=github_auth).json()
    touched = set([file['filename'] for file in compare2['files']])
    current_app.logger.debug('Touched files {}'.format(', '.join(touched)))

    return touched

def process_payload_files(payload, github_auth):
    ''' Return a dictionary of file paths to raw JSON contents and file IDs.
    '''
    files = dict()

    touched = get_touched_payload_files(payload)
    touched |= get_touched_branch_files(payload, github_auth)

    commit_sha = payload['head_commit']['id']

    for filename in touched:
        if relpath(filename, 'sources').startswith('..'):
            # Skip things outside of sources directory.
            continue

        if splitext(filename)[1] != '.json':
            # Skip non-JSON files.
            continue

        contents_url = payload['repository']['contents_url']
        contents_url = expand(contents_url, dict(path=filename))
        contents_url = '{contents_url}?ref={commit_sha}'.format(**locals())

        current_app.logger.debug('Contents URL {}'.format(contents_url))

        got = get(contents_url, auth=github_auth)
        contents = got.json()

        if got.status_code not in range(200, 299):
            current_app.logger.warning('Skipping {} - {}'.format(filename, got.status_code))
            continue

        content, encoding = contents['content'], contents['encoding']

        current_app.logger.debug('Contents SHA {}'.format(contents['sha']))

        if encoding == 'base64':
            files[filename] = b64decode(content).decode('utf8'), contents['sha']
        else:
            raise ValueError('Unrecognized encoding "{}"'.format(encoding))

    return files

def get_status_url(payload):
    ''' Get Github status API URL from webhook payload.
    '''
    commit_sha = payload['head_commit']['id']
    status_url = payload['repository']['statuses_url']
    status_url = expand(status_url, dict(sha=commit_sha))

    current_app.logger.debug('Status URL {}'.format(status_url))

    return status_url

def post_github_status(status_url, status_json, github_auth):
    ''' POST status JSON to Github status API.
    '''
    # Github only wants 140 chars of description.
    status_json['description'] = status_json['description'][:140]

    posted = post(status_url, data=json.dumps(status_json), auth=github_auth,
                  headers={'Content-Type': 'application/json'})

    if posted.status_code not in range(200, 299):
        raise ValueError('Failed status post to {}'.format(status_url))

    if posted.json()['state'] != status_json['state']:
        raise ValueError('Mismatched status post to {}'.format(status_url))

def update_pending_status(status_url, job_url, filenames, github_auth):
    ''' Push pending status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='pending',
                  description='Checking {}'.format(', '.join(filenames)),
                  target_url=job_url)

    return post_github_status(status_url, status, github_auth)

def update_error_status(status_url, message, filenames, github_auth):
    ''' Push error status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='error',
                  description='Errored on {}: {}'.format(', '.join(filenames), message))

    return post_github_status(status_url, status, github_auth)

def update_failing_status(status_url, job_url, bad_files, filenames, github_auth):
    ''' Push failing status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='failure',
                  description='Failed on {} from {}'.format(', '.join(bad_files), ', '.join(filenames)),
                  target_url=job_url)

    return post_github_status(status_url, status, github_auth)

def update_empty_status(status_url, github_auth):
    ''' Push success status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='success',
                  description='Nothing to check')

    return post_github_status(status_url, status, github_auth)

def update_success_status(status_url, job_url, filenames, github_auth):
    ''' Push success status for head commit to Github status API.
    '''
    status = dict(context='openaddresses/hooked', state='success',
                  description='Succeeded on {}'.format(', '.join(filenames)),
                  target_url=job_url)

    return post_github_status(status_url, status, github_auth)

def calculate_job_id(files):
    '''
    '''
    return str(uuid4())

    #
    # Previously, we created a deterministic hash of
    # the files, but for now that might be too optimistic.
    #
    blob = json.dumps(files, ensure_ascii=True, sort_keys=True)
    job_id = sha1(blob).hexdigest()

    return job_id

def create_queued_job(queue, files, job_url_template, status_url):
    ''' Create a new job, and add its files to the queue.
    '''
    filenames = list(files.keys())
    file_states = {name: None for name in filenames}
    file_results = {name: None for name in filenames}

    job_id = calculate_job_id(files)
    job_url = job_url_template and expand(job_url_template, dict(id=job_id))
    job_status = None

    with queue as db:
        task_files = add_files_to_queue(queue, job_id, job_url, files)
        add_job(db, job_id, None, task_files, file_states, file_results, status_url)

    return job_id

def add_files_to_queue(queue, job_id, job_url, files):
    ''' Make a new task for each file, return dict of file IDs to file names.
    '''
    tasks = {}

    for (file_name, (content, file_id)) in files.items():
        task_data = dict(id=job_id, url=job_url, name=file_name,
                         content=content, file_id=file_id)

        queue.put(task_data, expected_at=timedelta(0))
        tasks[file_id] = file_name

    return tasks

def add_job(db, job_id, status, task_files, file_states, file_results, status_url):
    ''' Save information about a job to the database.

        Throws an IntegrityError exception if the job ID exists.
    '''
    db.execute('''INSERT INTO jobs
                  (task_files, file_states, file_results, github_status_url, status, id)
                  VALUES (%s::json, %s::json, %s::json, %s, %s, %s)''',
               (json.dumps(task_files), json.dumps(file_states),
                json.dumps(file_results), status_url, status, job_id))

def write_job(db, job_id, status, task_files, file_states, file_results, status_url):
    ''' Save information about a job to the database.
    '''
    db.execute('''UPDATE jobs
                  SET task_files=%s::json, file_states=%s::json,
                      file_results=%s::json, github_status_url=%s, status=%s
                  WHERE id = %s''',
               (json.dumps(task_files), json.dumps(file_states),
                json.dumps(file_results), status_url, status, job_id))

def read_job(db, job_id):
    ''' Read information about a job from the database.

        Returns (status, task_files, file_states, file_results, github_status_url) or None.
    '''
    db.execute('''SELECT status, task_files, file_states, file_results, github_status_url
                  FROM jobs WHERE id = %s''', (job_id, ))

    try:
        status, task_files, states, file_results, github_status_url = db.fetchone()
    except TypeError:
        return None
    else:
        return status, task_files, states, file_results, github_status_url

def pop_task_from_taskqueue(task_queue, done_queue, due_queue, output_dir):
    '''
    '''
    with task_queue as db:
        task = task_queue.get()

        # PQ will return NULL after 1 second timeout if not ask
        if task is None:
            return

    _L.info("Got job {}".format(task.data))

    # Send a Due task, possibly for later.
    due_task_data = dict(task_data=task.data, file_id=task.data['file_id'])
    due_queue.put(due_task_data, schedule_at=JOB_TIMEOUT)

    # Run the task.
    from .worker import do_work # <-- TODO: un-suck this.
    result = do_work(task.data['id'], task.data['content'], output_dir)

    # Send a Done task
    done_task_data = {k: task.data[k] for k in ('id', 'url', 'name')}
    done_task_data['result'] = result
    done_queue.put(done_task_data, expected_at=timedelta(0))

def pop_task_from_donequeue(queue, github_auth):
    ''' Look for a completed job in the "done" task queue, update Github status.
    '''
    with queue as db:
        task = queue.get()

        if task is None:
            return

        results = task.data['result']
        message = results['message']
        run_state = results.get('output', None)
        job_url = task.data['url']
        filename = task.data['name']
        file_id = task.data['file_id']
        job_id = task.data['id']

        #
        # Add to the runs table.
        #
        db.execute('''INSERT INTO runs
                      (source_path, source_id, state, datetime)
                      VALUES (%s, %s, %s::json, NOW())''',
                   (filename, file_id, json.dumps(run_state)))

        try:
            _, task_files, file_states, file_results, status_url = read_job(db, job_id)
        except TypeError:
            raise Exception('Job {} not found'.format(job_id))

        if filename not in file_states:
            raise Exception('Unknown file from job {}: "{}"'.format(job_id, filename))

        filenames = list(task_files.values())
        file_states[filename] = bool(message == MAGIC_OK_MESSAGE)
        file_results[filename] = results

        if False in file_states.values():
            # Any task failure means the whole job has failed.
            job_status = False
        elif None in file_states.values():
            job_status = None
        else:
            job_status = True

        write_job(db, job_id, job_status, task_files, file_states, file_results, status_url)

        if not status_url:
            return

        if job_status is False:
            bad_files = [name for (name, state) in file_states.items() if state is False]
            update_failing_status(status_url, job_url, bad_files, filenames, github_auth)

        elif job_status is None:
            update_pending_status(status_url, job_url, filenames, github_auth)

        elif job_status is True:
            update_success_status(status_url, job_url, filenames, github_auth)

def db_connect(dsn):
    ''' Connect to database using DSN string.
    '''
    return connect(dsn)

def db_queue(conn, name):
    return PQ(conn, table='queue')[name]

def db_cursor(conn):
    return conn.cursor()

if __name__ == '__main__':
    app.run(debug=True)
<|endoftext|>


<commit_before>from beautifulhue.api import Bridge
from time import sleep
from sys import stdout
from numpy import clip, interp
from lib import ColorHelper, Converter
import os
import requests
import datetime
import colorsys
import copy
import argparse

# Number of seconds to wait between requests
interval = 60

# Transition time in ms
transition_time = 2000

# WOEIDs for Melbourne, New York, and Paris. WOEID Lookup: http://woeid.rosselliot.co.nz/lookup
locations = [1103816, 2459115, 615702]

# Minimum temperature
temp_min = 5

# Maximum temperature
temp_max = 35

# Start hue: Blue
hue_min = 175

# End hue: Red
hue_max = 360

# Default saturation
saturation = 100

# Default brightness
brightness = 100

# Yahoo Weather API endpoint
endpoint = 'https://query.yahooapis.com/v1/public/yql'

# Querystring object
querystring = { 'q': 'select * from weather.forecast where woeid={0}', 'format': 'json', 'env': 'store://datatables.org/alltableswithkeys' }

# Colour converter
converter = Converter()

# IP of your Hue Bridge. nano ~/.bash_profile and add: export HUE_BRIDGE_IP="XXX.XXX.XXX.XXX"
ip = os.environ['HUE_BRIDGE_IP']

# Username set up on your Hue Bridge.
username = os.environ['HUE_USERNAME']

# Bridge object
bridge = Bridge(device={ 'ip': ip }, user={ 'name': username })

# Colour map
hsvmap = [	{ 'text': ['Thunderstorms', 'Rain'], 'saturation': 0.01, 'brightness': 0.01 },
	{ 'text': ['Cloudy', 'Mostly Cloudy', 'Showers'], 'saturation': 0.5, 'brightness': 0.05 },
	{ 'text': ['Partly Cloudy', 'Scattered Showers'], 'saturation': 0.95, 'brightness': 0.4 },
	{ 'text': ['Mostly Sunny'], 'saturation': 0.9, 'brightness': 0.5 },
	{ 'text': ['Clear', 'Breezy'], 'saturation': 1.0, 'brightness': 0.8 },
	{ 'text': ['Sunny'], 'saturation': 1.0, 'brightness': 1.0} ]

# Convert farenheit to celcius because Americans are barbarians
def farenheit_to_celsius(temp):
	return (temp-32)/1.8

# Convert temperature to HSV
def temp_to_hsv(temp):
	return (interp(clip(temp, temp_min, temp_max), [temp_min, temp_max], [hue_min, hue_max]), saturation, brightness)

def run(debug):
	try:
		globe = 1
		# Loop through each location
		for woeid in locations:
			# Deep copy the querystring object
			query = copy.deepcopy(querystring)
			query['q'] = query['q'].format(woeid)

			# Send the request
			response = requests.get(endpoint, params=query)

			# Get the response
			json = response.json()

			# Get the first day of the forecast
			day = json['query']['results']['channel']['item']['forecast'][0]

			# Get the location
			location = json['query']['results']['channel']['location']['city']

			text = day['text']
			date = day['date']
			high = int(day['high'])
			temp = farenheit_to_celsius(high)
			h,s,b = temp_to_hsv(temp)

			# Map the weather description to the colour
			lookup = [item for item in hsvmap if text in item['text']]

			# If there's a match, then change the saturation and brightness
			if lookup:
				s = s*lookup[0]['saturation']
				b = b*lookup[0]['brightness']

			# Convert to RGB
			red,green,blue = colorsys.hsv_to_rgb(h/360, b/100, s/100)

			# Convert to XY
			xy = converter.rgbToCIE1931(red,green,blue)

			# Debugging
			if debug:
				print(location)
				print(date)
				print(text)
				print("%.1fc" % round(temp,1))
				print('-----')

			# JSON sent to the Hue Bridge
			resource = {
				'which': globe,
				'data':{
					'state':{ 'on': True, 'xy':xy, 'bri': int((b/100)*254), 'transitiontime': (transition_time/100) }
				}
			}

			# Update the globe
			bridge.light.update(resource)
			globe+=1

	except Exception as e:
		print("Exception: %s" % str(e))

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('-d', action='store_true')
	args = parser.parse_args()

	while True:
		run(args.d)
		sleep(interval)
<commit_msg>Removed tab<commit_after>from beautifulhue.api import Bridge
from time import sleep
from sys import stdout
from numpy import clip, interp
from lib import ColorHelper, Converter
import os
import requests
import datetime
import colorsys
import copy
import argparse

# Number of seconds to wait between requests
interval = 60

# Transition time in ms
transition_time = 2000

# WOEIDs for Melbourne, New York, and Paris. WOEID Lookup: http://woeid.rosselliot.co.nz/lookup
locations = [1103816, 2459115, 615702]

# Minimum temperature
temp_min = 5

# Maximum temperature
temp_max = 35

# Start hue: Blue
hue_min = 175

# End hue: Red
hue_max = 360

# Default saturation
saturation = 100

# Default brightness
brightness = 100

# Yahoo Weather API endpoint
endpoint = 'https://query.yahooapis.com/v1/public/yql'

# Querystring object
querystring = { 'q': 'select * from weather.forecast where woeid={0}', 'format': 'json', 'env': 'store://datatables.org/alltableswithkeys' }

# Colour converter
converter = Converter()

# IP of your Hue Bridge. nano ~/.bash_profile and add: export HUE_BRIDGE_IP="XXX.XXX.XXX.XXX"
ip = os.environ['HUE_BRIDGE_IP']

# Username set up on your Hue Bridge.
username = os.environ['HUE_USERNAME']

# Bridge object
bridge = Bridge(device={ 'ip': ip }, user={ 'name': username })

# Colour map
hsvmap = [ { 'text': ['Thunderstorms', 'Rain'], 'saturation': 0.01, 'brightness': 0.01 },
	{ 'text': ['Cloudy', 'Mostly Cloudy', 'Showers'], 'saturation': 0.5, 'brightness': 0.05 },
	{ 'text': ['Partly Cloudy', 'Scattered Showers'], 'saturation': 0.95, 'brightness': 0.4 },
	{ 'text': ['Mostly Sunny'], 'saturation': 0.9, 'brightness': 0.5 },
	{ 'text': ['Clear', 'Breezy'], 'saturation': 1.0, 'brightness': 0.8 },
	{ 'text': ['Sunny'], 'saturation': 1.0, 'brightness': 1.0} ]

# Convert farenheit to celcius because Americans are barbarians
def farenheit_to_celsius(temp):
	return (temp-32)/1.8

# Convert temperature to HSV
def temp_to_hsv(temp):
	return (interp(clip(temp, temp_min, temp_max), [temp_min, temp_max], [hue_min, hue_max]), saturation, brightness)

def run(debug):
	try:
		globe = 1
		# Loop through each location
		for woeid in locations:
			# Deep copy the querystring object
			query = copy.deepcopy(querystring)
			query['q'] = query['q'].format(woeid)

			# Send the request
			response = requests.get(endpoint, params=query)

			# Get the response
			json = response.json()

			# Get the first day of the forecast
			day = json['query']['results']['channel']['item']['forecast'][0]

			# Get the location
			location = json['query']['results']['channel']['location']['city']

			text = day['text']
			date = day['date']
			high = int(day['high'])
			temp = farenheit_to_celsius(high)
			h,s,b = temp_to_hsv(temp)

			# Map the weather description to the colour
			lookup = [item for item in hsvmap if text in item['text']]

			# If there's a match, then change the saturation and brightness
			if lookup:
				s = s*lookup[0]['saturation']
				b = b*lookup[0]['brightness']

			# Convert to RGB
			red,green,blue = colorsys.hsv_to_rgb(h/360, b/100, s/100)

			# Convert to XY
			xy = converter.rgbToCIE1931(red,green,blue)

			# Debugging
			if debug:
				print(location)
				print(date)
				print(text)
				print("%.1fc" % round(temp,1))
				print('-----')

			# JSON sent to the Hue Bridge
			resource = {
				'which': globe,
				'data':{
					'state':{ 'on': True, 'xy':xy, 'bri': int((b/100)*254), 'transitiontime': (transition_time/100) }
				}
			}

			# Update the globe
			bridge.light.update(resource)
			globe+=1

	except Exception as e:
		print("Exception: %s" % str(e))

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('-d', action='store_true')
	args = parser.parse_args()

	while True:
		run(args.d)
		sleep(interval)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from AbstractVisitor import AbstractVisitor

from duralex.alinea_parser import *
from duralex.AddCommitMessageVisitor import int_to_roman
import duralex.node_type
import duralex.diff

from bs4 import BeautifulSoup
import jinja2

import os
import subprocess
import tempfile
from distutils.dir_util import copy_tree

class CreateGitBookVisitor(AbstractVisitor):
    def __init__(self, args):
        self.gitbook_dir = args.gitbook
        self.tmp_dir = tempfile.mkdtemp()
        self.formats = args.gitbook_format

        super(CreateGitBookVisitor, self).__init__()

    def write_file(self, filename, data):
        f = open(self.tmp_dir + '/' + filename, 'w')
        f.write(data.encode('utf-8'))
        f.close()

    def get_article_commit_title(self, node):
        ancestors = get_node_ancestors(node)
        messages = []
        for ancestor in ancestors:
            if 'type' not in ancestor:
                continue;
            if ancestor['type'] == 'article':
                messages.append('Article ' + str(ancestor['order']))
                if 'githubIssue' in node:
                    link = ancestor['githubIssue']
            if ancestor['type'] == 'bill-header1' and 'implicit' not in ancestor:
                messages.append(int_to_roman(ancestor['order']))
            if ancestor['type'] == 'bill-header2':
                messages.append(unicode(ancestor['order']) + u'°')
        return ', '.join(messages[::-1])

    def get_article_commit_diff(self, edit, target_title, target_href):
        if 'htmlDiff' in edit:
            soup = BeautifulSoup(edit['htmlDiff'], "html5lib")
            filename_div = soup.find('div', {'class': 'diff-filename'})
            a_tag = soup.new_tag('a', href=target_href)
            a_tag.string = target_title
            filename_div.string = ''
            filename_div.append(a_tag)
            return unicode(soup.body.div)
        elif 'diff' in edit:
            process = subprocess.Popen(
                'diff2html -i stdin -d word -o stdout --su hidden -s line',
                shell=True,
                stdout=subprocess.PIPE,
                stdin=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            out, err = process.communicate(input=edit['diff'].encode('utf-8') + '\n')
            soup = BeautifulSoup(out, "html5lib")
            return (str(list(soup.find_all('style'))[0]) + '\n\n'
                + unicode(soup.find('div', {'id': 'diff'})))

    def get_articles(self, node):
        articles = []
        article_nodes = filter_nodes(node, lambda n: n['type'] == 'article')
        for article_node in article_nodes:
            edit_nodes = filter_nodes(article_node, lambda n: 'type' in n and n['type'] == 'edit')
            commits = []
            for edit_node in edit_nodes:
                article_ref = filter_nodes(edit_node, lambda n: n['type'] == 'article-reference')[0]
                target_title, target_href = self.get_deep_link(self.get_edit_target_nodes(article_ref))
                commits.append({
                    'title': self.get_article_commit_title(edit_node),
                    # remove the " ({reference list})" from the commit message since its already printed
                    # in the header above
                    'description': re.sub(r' \(.*\)', '', edit_node['commitMessage'].splitlines()[0]) if 'commitMessage' in edit_node else None,
                    'diff': self.get_article_commit_diff(edit_node, target_title, target_href),
                    'target': {
                        'title': target_title,
                        'link': target_href
                    }
                })

            articles.append({
                'order': article_node['order'],
                'content': article_node['content'].replace('\n', '\n\n'),
                'commits': commits,
                'githubIssue': article_node['githubIssue'] if 'githubIssue' in article_node else None
            })
        return articles

    def merge_dicts(self, *dict_args):
        """
        Given any number of dicts, shallow copy and merge into a new dict,
        precedence goes to key value pairs in latter dicts.
        """
        result = {}
        for dictionary in dict_args:
            result.update(dictionary)
        return result

    def visit_node(self, node):
        super(CreateGitBookVisitor, self).visit_node(node)

        if 'parent' not in node:
            edits = self.build_edit_matrix(node)
            articles = self.get_articles(node)
            modified_texts = self.get_modified_texts(edits)
            template_data = {
                'title': self.get_book_title(node),
                'url': node['url'],
                'type': node['type'],
                'modified': modified_texts,
                'articles': articles
            }

            if 'cocoricoVote' in node:
                template_data['cocorico_vote'] = node['cocoricoVote']

            self.template_file('book.json', template_data)
            self.template_file('styles/website.css', template_data)
            self.template_file('SUMMARY.md', template_data)
            self.template_file('README.md', template_data)
            current_article = 0
            for article in articles:
                self.template_file(
                    'article.md',
                    self.merge_dicts(template_data, {'current_article': current_article}),
                    'article-' + str(article['order']) + '.md'
                )
                current_article += 1

            current_article = 0
            current_law = 0
            for modified in modified_texts:
                self.template_file(
                    'law.md',
                    self.merge_dicts(template_data, {
                        'current_law': current_law,
                    }),
                    modified['law'] + '.md'
                )
                for article in modified['articles']:
                    self.template_file(
                        'text.md',
                        self.merge_dicts(template_data, {
                            'current_law': current_law,
                            'current_article': current_article
                        }),
                        modified['law'] + '-' + article['id'] + '.md'
                    )
                    current_article += 1
                current_law += 1

            if 'html' in self.formats:
                self.cmd('gitbook install')
                self.cmd('gitbook build')

                if 'markdown' in self.formats:
                    copy_tree(self.tmp_dir, self.gitbook_dir)
                else:
                    copy_tree(os.path.join(self.tmp_dir, '_book'), self.gitbook_dir)
            else:
                copy_tree(self.tmp_dir, self.gitbook_dir)

    def cmd(self, command):
        process = subprocess.Popen(
            command,
            cwd=self.tmp_dir,
            shell=True,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        return process.communicate()

    def get_book_title(self, root_node):
        title = root_node['type'].title()
        if 'id' in root_node:
            title += u' N°' + str(root_node['id'])
        if 'legislature' in root_node:
            title += ', ' + str(root_node['legislature']) + u'ème législature'
        return title

    def patch(self, original, unified_diff):
        fd, filename = input_file = tempfile.mkstemp()
        os.write(fd, original.encode('utf-8'))
        process = subprocess.Popen(
            'patch -p0 --output=- ' + filename,
            shell=True,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        out, err = process.communicate(input=unified_diff.encode('utf-8') + '\n')
        return ''.join(out).decode('utf-8')

    def get_deep_link(self, nodes):
        href = []
        title = []
        for node in nodes:
            if node['type'] == 'law-reference':
                title.append(u'Loi N°' + node['lawId'])
                href.append(node['lawId'])
            elif node['type'] == 'article':
                title.append(u'Article ' + str(node['order']))
                href.append(u'article-' + str(node['order']) + '.md#article-' + str(node['order']))
            elif node['type'] == 'article-reference':
                title.append(u'Article ' + node['id'])
                href.append(node['id'] + '.md')
            elif node['type'] == 'bill-header1' and 'implicit' not in node:
                title.append(int_to_roman(node['order']))
                href.append(int_to_roman(node['order']))
            elif node['type'] == 'bill-header2':
                title.append(unicode(node['order']) + u'°')
                href.append(str(node['order']) + u'°')
        return (', '.join(title), '-'.join(href))

    def get_edit_target_nodes(self, node):
        nodes = []

        if node_type.is_reference(node):
            nodes.append(node)

        nodes += filter(
            lambda n: node_type.is_reference(n),
            get_node_ancestors(node)
        )

        return sorted(
            nodes,
            key=lambda n: node_type.REFERENCE.index(n['type'])
        )

    def get_edit_source_nodes(self, node):
        edit_source_types = [
            'article',
            'bill-header1',
            'bill-header2'
        ]

        return sorted(
            filter(
                lambda n: 'type' in n and n['type'] in edit_source_types,
                get_node_ancestors(node)
            ),
            key=lambda n: edit_source_types.index(n['type'])
        )

    def get_modified_texts(self, edits):
        modified = []
        law_ids = set([i[0] for i in edits.keys()])
        for law_id in law_ids:
            law_edits = {k: v for k, v in edits.iteritems() if k[0] == law_id}
            articles = []
            for k, v in edits.iteritems():
                law_ref = filter_nodes(v[0][-1], lambda n: n['type'] == 'law-reference' and n['lawId'] == k[0])[0]
                article_ref = filter_nodes(law_ref, lambda n: n['type'] == 'article-reference' and n['id'] == k[1])[0]

                f = open(article_ref['filename'], 'r')
                text = f.read().decode('utf-8')
                original_text = text
                f.close()

                commits = []
                for edit_source in v:
                    title, href = self.get_deep_link(edit_source)
                    commits.append({'title': title, 'link': href})
                    edit_refs = filter_nodes(edit_source[-1], lambda n: n['type'] == 'edit')
                    for edit_ref in edit_refs:
                        text = self.patch(text, edit_ref['diff'])
                article = {
                    'id': k[1],
                    'diff': duralex.diff.make_html_rich_diff(original_text, text),
                    'commits': commits
                }
                if 'gitlabHistory' in article_ref:
                    article['gitlabHistory'] = article_ref['gitlabHistory']
                if 'githubHistory' in article_ref:
                    article['githubHistory'] = article_ref['githubHistory']
                articles.append(article)
            articles = sorted(articles, key=lambda x: x['id'].replace('-', ' '))
            modified.append({'law': law_id, 'articles': articles})
        return modified

    def build_edit_matrix(self, node):
        edits = {}
        article_refs = filter_nodes(node, lambda n: 'type' in n and n['type'] == 'article-reference')
        for article_ref in article_refs:
            law_ref = filter(
                lambda n: 'type' in n and n['type'] == 'law-reference',
                get_node_ancestors(article_ref)
            )[0]
            t = (law_ref['lawId'], article_ref['id'])
            if t not in edits:
                edits[t] = []
            edits[t].append(self.get_edit_source_nodes(article_ref))
        return edits

    def template_file(self, template, values, out=None):
        f = open(os.path.join('./template/gitbook', template), 'r')
        t = jinja2.Environment(loader=jinja2.FileSystemLoader('./template/gitbook')).from_string(f.read().decode('utf-8'))
        f.close()

        filename = os.path.join(self.tmp_dir, out if out else template)
        path = os.path.dirname(filename)
        if not os.path.exists(path):
            os.makedirs(path)
        f = open(filename, 'w')
        f.write(t.render(values).encode('utf-8'))
        f.truncate()
        f.close()
<commit_msg>Fix broken reference to the gitlabIssue used in the Gitbook template.<commit_after># -*- coding: utf-8 -*-

from AbstractVisitor import AbstractVisitor

from duralex.alinea_parser import *
from duralex.AddCommitMessageVisitor import int_to_roman
import duralex.node_type
import duralex.diff

from bs4 import BeautifulSoup
import jinja2

import os
import subprocess
import tempfile
from distutils.dir_util import copy_tree

class CreateGitBookVisitor(AbstractVisitor):
    def __init__(self, args):
        self.gitbook_dir = args.gitbook
        self.tmp_dir = tempfile.mkdtemp()
        self.formats = args.gitbook_format

        super(CreateGitBookVisitor, self).__init__()

    def write_file(self, filename, data):
        f = open(self.tmp_dir + '/' + filename, 'w')
        f.write(data.encode('utf-8'))
        f.close()

    def get_article_commit_title(self, node):
        ancestors = get_node_ancestors(node)
        messages = []
        for ancestor in ancestors:
            if 'type' not in ancestor:
                continue;
            if ancestor['type'] == 'article':
                messages.append('Article ' + str(ancestor['order']))
            if ancestor['type'] == 'bill-header1' and 'implicit' not in ancestor:
                messages.append(int_to_roman(ancestor['order']))
            if ancestor['type'] == 'bill-header2':
                messages.append(unicode(ancestor['order']) + u'°')
        return ', '.join(messages[::-1])

    def get_article_commit_diff(self, edit, target_title, target_href):
        if 'htmlDiff' in edit:
            soup = BeautifulSoup(edit['htmlDiff'], "html5lib")
            filename_div = soup.find('div', {'class': 'diff-filename'})
            a_tag = soup.new_tag('a', href=target_href)
            a_tag.string = target_title
            filename_div.string = ''
            filename_div.append(a_tag)
            return unicode(soup.body.div)
        elif 'diff' in edit:
            process = subprocess.Popen(
                'diff2html -i stdin -d word -o stdout --su hidden -s line',
                shell=True,
                stdout=subprocess.PIPE,
                stdin=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            out, err = process.communicate(input=edit['diff'].encode('utf-8') + '\n')
            soup = BeautifulSoup(out, "html5lib")
            return (str(list(soup.find_all('style'))[0]) + '\n\n'
                + unicode(soup.find('div', {'id': 'diff'})))

    def get_articles(self, node):
        articles = []
        article_nodes = filter_nodes(node, lambda n: n['type'] == 'article')
        for article_node in article_nodes:
            edit_nodes = filter_nodes(article_node, lambda n: 'type' in n and n['type'] == 'edit')
            commits = []
            for edit_node in edit_nodes:
                article_ref = filter_nodes(edit_node, lambda n: n['type'] == 'article-reference')[0]
                target_title, target_href = self.get_deep_link(self.get_edit_target_nodes(article_ref))
                commits.append({
                    'title': self.get_article_commit_title(edit_node),
                    # remove the " ({reference list})" from the commit message since its already printed
                    # in the header above
                    'description': re.sub(r' \(.*\)', '', edit_node['commitMessage'].splitlines()[0]) if 'commitMessage' in edit_node else None,
                    'diff': self.get_article_commit_diff(edit_node, target_title, target_href),
                    'target': {
                        'title': target_title,
                        'link': target_href
                    }
                })

            articles.append({
                'order': article_node['order'],
                'content': article_node['content'].replace('\n', '\n\n'),
                'commits': commits,
                'githubIssue': article_node['githubIssue'] if 'githubIssue' in article_node else None,
                'gitlabIssue': article_node['gitlabIssue'] if 'gitlabIssue' in article_node else None
            })
        return articles

    def merge_dicts(self, *dict_args):
        """
        Given any number of dicts, shallow copy and merge into a new dict,
        precedence goes to key value pairs in latter dicts.
        """
        result = {}
        for dictionary in dict_args:
            result.update(dictionary)
        return result

    def visit_node(self, node):
        super(CreateGitBookVisitor, self).visit_node(node)

        if 'parent' not in node:
            edits = self.build_edit_matrix(node)
            articles = self.get_articles(node)
            modified_texts = self.get_modified_texts(edits)
            template_data = {
                'title': self.get_book_title(node),
                'url': node['url'],
                'type': node['type'],
                'modified': modified_texts,
                'articles': articles
            }

            if 'cocoricoVote' in node:
                template_data['cocorico_vote'] = node['cocoricoVote']

            self.template_file('book.json', template_data)
            self.template_file('styles/website.css', template_data)
            self.template_file('SUMMARY.md', template_data)
            self.template_file('README.md', template_data)
            current_article = 0
            for article in articles:
                self.template_file(
                    'article.md',
                    self.merge_dicts(template_data, {'current_article': current_article}),
                    'article-' + str(article['order']) + '.md'
                )
                current_article += 1

            current_article = 0
            current_law = 0
            for modified in modified_texts:
                self.template_file(
                    'law.md',
                    self.merge_dicts(template_data, {
                        'current_law': current_law,
                    }),
                    modified['law'] + '.md'
                )
                for article in modified['articles']:
                    self.template_file(
                        'text.md',
                        self.merge_dicts(template_data, {
                            'current_law': current_law,
                            'current_article': current_article
                        }),
                        modified['law'] + '-' + article['id'] + '.md'
                    )
                    current_article += 1
                current_law += 1

            if 'html' in self.formats:
                self.cmd('gitbook install')
                self.cmd('gitbook build')

                if 'markdown' in self.formats:
                    copy_tree(self.tmp_dir, self.gitbook_dir)
                else:
                    copy_tree(os.path.join(self.tmp_dir, '_book'), self.gitbook_dir)
            else:
                copy_tree(self.tmp_dir, self.gitbook_dir)

    def cmd(self, command):
        process = subprocess.Popen(
            command,
            cwd=self.tmp_dir,
            shell=True,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        return process.communicate()

    def get_book_title(self, root_node):
        title = root_node['type'].title()
        if 'id' in root_node:
            title += u' N°' + str(root_node['id'])
        if 'legislature' in root_node:
            title += ', ' + str(root_node['legislature']) + u'ème législature'
        return title

    def patch(self, original, unified_diff):
        fd, filename = input_file = tempfile.mkstemp()
        os.write(fd, original.encode('utf-8'))
        process = subprocess.Popen(
            'patch -p0 --output=- ' + filename,
            shell=True,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        out, err = process.communicate(input=unified_diff.encode('utf-8') + '\n')
        return ''.join(out).decode('utf-8')

    def get_deep_link(self, nodes):
        href = []
        title = []
        for node in nodes:
            if node['type'] == 'law-reference':
                title.append(u'Loi N°' + node['lawId'])
                href.append(node['lawId'])
            elif node['type'] == 'article':
                title.append(u'Article ' + str(node['order']))
                href.append(u'article-' + str(node['order']) + '.md#article-' + str(node['order']))
            elif node['type'] == 'article-reference':
                title.append(u'Article ' + node['id'])
                href.append(node['id'] + '.md')
            elif node['type'] == 'bill-header1' and 'implicit' not in node:
                title.append(int_to_roman(node['order']))
                href.append(int_to_roman(node['order']))
            elif node['type'] == 'bill-header2':
                title.append(unicode(node['order']) + u'°')
                href.append(str(node['order']) + u'°')
        return (', '.join(title), '-'.join(href))

    def get_edit_target_nodes(self, node):
        nodes = []

        if node_type.is_reference(node):
            nodes.append(node)

        nodes += filter(
            lambda n: node_type.is_reference(n),
            get_node_ancestors(node)
        )

        return sorted(
            nodes,
            key=lambda n: node_type.REFERENCE.index(n['type'])
        )

    def get_edit_source_nodes(self, node):
        edit_source_types = [
            'article',
            'bill-header1',
            'bill-header2'
        ]

        return sorted(
            filter(
                lambda n: 'type' in n and n['type'] in edit_source_types,
                get_node_ancestors(node)
            ),
            key=lambda n: edit_source_types.index(n['type'])
        )

    def get_modified_texts(self, edits):
        modified = []
        law_ids = set([i[0] for i in edits.keys()])
        for law_id in law_ids:
            law_edits = {k: v for k, v in edits.iteritems() if k[0] == law_id}
            articles = []
            for k, v in edits.iteritems():
                law_ref = filter_nodes(v[0][-1], lambda n: n['type'] == 'law-reference' and n['lawId'] == k[0])[0]
                article_ref = filter_nodes(law_ref, lambda n: n['type'] == 'article-reference' and n['id'] == k[1])[0]

                f = open(article_ref['filename'], 'r')
                text = f.read().decode('utf-8')
                original_text = text
                f.close()

                commits = []
                for edit_source in v:
                    title, href = self.get_deep_link(edit_source)
                    commits.append({'title': title, 'link': href})
                    edit_refs = filter_nodes(edit_source[-1], lambda n: n['type'] == 'edit')
                    for edit_ref in edit_refs:
                        text = self.patch(text, edit_ref['diff'])
                article = {
                    'id': k[1],
                    'diff': duralex.diff.make_html_rich_diff(original_text, text),
                    'commits': commits
                }
                if 'gitlabHistory' in article_ref:
                    article['gitlabHistory'] = article_ref['gitlabHistory']
                if 'githubHistory' in article_ref:
                    article['githubHistory'] = article_ref['githubHistory']
                articles.append(article)
            articles = sorted(articles, key=lambda x: x['id'].replace('-', ' '))
            modified.append({'law': law_id, 'articles': articles})
        return modified

    def build_edit_matrix(self, node):
        edits = {}
        article_refs = filter_nodes(node, lambda n: 'type' in n and n['type'] == 'article-reference')
        for article_ref in article_refs:
            law_ref = filter(
                lambda n: 'type' in n and n['type'] == 'law-reference',
                get_node_ancestors(article_ref)
            )[0]
            t = (law_ref['lawId'], article_ref['id'])
            if t not in edits:
                edits[t] = []
            edits[t].append(self.get_edit_source_nodes(article_ref))
        return edits

    def template_file(self, template, values, out=None):
        f = open(os.path.join('./template/gitbook', template), 'r')
        t = jinja2.Environment(loader=jinja2.FileSystemLoader('./template/gitbook')).from_string(f.read().decode('utf-8'))
        f.close()

        filename = os.path.join(self.tmp_dir, out if out else template)
        path = os.path.dirname(filename)
        if not os.path.exists(path):
            os.makedirs(path)
        f = open(filename, 'w')
        f.write(t.render(values).encode('utf-8'))
        f.truncate()
        f.close()
<|endoftext|>


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

from subprocess import check_output, PIPE
from re import findall

def traceroute(dest):
	"""
		Finds IP adresses along route to dest

		Shells out to existing 'traceroute' program on system.
		May take up to 5 seconds per traceroute (default timeout).
		IPs are gathered with quick and dirty regex,
			all other info (like name, ping) is discarded.
	"""

	raw_result = check_output(
		['traceroute', '-n', str(dest)],
		stdin = None,
		stderr = PIPE,  # stops subprocess from writting its stderr to our own
		shell = False
	)

	# extract IP adresses
	ipv4_pattern = r"(\d+\.\d+\.\d+\.\d+)"  # close enough for our purpose
	hops = findall(ipv4_pattern, raw_result)

	return hops

if __name__ == '__main__':
	import sys
	# run traceroute on all given arguments
	for dest in sys.argv[1:]:
		print "{}: {}".format(dest, traceroute(dest))<commit_msg>Fixed some pep8 style rrors in src/traceroute.py<commit_after>#!/usr/bin/env python

from subprocess import check_output, PIPE
from re import findall

def traceroute(dest):
    """
    Finds IP adresses along route to dest.
    Shells out to existing 'traceroute' program on system.
    May take up to 5 seconds per traceroute (default timeout).
    IPs are gathered with quick and dirty regex,
    all other info (like name, ping) is discarded.
    """

    raw_result = check_output(
        ['traceroute', '-n', str(dest)],
        stdin=None,
        stderr=PIPE,  # stops subprocess from writting its stderr to our own
        shell=False
    )

    # extract IP adresses
    ipv4_pattern = r"(\d+\.\d+\.\d+\.\d+)"  # close enough for our purpose
    hops = findall(ipv4_pattern, raw_result)

    return hops

if __name__ == '__main__':
    import sys
    # run traceroute on all given arguments
    for dest in sys.argv[1:]:
        print "{}: {}".format(dest, traceroute(dest))
<|endoftext|>


<commit_before>def parseCSVData(csvFile):
    """
    This function parses a typical csv file outputted from a simulation or
    sensitivity analysis in the form of

    Time (s)  Header1  Header2  Header....
    t0        val1_0   val2_0   val...
    t1..
    ..

    It returns the data in the form

    Time, DataList

    Where Time is returned as a GenericData object, and DataList is list of GenericData objects
    """
    import csv
    import numpy
    import re

    # Pattern for matching indices or units
    indexPattern = re.compile(r'^\S+\(\d+\)$')
    unitsPattern = re.compile(r'\s\(.+\)$')
    rxnSensPattern = re.compile('^dln\[\S+\]\/dln\[k\d+\]:\s\S+$')
    thermoSensPattern = re.compile('^dln\[\S+\]\/dG\[\S+\]$')

    timeData = []; data = {}
    f = csv.reader(open(csvFile, 'r'))

    columns = zip(*f)
    time = GenericData(label = columns[0][0],
                       data = numpy.array(columns[0][1:],dtype=numpy.float64),
                      )

    # Parse the units from the Time header
    if unitsPattern.search(time.label):
        label, sep, units = time.label[:-1].rpartition('(')
        time.label = label
        time.units = units

    dataList = []
    for col in columns[1:]:
        header = col[0]
        values = numpy.array(col[1:],dtype=numpy.float64)
        data = GenericData(label=header,data=values)

        # Parse the index or the label from the header
        if indexPattern.search(data.label):
            species, sep, index = data.label[:-1].rpartition('(')
            # Save the species attribute if an index was found
            data.species = species
            data.index = int(index)
        elif unitsPattern.search(data.label):
            label, sep, units = data.label[:-1].rpartition('(')
            data.label = label
            data.units = units
        elif rxnSensPattern.search(data.label):
            rxn = data.label.split()[1]
            index = data.label.split()[0][:-2].rpartition('dln[k')[2]
            data.reaction = rxn
            data.index = index
        elif thermoSensPattern.search(data.label):
            species = data.label[:-1].rpartition('dG[')[2]
            data.species = species
            if indexPattern.search(species):
                data.index = species[:-1].rpartition('(')[2]

        dataList.append(data)

    return time, dataList

def findNearest(array, value):
    """
    Returns the index of the closest value in a sorted array
    """
    import numpy
    idx = (numpy.abs(array-value)).argmin()
    return idx

class GenericData(object):
    """
    A generic data class for the purpose of plotting.
    label is the original label for the data
    data is the numpy array containing the data
    The other parameters are additional common flags for data found in RMG
    """
    def __init__(self, label='', data=None, species=None, reaction=None, units=None, index=None):
        self.label = label
        self.data = data
        self.species = species
        self.reaction = reaction
        self.units = units
        self.index = index

class GenericPlot(object):
    """
    A generic plotting class that can be extended to plot other things.
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel=''):
        self.xVar = xVar
        self.yVar = yVar
        self.title = title
        self.xlabel = xlabel
        self.ylabel = ylabel

    def plot(self, filename=''):
        """
        Execute the actual plotting
        """
        import matplotlib as mpl
        import matplotlib.pyplot as plt

        mpl.rc('font',family='monospace')
        fig=plt.figure()

        ax = fig.add_subplot(111)
        xVar = self.xVar
        yVar = self.yVar
        # Convert yVar to a list if it wasn't one already
        if isinstance(yVar, GenericData):
            yVar = [yVar]

        if len(yVar) == 1:
            y = yVar[0]
            ax.plot(xVar.data, y.data)

            # Create a ylabel for the label of the y variable
            if not self.ylabel and y.label:
                ylabel = y.label
                if y.units: ylabel += ' ({0})'.format(y.units)
                plt.ylabel(ylabel)
        else:
            for y in yVar:
                ax.plot(xVar.data, y.data, '-', label=y.label)

        if self.xlabel:
            plt.xlabel(self.xlabel)
        elif xVar.label:
            xlabel = xVar.label
            if xVar.units: xlabel += ' ({0})'.format(xVar.units)
            plt.xlabel(xlabel)

        if self.ylabel:
            plt.ylabel(self.ylabel)

        if self.title:
            plt.title(self.title)

        ax.grid('on')
        handles, labels = ax.get_legend_handles_labels()
        if labels:
            # Create a legend outside the plot and adjust width based off of longest legend label
            maxStringLength = max([len(label) for label in labels])
            width = 1.1 + .011*maxStringLength
            legend = ax.legend(handles,labels,loc='upper center', numpoints=1, bbox_to_anchor=(width,1)) #bbox_to_anchor=(1.01,.9)
            fig.savefig(filename, bbox_extra_artists=(legend,), bbox_inches='tight')
        else:
            fig.savefig(filename)

    def barplot(self, filename='', idx=None):
        """
        Plot a generic barplot using just the yVars.
        idx is the index of the each y-variable to be plotted. if not given, the last value will be used
        """
        import matplotlib as mpl
        import matplotlib.pyplot as plt
        import numpy

        mpl.rc('font',family='monospace')
        fig = plt.figure()
        ax = fig.add_subplot(111)

        position = numpy.arange(len(self.yVar),0,-1)+0.5
        # Reverse in order to go front top to bottom

        if not idx:
            idx = -1
        ax.barh(position, numpy.array([y.data[idx] for y in self.yVar]), align='center', alpha=0.5)
        plt.yticks(position, [y.label for y in self.yVar])

        # If any labels or titles are explicitly specified, write them
        if self.xlabel:
            plt.xlabel(self.xlabel)

        if self.ylabel:
            plt.ylabel(self.ylabel)

        if self.title:
            plt.title(self.title)

        fig.savefig(filename)

class SimulationPlot(GenericPlot):
    """
    A class for plotting simulations containing mole fraction vs time data.
    Can plot the top species in generic simulation csv generated by RMG-Py
    i.e. simulation_1_19.csv, found in the solver folder of an RMG job

    Use numSpecies as a flag to dictate how many species to plot.
    This function will plot the top species, based on maximum mole fraction at
    any point in the simulation.

    Alternatively, the `species` flag can be used as a dictionary for
    plotting specific species within the csvFile
    This should be formulated as
    {'desired_name_for_species': 'corresponding_chemkin_name_of_species'}
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='', csvFile='', numSpecies=None, species={}):
        GenericPlot.__init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='')
        self.csvFile = csvFile
        self.numSpecies = numSpecies
        self.species = species

    def load(self):
        time, dataList = parseCSVData(self.csvFile)
        speciesData = []
        if self.species:
            # A specific set of species was specified to be plotted
            for speciesLabel, chemkinLabel in self.species.iteritems():
                for data in dataList:
                    if chemkinLabel == data.label:
                        # replace the data label with the desired species label
                        data.label = speciesLabel
                        speciesData.append(data)
                        break
            speciesData.sort(key=lambda x: max(x.data), reverse=True)
            self.xVar = time
            self.yVar = speciesData

        else:
            for data in dataList:
                # Only plot if RMG detects that the data corresponds with a species
                # This will not include bath gases
                if data.species:
                    speciesData.append(data)
            speciesData.sort(key=lambda x: max(x.data), reverse=True)

            self.xVar = time
            self.yVar = speciesData[:self.numSpecies]

    def plot(self, filename=''):
        filename = filename if filename else 'simulation.png'
        self.load()
        self.yVar.sort(key=lambda x: max(x.data), reverse=True)

        if self.numSpecies:
            self.yVar = self.yVar[:self.numSpecies]

        GenericPlot.plot(self, filename=filename)

class ReactionSensitivityPlot(GenericPlot):
    """
    A class for plotting the top reaction sensitivites in a generic sensitivity csv file generated by RMG-Py.

    `numReactions` is a flag indicating the number of reaction sensitivities to plot.
    This function will plot the top sensitivities based on this number, based on the
    magnitude of the sensitivity at the final time step.

    `reactions` can be used to plot the sensitivities of a specific set of reactions.
    This should be formulated as
    {'desired_rxn_string': 'corresponding chemkin rxn string'}

    barplot() will instead plot a horizontal bar plot of the sensitivities at a given
    time step.  If time step is not given, the end step will automatically be chosen
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='', csvFile='', numReactions=None, reactions={}):
        GenericPlot.__init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='')
        self.csvFile = csvFile
        self.numReactions = numReactions
        self.reactions = reactions

    def load(self):
        time, dataList = parseCSVData(self.csvFile)
        reactionData = []
        if self.reactions:
            # A specific set of reaction sensitivities was specified to be plotted
            for reactionLabel, chemkinLabel in self.reactions.iteritems():
                for data in dataList:
                    if chemkinLabel == data.reaction:
                        # replace the data label with the desired species label
                        data.label = reactionLabel
                        reactionData.append(data)
                        break
        else:
            for data in dataList:
                if data.reaction:
                    reactionData.append(data)
        self.xVar = time
        self.yVar = reactionData

    def plot(self, filename=''):
        filename = filename if filename else "reaction_sensitivity.png"
        self.load()
        # Sort reactions according to absolute max value of final time point
        self.yVar.sort(key=lambda x: abs(x.data[-1]), reverse=True)
        if self.numReactions:
            self.yVar = self.yVar[:self.numReactions]
        self.ylabel = 'dlnc/dlnk'
        GenericPlot.plot(self, filename=filename)

    def barplot(self, filename='', t=None):
        """
        Time must be indicated in seconds
        The closest timepoint will then be used, otherwise if no time point is given,
        the end time step will be used
        """
        filename = filename if filename else "reaction_sensitivity.png"
        self.load()
        # Sort reactions according to absolute max value at the specified time point
        # if the time point is not given, use the final time point
        if t:
            idx = findNearest(self.xVar.data, t)
        else:
            idx = -1

        self.yVar.sort(key=lambda x: abs(x.data[idx]), reverse=True)
        if self.numReactions:
            self.yVar = self.yVar[:self.numReactions]

        self.xlabel = 'dlnc/dlnk_i'
        GenericPlot.barplot(self, filename=filename, idx=idx)

class ThermoSensitivityPlot(GenericPlot):
    """
    A class for plotting the top sensitivities to a thermo DeltaG value of species within the model.
    The value used is the sensitivity at the final time point.
    `numSpecies` indicates the number of species to plot.

    `species` is a dictionary corresponding to specific species thermo sensitivities to be plotted

    barplot() will instead plot a horizontal bar plot of the sensitivities at a given
    time step.  If time step is not given, the end step will automatically be chosen
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='', csvFile='', numSpecies=None, species={}):
        GenericPlot.__init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='')
        self.csvFile = csvFile
        self.numSpecies = numSpecies
        self.species = species

    def load(self):
        time, dataList = parseCSVData(self.csvFile)
        thermoData = []
        if self.species:
            # A specific set of species sensitivities was specified to be plotted
            for speciesLabel, chemkinLabel in self.species.iteritems():
                for data in dataList:
                    if chemkinLabel == data.species:
                        # replace the data label with the desired species label
                        data.label = speciesLabel
                        thermoData.append(data)
                        break
        else:
            for data in dataList:
                if data.species:
                    thermoData.append(data)

        self.xVar = time
        self.yVar = thermoData

    def plot(self, filename=''):
        filename = filename if filename else "thermo_sensitivity.png"
        self.load()
        self.yVar.sort(key=lambda x: abs(x.data[-1]), reverse = True)
        if self.numSpecies:
            self.yVar = self.yVar[:self.numSpecies]
        self.ylabel = 'dlnc/dlnG'
        GenericPlot.plot(self, filename=filename)

    def barplot(self, filename='', t=None):
        filename = filename if filename else "thermo_sensitivity.png"
        self.load()
        if t:
            idx = findNearest(self.xVar.data, t)
        else:
            idx = -1
        self.yVar.sort(key=lambda x: abs(x.data[idx]), reverse = True)
        if self.numSpecies:
            self.yVar = self.yVar[:self.numSpecies]
        self.xlabel = 'dlnc/dlnG'
        GenericPlot.barplot(self, filename=filename, idx=idx)<commit_msg>Load from csv only when the data isn't already in the object.<commit_after>def parseCSVData(csvFile):
    """
    This function parses a typical csv file outputted from a simulation or
    sensitivity analysis in the form of

    Time (s)  Header1  Header2  Header....
    t0        val1_0   val2_0   val...
    t1..
    ..

    It returns the data in the form

    Time, DataList

    Where Time is returned as a GenericData object, and DataList is list of GenericData objects
    """
    import csv
    import numpy
    import re

    # Pattern for matching indices or units
    indexPattern = re.compile(r'^\S+\(\d+\)$')
    unitsPattern = re.compile(r'\s\(.+\)$')
    rxnSensPattern = re.compile('^dln\[\S+\]\/dln\[k\d+\]:\s\S+$')
    thermoSensPattern = re.compile('^dln\[\S+\]\/dG\[\S+\]$')

    timeData = []; data = {}
    f = csv.reader(open(csvFile, 'r'))

    columns = zip(*f)
    time = GenericData(label = columns[0][0],
                       data = numpy.array(columns[0][1:],dtype=numpy.float64),
                      )

    # Parse the units from the Time header
    if unitsPattern.search(time.label):
        label, sep, units = time.label[:-1].rpartition('(')
        time.label = label
        time.units = units

    dataList = []
    for col in columns[1:]:
        header = col[0]
        values = numpy.array(col[1:],dtype=numpy.float64)
        data = GenericData(label=header,data=values)

        # Parse the index or the label from the header
        if indexPattern.search(data.label):
            species, sep, index = data.label[:-1].rpartition('(')
            # Save the species attribute if an index was found
            data.species = species
            data.index = int(index)
        elif unitsPattern.search(data.label):
            label, sep, units = data.label[:-1].rpartition('(')
            data.label = label
            data.units = units
        elif rxnSensPattern.search(data.label):
            rxn = data.label.split()[1]
            index = data.label.split()[0][:-2].rpartition('dln[k')[2]
            data.reaction = rxn
            data.index = index
        elif thermoSensPattern.search(data.label):
            species = data.label[:-1].rpartition('dG[')[2]
            data.species = species
            if indexPattern.search(species):
                data.index = species[:-1].rpartition('(')[2]

        dataList.append(data)

    return time, dataList

def findNearest(array, value):
    """
    Returns the index of the closest value in a sorted array
    """
    import numpy
    idx = (numpy.abs(array-value)).argmin()
    return idx

class GenericData(object):
    """
    A generic data class for the purpose of plotting.
    label is the original label for the data
    data is the numpy array containing the data
    The other parameters are additional common flags for data found in RMG
    """
    def __init__(self, label='', data=None, species=None, reaction=None, units=None, index=None):
        self.label = label
        self.data = data
        self.species = species
        self.reaction = reaction
        self.units = units
        self.index = index

class GenericPlot(object):
    """
    A generic plotting class that can be extended to plot other things.
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel=''):
        self.xVar = xVar
        self.yVar = yVar
        self.title = title
        self.xlabel = xlabel
        self.ylabel = ylabel

    def plot(self, filename=''):
        """
        Execute the actual plotting
        """
        import matplotlib as mpl
        import matplotlib.pyplot as plt

        mpl.rc('font',family='monospace')
        fig=plt.figure()

        ax = fig.add_subplot(111)
        xVar = self.xVar
        yVar = self.yVar
        # Convert yVar to a list if it wasn't one already
        if isinstance(yVar, GenericData):
            yVar = [yVar]

        if len(yVar) == 1:
            y = yVar[0]
            ax.plot(xVar.data, y.data)

            # Create a ylabel for the label of the y variable
            if not self.ylabel and y.label:
                ylabel = y.label
                if y.units: ylabel += ' ({0})'.format(y.units)
                plt.ylabel(ylabel)
        else:
            for y in yVar:
                ax.plot(xVar.data, y.data, '-', label=y.label)

        if self.xlabel:
            plt.xlabel(self.xlabel)
        elif xVar.label:
            xlabel = xVar.label
            if xVar.units: xlabel += ' ({0})'.format(xVar.units)
            plt.xlabel(xlabel)

        if self.ylabel:
            plt.ylabel(self.ylabel)

        if self.title:
            plt.title(self.title)

        ax.grid('on')
        handles, labels = ax.get_legend_handles_labels()
        if labels:
            # Create a legend outside the plot and adjust width based off of longest legend label
            maxStringLength = max([len(label) for label in labels])
            width = 1.1 + .011*maxStringLength
            legend = ax.legend(handles,labels,loc='upper center', numpoints=1, bbox_to_anchor=(width,1)) #bbox_to_anchor=(1.01,.9)
            fig.savefig(filename, bbox_extra_artists=(legend,), bbox_inches='tight')
        else:
            fig.savefig(filename)

    def barplot(self, filename='', idx=None):
        """
        Plot a generic barplot using just the yVars.
        idx is the index of the each y-variable to be plotted. if not given, the last value will be used
        """
        import matplotlib as mpl
        import matplotlib.pyplot as plt
        import numpy

        mpl.rc('font',family='monospace')
        fig = plt.figure()
        ax = fig.add_subplot(111)

        position = numpy.arange(len(self.yVar),0,-1)+0.5
        # Reverse in order to go front top to bottom

        if not idx:
            idx = -1
        ax.barh(position, numpy.array([y.data[idx] for y in self.yVar]), align='center', alpha=0.5)
        plt.yticks(position, [y.label for y in self.yVar])

        # If any labels or titles are explicitly specified, write them
        if self.xlabel:
            plt.xlabel(self.xlabel)

        if self.ylabel:
            plt.ylabel(self.ylabel)

        if self.title:
            plt.title(self.title)

        fig.savefig(filename)

class SimulationPlot(GenericPlot):
    """
    A class for plotting simulations containing mole fraction vs time data.
    Can plot the top species in generic simulation csv generated by RMG-Py
    i.e. simulation_1_19.csv, found in the solver folder of an RMG job

    Use numSpecies as a flag to dictate how many species to plot.
    This function will plot the top species, based on maximum mole fraction at
    any point in the simulation.

    Alternatively, the `species` flag can be used as a dictionary for
    plotting specific species within the csvFile
    This should be formulated as
    {'desired_name_for_species': 'corresponding_chemkin_name_of_species'}
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='', csvFile='', numSpecies=None, species={}):
        GenericPlot.__init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='')
        self.csvFile = csvFile
        self.numSpecies = numSpecies
        self.species = species

    def load(self):
        if self.xVar == None and self.yVar == None:
            time, dataList = parseCSVData(self.csvFile)

        speciesData = []
        if self.species:
            # A specific set of species was specified to be plotted
            for speciesLabel, chemkinLabel in self.species.iteritems():
                for data in dataList:
                    if chemkinLabel == data.label:
                        # replace the data label with the desired species label
                        data.label = speciesLabel
                        speciesData.append(data)
                        break
            speciesData.sort(key=lambda x: max(x.data), reverse=True)
            self.xVar = time
            self.yVar = speciesData

        else:
            for data in dataList:
                # Only plot if RMG detects that the data corresponds with a species
                # This will not include bath gases
                if data.species:
                    speciesData.append(data)
            speciesData.sort(key=lambda x: max(x.data), reverse=True)

            self.xVar = time
            self.yVar = speciesData[:self.numSpecies]

    def plot(self, filename=''):
        filename = filename if filename else 'simulation.png'
        self.load()
        self.yVar.sort(key=lambda x: max(x.data), reverse=True)

        if self.numSpecies:
            self.yVar = self.yVar[:self.numSpecies]

        GenericPlot.plot(self, filename=filename)

class ReactionSensitivityPlot(GenericPlot):
    """
    A class for plotting the top reaction sensitivites in a generic sensitivity csv file generated by RMG-Py.

    `numReactions` is a flag indicating the number of reaction sensitivities to plot.
    This function will plot the top sensitivities based on this number, based on the
    magnitude of the sensitivity at the final time step.

    `reactions` can be used to plot the sensitivities of a specific set of reactions.
    This should be formulated as
    {'desired_rxn_string': 'corresponding chemkin rxn string'}

    barplot() will instead plot a horizontal bar plot of the sensitivities at a given
    time step.  If time step is not given, the end step will automatically be chosen
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='', csvFile='', numReactions=None, reactions={}):
        GenericPlot.__init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='')
        self.csvFile = csvFile
        self.numReactions = numReactions
        self.reactions = reactions

    def load(self):
        if self.xVar == None and self.yVar == None:
            time, dataList = parseCSVData(self.csvFile)
        reactionData = []
        if self.reactions:
            # A specific set of reaction sensitivities was specified to be plotted
            for reactionLabel, chemkinLabel in self.reactions.iteritems():
                for data in dataList:
                    if chemkinLabel == data.reaction:
                        # replace the data label with the desired species label
                        data.label = reactionLabel
                        reactionData.append(data)
                        break
        else:
            for data in dataList:
                if data.reaction:
                    reactionData.append(data)
        self.xVar = time
        self.yVar = reactionData

    def plot(self, filename=''):
        filename = filename if filename else "reaction_sensitivity.png"
        self.load()
        # Sort reactions according to absolute max value of final time point
        self.yVar.sort(key=lambda x: abs(x.data[-1]), reverse=True)
        if self.numReactions:
            self.yVar = self.yVar[:self.numReactions]
        self.ylabel = 'dlnc/dlnk'
        GenericPlot.plot(self, filename=filename)

    def barplot(self, filename='', t=None):
        """
        Time must be indicated in seconds
        The closest timepoint will then be used, otherwise if no time point is given,
        the end time step will be used
        """
        filename = filename if filename else "reaction_sensitivity.png"
        self.load()
        # Sort reactions according to absolute max value at the specified time point
        # if the time point is not given, use the final time point
        if t:
            idx = findNearest(self.xVar.data, t)
        else:
            idx = -1

        self.yVar.sort(key=lambda x: abs(x.data[idx]), reverse=True)
        if self.numReactions:
            self.yVar = self.yVar[:self.numReactions]

        self.xlabel = 'dlnc/dlnk_i'
        GenericPlot.barplot(self, filename=filename, idx=idx)

class ThermoSensitivityPlot(GenericPlot):
    """
    A class for plotting the top sensitivities to a thermo DeltaG value of species within the model.
    The value used is the sensitivity at the final time point.
    `numSpecies` indicates the number of species to plot.

    `species` is a dictionary corresponding to specific species thermo sensitivities to be plotted

    barplot() will instead plot a horizontal bar plot of the sensitivities at a given
    time step.  If time step is not given, the end step will automatically be chosen
    """
    def __init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='', csvFile='', numSpecies=None, species={}):
        GenericPlot.__init__(self, xVar=None, yVar=None, title='', xlabel='', ylabel='')
        self.csvFile = csvFile
        self.numSpecies = numSpecies
        self.species = species

    def load(self):
        if self.xVar == None and self.yVar == None:
            time, dataList = parseCSVData(self.csvFile)
        thermoData = []
        if self.species:
            # A specific set of species sensitivities was specified to be plotted
            for speciesLabel, chemkinLabel in self.species.iteritems():
                for data in dataList:
                    if chemkinLabel == data.species:
                        # replace the data label with the desired species label
                        data.label = speciesLabel
                        thermoData.append(data)
                        break
        else:
            for data in dataList:
                if data.species:
                    thermoData.append(data)

        self.xVar = time
        self.yVar = thermoData

    def plot(self, filename=''):
        filename = filename if filename else "thermo_sensitivity.png"
        self.load()
        self.yVar.sort(key=lambda x: abs(x.data[-1]), reverse = True)
        if self.numSpecies:
            self.yVar = self.yVar[:self.numSpecies]
        self.ylabel = 'dlnc/dlnG'
        GenericPlot.plot(self, filename=filename)

    def barplot(self, filename='', t=None):
        filename = filename if filename else "thermo_sensitivity.png"
        self.load()
        if t:
            idx = findNearest(self.xVar.data, t)
        else:
            idx = -1
        self.yVar.sort(key=lambda x: abs(x.data[idx]), reverse = True)
        if self.numSpecies:
            self.yVar = self.yVar[:self.numSpecies]
        self.xlabel = 'dlnc/dlnG'
        GenericPlot.barplot(self, filename=filename, idx=idx)<|endoftext|>