<commit_before>import asyncio
import logging
import os
import sys

import aiohttp.web
import aiojobs
import aiojobs.aiohttp

import virtool.config
import virtool.db.core
import virtool.db.migrate
import virtool.db.utils
import virtool.dispatcher
import virtool.errors
import virtool.hmm.db
import virtool.http.accept
import virtool.http.auth
import virtool.http.csp
import virtool.http.errors
import virtool.http.proxy
import virtool.http.query
import virtool.jobs.runner
import virtool.logs
import virtool.references.db
import virtool.resources
import virtool.routes
import virtool.sentry
import virtool.settings.db
import virtool.settings.schema
import virtool.software.db
import virtool.startup
import virtool.utils
import virtool.version

logger = logging.getLogger(__name__)

async def on_shutdown(app: aiohttp.web.Application):
    """
    A function called when the app is shutting down.

    :param app: the app object
    :type app: :class:`aiohttp.web.Application`

    """
    logger.debug("Shutting down")

    try:
        await app["client"].close()
    except KeyError:
        pass

    try:
        await app["dispatcher"].close()
    except KeyError:
        pass

    try:
        app["executor"].shutdown(wait=True)
    except KeyError:
        pass

    try:
        app["process_executor"].shutdown(wait=True)
    except KeyError:
        pass

    scheduler = virtool.startup.get_scheduler_from_app(app)
    await scheduler.close()

    try:
        app["redis"].close()
        await app["redis"].wait_closed()
    except KeyError:
        pass

def create_app(config):
    """
    Creates the Virtool application.

    """
    middlewares = [
        virtool.http.auth.middleware,
        virtool.http.accept.middleware,
        virtool.http.csp.middleware,
        virtool.http.errors.middleware,
        virtool.http.proxy.middleware,
        virtool.http.query.middleware
    ]

    app = aiohttp.web.Application(middlewares=middlewares)

    app["config"] = config
    app["change_queue"] = asyncio.Queue()
    app["mode"] = "server"

    aiojobs.aiohttp.setup(app)

    app.on_startup.extend([
        virtool.startup.init_version,
        virtool.startup.init_events,
        virtool.startup.init_db,
        virtool.startup.init_dispatcher,
        virtool.startup.init_settings,
        virtool.startup.init_client_path,
        virtool.startup.init_http_client,
        virtool.startup.init_paths,
        virtool.startup.init_postgres,
        virtool.startup.init_routes,
        virtool.startup.init_executors,
        virtool.startup.init_tasks,
        virtool.startup.init_redis,
        virtool.startup.init_listen_for_changes,
        virtool.startup.init_sentry,
        virtool.startup.init_check_db,
        virtool.startup.init_resources,
        virtool.startup.init_job_interface,
        virtool.startup.init_refresh
    ])

    app.on_response_prepare.append(virtool.http.csp.on_prepare)
    app.on_shutdown.append(on_shutdown)

    return app

async def create_app_runner(app: aiohttp.web.Application, host: str, port: int) -> aiohttp.web.AppRunner:
    """
    Create an :class:`aiohttp.web.AppRunner` to allow customization of signal handlers.

    The basic :func:`aiohttp.web.run_app` sets up handlers for `SIGINT` and `SIGTERM` that can interfere with Virtool
    code such as that for restarting the server after software update. This custom runner allows handling of signals
    as well as restart and shutdown events from users.

    https://docs.aiohttp.org/en/stable/web_advanced.html#application-runners

    :param app: the application
    :param host: the host to listen on
    :param port: the port to listen on
    :return: a custom :class:`~aiohttp.web.AppRunner`

    """
    runner = aiohttp.web.AppRunner(app)

    await runner.setup()

    site = aiohttp.web.TCPSite(runner, host, port)

    try:
        await site.start()
    except OSError as err:
        if err.args[0] == 48:
            logger.fatal(f"Could not bind address {(host, port)}")
            sys.exit(1)

    logger.info(f"Listening at http://{host}:{port}")

    return runner

async def run_app(config):
    app = create_app(config)

    runner = await create_app_runner(
        app,
        config["host"],
        config["port"]
    )

    _, pending = await asyncio.wait(
        [wait_for_restart(runner, app["events"]), wait_for_shutdown(runner, app["events"])],
        return_when=asyncio.FIRST_COMPLETED
    )

    for task in pending:
        task.cancel()

async def wait_for_restart(runner: aiohttp.web.AppRunner, events: dict):
    """
    Wait for the shutdown event and restart if it is encountered.

    Restart is accomplished using :func:`os.execl` or :func:`os.execv` after cleaning up the `runner`.

    :param runner: the :class:`~aiohttp.web.AppRunner` returned by :func:`.create_app_runner`
    :param events: a dict containing the `restart` and `shutdown` :class:`~asyncio.Event` objects

    """
    await events["restart"].wait()
    await asyncio.sleep(0.5)
    await runner.cleanup()

    exe = sys.executable

    if exe.endswith("python") or "python3" in exe:
        os.execl(exe, exe, *sys.argv)
        return

    if exe.endswith("run"):
        os.execv(exe, sys.argv)
        return

    raise SystemError("Could not determine executable type")

async def wait_for_shutdown(runner: aiohttp.web.AppRunner, events: dict):
    """
    Wait for the shutdown event and terminate if it is encountered.

    :param runner: the :class:`~aiohttp.web.AppRunner` returned by :func:`.create_app_runner`
    :param events: a dict containing the `restart` and `shutdown` :class:`~asyncio.Event` objects

    """
    await events["shutdown"].wait()
    await asyncio.sleep(0.5)
    await runner.cleanup()
<commit_msg>Use nonce generation middleware first<commit_after>import asyncio
import logging
import os
import sys

import aiohttp.web
import aiojobs
import aiojobs.aiohttp

import virtool.config
import virtool.db.core
import virtool.db.migrate
import virtool.db.utils
import virtool.dispatcher
import virtool.errors
import virtool.hmm.db
import virtool.http.accept
import virtool.http.auth
import virtool.http.csp
import virtool.http.errors
import virtool.http.proxy
import virtool.http.query
import virtool.jobs.runner
import virtool.logs
import virtool.references.db
import virtool.resources
import virtool.routes
import virtool.sentry
import virtool.settings.db
import virtool.settings.schema
import virtool.software.db
import virtool.startup
import virtool.utils
import virtool.version

logger = logging.getLogger(__name__)

async def on_shutdown(app: aiohttp.web.Application):
    """
    A function called when the app is shutting down.

    :param app: the app object
    :type app: :class:`aiohttp.web.Application`

    """
    logger.debug("Shutting down")

    try:
        await app["client"].close()
    except KeyError:
        pass

    try:
        await app["dispatcher"].close()
    except KeyError:
        pass

    try:
        app["executor"].shutdown(wait=True)
    except KeyError:
        pass

    try:
        app["process_executor"].shutdown(wait=True)
    except KeyError:
        pass

    scheduler = virtool.startup.get_scheduler_from_app(app)
    await scheduler.close()

    try:
        app["redis"].close()
        await app["redis"].wait_closed()
    except KeyError:
        pass

def create_app(config):
    """
    Creates the Virtool application.

    """
    middlewares = [
        virtool.http.csp.middleware,
        virtool.http.auth.middleware,
        virtool.http.accept.middleware,
        virtool.http.errors.middleware,
        virtool.http.proxy.middleware,
        virtool.http.query.middleware
    ]

    app = aiohttp.web.Application(middlewares=middlewares)

    app["config"] = config
    app["change_queue"] = asyncio.Queue()
    app["mode"] = "server"

    aiojobs.aiohttp.setup(app)

    app.on_startup.extend([
        virtool.startup.init_version,
        virtool.startup.init_events,
        virtool.startup.init_db,
        virtool.startup.init_dispatcher,
        virtool.startup.init_settings,
        virtool.startup.init_client_path,
        virtool.startup.init_http_client,
        virtool.startup.init_paths,
        virtool.startup.init_postgres,
        virtool.startup.init_routes,
        virtool.startup.init_executors,
        virtool.startup.init_tasks,
        virtool.startup.init_redis,
        virtool.startup.init_listen_for_changes,
        virtool.startup.init_sentry,
        virtool.startup.init_check_db,
        virtool.startup.init_resources,
        virtool.startup.init_job_interface,
        virtool.startup.init_refresh
    ])

    app.on_response_prepare.append(virtool.http.csp.on_prepare)
    app.on_shutdown.append(on_shutdown)

    return app

async def create_app_runner(app: aiohttp.web.Application, host: str, port: int) -> aiohttp.web.AppRunner:
    """
    Create an :class:`aiohttp.web.AppRunner` to allow customization of signal handlers.

    The basic :func:`aiohttp.web.run_app` sets up handlers for `SIGINT` and `SIGTERM` that can interfere with Virtool
    code such as that for restarting the server after software update. This custom runner allows handling of signals
    as well as restart and shutdown events from users.

    https://docs.aiohttp.org/en/stable/web_advanced.html#application-runners

    :param app: the application
    :param host: the host to listen on
    :param port: the port to listen on
    :return: a custom :class:`~aiohttp.web.AppRunner`

    """
    runner = aiohttp.web.AppRunner(app)

    await runner.setup()

    site = aiohttp.web.TCPSite(runner, host, port)

    try:
        await site.start()
    except OSError as err:
        if err.args[0] == 48:
            logger.fatal(f"Could not bind address {(host, port)}")
            sys.exit(1)

    logger.info(f"Listening at http://{host}:{port}")

    return runner

async def run_app(config):
    app = create_app(config)

    runner = await create_app_runner(
        app,
        config["host"],
        config["port"]
    )

    _, pending = await asyncio.wait(
        [wait_for_restart(runner, app["events"]), wait_for_shutdown(runner, app["events"])],
        return_when=asyncio.FIRST_COMPLETED
    )

    for task in pending:
        task.cancel()

async def wait_for_restart(runner: aiohttp.web.AppRunner, events: dict):
    """
    Wait for the shutdown event and restart if it is encountered.

    Restart is accomplished using :func:`os.execl` or :func:`os.execv` after cleaning up the `runner`.

    :param runner: the :class:`~aiohttp.web.AppRunner` returned by :func:`.create_app_runner`
    :param events: a dict containing the `restart` and `shutdown` :class:`~asyncio.Event` objects

    """
    await events["restart"].wait()
    await asyncio.sleep(0.5)
    await runner.cleanup()

    exe = sys.executable

    if exe.endswith("python") or "python3" in exe:
        os.execl(exe, exe, *sys.argv)
        return

    if exe.endswith("run"):
        os.execv(exe, sys.argv)
        return

    raise SystemError("Could not determine executable type")

async def wait_for_shutdown(runner: aiohttp.web.AppRunner, events: dict):
    """
    Wait for the shutdown event and terminate if it is encountered.

    :param runner: the :class:`~aiohttp.web.AppRunner` returned by :func:`.create_app_runner`
    :param events: a dict containing the `restart` and `shutdown` :class:`~asyncio.Event` objects

    """
    await events["shutdown"].wait()
    await asyncio.sleep(0.5)
    await runner.cleanup()
<|endoftext|>


<commit_before>import apscheduler.schedulers.background
import datetime
import distutils.spawn
import importlib
import json
import mysql.connector
import queue
import time
import sys
import pprint
from utilities import url_to_json

tokens = []

class Tokenizer():
    def __init__(self):
        self.have_tokens = bool(tokens)
        self.available_tokens = queue.Queue()
        self.scheduler = apscheduler.schedulers.background.BackgroundScheduler(
            )
        self.scheduler.start()

        if not self.have_tokens:
            self.print_warning(
                'No GitHub OAuth tokens provided. Proceeding without '
                'authentication.'
            )

        for token in tokens:
            self.available_tokens.put(token)

    def tokenize(self, url):
        if url.startswith('https://api.github.com'):
            if self.have_tokens:
                token = self.get_token()
                if token is not None:
                    return '{0}?access_token={1}'.format(url, token)
                else:
                    return url
            else:
                return url
        else:
            raise ValueError('url must be for the GitHub API')

    def get_token(self):
        while True:
            if not self.scheduler.get_jobs() and self.available_tokens.empty():
                self.print_warning('No more valid OAuth tokens available.')
                return None

            token = self.available_tokens.get(block=True)

            rate_limit_url = (
                'https://api.github.com/rate_limit?access_token={0}'
            ).format(token)
            status = url_to_json(rate_limit_url)

            # Throw away bad OAuth keys.
            if 'resources' not in status:
                self.print_warning(
                    'Invalid OAuth token supplied. Trying again...'
                )
                continue

            if status['resources']['core']['remaining'] > 0:
                self.available_tokens.put_nowait(token)
                return token
            else:
                self.scheduler.add_job(
                    self.available_tokens.put_nowait,
                    'date',
                    args=[token],
                    run_date=datetime.datetime.fromtimestamp(
                        status['resources']['core']['reset']
                    )
                )

    def print_warning(self, message):
        formatted_message = '\033[91mWARNING\033[0m: {0}'.format(message)
        print(formatted_message)

def get_run_id(cursor):
    '''
    Generate an id from the datasource that uniquely identifies the run.

    Args:
        cursor: mysql.cursor.MySQLCursor
            Cursor object used to insert data.

    Return:
        Identifier of the run to persist.
    '''
    cursor.execute('INSERT INTO reaper_runs VALUES()')

    return cursor.lastrowid

def save_result(run_id, repo_id, results, cursor):
    """
    Save the results to the specified data source, creating the results table
    as necessary.

    Args:
        run_id: int
            Identifier of the run to persist.
        repo_id: int
            Identifier of the repository to save.
        results: dict
            Key value pair of results to be saved.
        cursor: mysql.cursor.MySQLCursor
            Cursor object used to insert data.

    Return:
        True if successful, False otherwise.
    """
    query = '''
        INSERT INTO reaper_results
            (project_id, results, run_id)
        VALUES
            ({0}, '{1}', {2})
        '''.format(repo_id, json.dumps(results), run_id)

    cursor.execute(query)

def init_attribute_plugins(attributes, connection):
    for attribute in attributes:
        if 'implementation' in attribute:
            try:
                cursor = connection.cursor()
                attribute['implementation'].init(cursor)
            except:
                pass
            finally:
                cursor.close()

def global_init_attribute_plugins(attributes, connection, sample):
    for attribute in attributes:
        if 'implementation' in attribute:
            try:
                cursor = connection.cursor()
                attribute['implementation'].global_init(cursor, sample)
            except:
                pass
            finally:
                cursor.close()

def process_repository(project_id, repo_path, attributes, connection):
    """
    Calculate a score for the given repository.

    Args:
        project_id: int
            GHTorrent dataset identifier for the repository.
        repo_path: string
            Path to the repository contents.
        attributes: list
            List of attributes to be executed against the repository.
        connection: mysql.connector.MySQLConnection
            Database connection used for querying the GHTorrent dataset.
    """
    score = 0
    results = {}
    for attribute in attributes:
        if 'implementation' in attribute:
            cursor = connection.cursor()

            binary_result, raw_result = attribute['implementation'].run(
                project_id,
                repo_path,
                cursor,
                **attribute.get('options', {})
            )

            cursor.close()

            score += binary_result * attribute['weight']
            results[attribute['name']] = raw_result

            if ('essential' in attribute and
                    attribute['essential'] and
                    not binary_result):
                score = 0
                break

    return score, results

def load_attribute_plugins(plugins_dir, attributes):
    """
    Attempt to load each of the attributes as defined in the configuration
    file.

    Args:
        attributes: list
            List of attribute dictionaries with the specific configuration
            data.
    """
    for attribute in attributes:
        if attribute['enabled']:
            try:
                attribute['implementation'] = importlib.import_module(
                    '{0}.{1}.main'.format(plugins_dir, attribute['name'])
                )

            except ImportError:
                print('Failed to load the {0} attribute.'.format(
                        attribute['name']
                    )
                )

def establish_database_connection(config):
    """
    Attempt to establish a connection to the specified database. Exit with an
    error message on failure.

    Args:
        config: dict
            Settings for the database connection.
    """
    try:
        config['autocommit'] = True
        connection = mysql.connector.connect(**config)
        connection.connect()
        return connection
    except mysql.connector.Error:
        print('\rUnable to establish connection to database.')
        sys.exit(1)

def process_configuration(config_file):
    """
    Load and validate the given configuration file.

    Args:
        config_file: File
            File object with the configuration contents.

    Returns:
        Validated dictionary with configuration parameters.
    """
    global tokens

    try:
        config = json.load(config_file)
    except:
        print('Your config file is malformed, or does not exist!')
        sys.exit(1)
    finally:
        tokens = config['options'].get('github_tokens', [])
        return config

def process_plugins(plugins_dir):
    manifest = plugins_dir + '/manifest.json'
    try:
        with open(manifest, 'r') as f:
            plugins = json.load(f)

            attributes = plugins['attributes']
            threshold = plugins['threshold']

            for attribute in attributes:
                exit_if_dep_missing(attribute)

            return (threshold, attributes)

    except:
        print('{} is malformed, or does not exist!'.format(manifest))
        sys.exit(1)

def exit_if_dep_missing(attribute):
    if not attribute['enabled'] or 'dependencies' not in attribute:
        dependencies = []
    else:
        dependencies = attribute['dependencies']

    for dependency in dependencies:
        if not distutils.spawn.find_executable(dependency):
            print('Missing dependency for attribute {}: {}'
                .format(attribute['name'], dependency)
            )
            sys.exit(1)

def process_key_string(attributes, key_string):
    """Reprocesses the attribute configuration based on a keystring

    The keystring is a single string interpreted as a list of individual
    initials. Each initial indicates a particular attribute, as defined in
    config.json.

    If no keystring is provided, the attribute configuration is left untouched.
    If a keystring is provided, then first, all attributes are disabled. Only
    those attribtues referred by their initials in the keystring are
    re-enabled. Further, if the attribute initial in the keystring is
    capitalized, then that attribute's data will be persisted.
    """
    if key_string is None:
        return

    initials = dict()  # Attribute initial => attribute dict
    for attribute in attributes:
        initial = attribute['initial'].lower()
        initials[initial] = attribute
        attribute['enabled'] = False  # Disable all attributes first
        attribute['persistResult'] = False

    if key_string == "help":
        for key_low, attribute in initials.items():
            print(key_low + " " + attribute['name'])
        sys.exit()

    for key_raw in key_string:
        key_low = key_raw.lower()
        attribute = initials[key_low]
        # Enable attribute since its initial is in the keystring
        attribute['enabled'] = True
        # Persist results if the raw key is capitalized
        attribute['persistResult'] = (not key_raw == key_low)

def get_persist_attrs(attributes):
    persist_attrs = []
    for attribute in attributes:
        if attribute['persistResult']:
            persist_attrs.append(attribute['name'])

    return (len(persist_attrs) > 0, persist_attrs)
<commit_msg>Fix issue related to long running processes<commit_after>import apscheduler.schedulers.background
import datetime
import distutils.spawn
import importlib
import json
import mysql.connector
import queue
import time
import sys
import pprint
from utilities import url_to_json

tokens = []

class Tokenizer():
    def __init__(self):
        self.have_tokens = bool(tokens)
        self.available_tokens = queue.Queue()
        self.scheduler = apscheduler.schedulers.background.BackgroundScheduler(
            )
        self.scheduler.start()

        if not self.have_tokens:
            self.print_warning(
                'No GitHub OAuth tokens provided. Proceeding without '
                'authentication.'
            )

        for token in tokens:
            self.available_tokens.put(token)

    def tokenize(self, url):
        if url.startswith('https://api.github.com'):
            if self.have_tokens:
                token = self.get_token()
                if token is not None:
                    return '{0}?access_token={1}'.format(url, token)
                else:
                    return url
            else:
                return url
        else:
            raise ValueError('url must be for the GitHub API')

    def get_token(self):
        while True:
            if not self.scheduler.get_jobs() and self.available_tokens.empty():
                self.print_warning('No more valid OAuth tokens available.')
                return None

            token = self.available_tokens.get(block=True)

            rate_limit_url = (
                'https://api.github.com/rate_limit?access_token={0}'
            ).format(token)
            status = url_to_json(rate_limit_url)

            # Throw away bad OAuth keys.
            if 'resources' not in status:
                self.print_warning(
                    'Invalid OAuth token supplied. Trying again...'
                )
                continue

            if status['resources']['core']['remaining'] > 0:
                self.available_tokens.put_nowait(token)
                return token
            else:
                self.scheduler.add_job(
                    self.available_tokens.put_nowait,
                    'date',
                    args=[token],
                    run_date=datetime.datetime.fromtimestamp(
                        status['resources']['core']['reset']
                    )
                )

    def print_warning(self, message):
        formatted_message = '\033[91mWARNING\033[0m: {0}'.format(message)
        print(formatted_message)

def get_run_id(cursor):
    '''
    Generate an id from the datasource that uniquely identifies the run.

    Args:
        cursor: mysql.cursor.MySQLCursor
            Cursor object used to insert data.

    Return:
        Identifier of the run to persist.
    '''
    cursor.execute('INSERT INTO reaper_runs VALUES()')

    return cursor.lastrowid

def save_result(run_id, repo_id, results, cursor):
    """
    Save the results to the specified data source, creating the results table
    as necessary.

    Args:
        run_id: int
            Identifier of the run to persist.
        repo_id: int
            Identifier of the repository to save.
        results: dict
            Key value pair of results to be saved.
        cursor: mysql.cursor.MySQLCursor
            Cursor object used to insert data.

    Return:
        True if successful, False otherwise.
    """
    query = '''
        INSERT INTO reaper_results
            (project_id, results, run_id)
        VALUES
            ({0}, '{1}', {2})
        '''.format(repo_id, json.dumps(results), run_id)

    cursor.execute(query)

def init_attribute_plugins(attributes, connection):
    for attribute in attributes:
        if 'implementation' in attribute:
            try:
                cursor = connection.cursor()
                attribute['implementation'].init(cursor)
            except:
                pass
            finally:
                cursor.close()

def global_init_attribute_plugins(attributes, connection, sample):
    for attribute in attributes:
        if 'implementation' in attribute:
            try:
                cursor = connection.cursor()
                attribute['implementation'].global_init(cursor, sample)
            except:
                pass
            finally:
                cursor.close()

def process_repository(project_id, repo_path, attributes, connection):
    """
    Calculate a score for the given repository.

    Args:
        project_id: int
            GHTorrent dataset identifier for the repository.
        repo_path: string
            Path to the repository contents.
        attributes: list
            List of attributes to be executed against the repository.
        connection: mysql.connector.MySQLConnection
            Database connection used for querying the GHTorrent dataset.
    """
    score = 0
    results = {}
    for attribute in attributes:
        if 'implementation' in attribute:
            try:
                cursor = connection.cursor()
            except mysql.connector.Error:
                sys.stderr.write(
                    'Error while creating cursor, attempting to reconnect.\n'
                )
                connection.reconnect()
            finally:
                cursor = connection.cursor()

            binary_result, raw_result = attribute['implementation'].run(
                project_id,
                repo_path,
                cursor,
                **attribute.get('options', {})
            )

            cursor.close()

            score += binary_result * attribute['weight']
            results[attribute['name']] = raw_result

            if ('essential' in attribute and
                    attribute['essential'] and
                    not binary_result):
                score = 0
                break

    return score, results

def load_attribute_plugins(plugins_dir, attributes):
    """
    Attempt to load each of the attributes as defined in the configuration
    file.

    Args:
        attributes: list
            List of attribute dictionaries with the specific configuration
            data.
    """
    for attribute in attributes:
        if attribute['enabled']:
            try:
                attribute['implementation'] = importlib.import_module(
                    '{0}.{1}.main'.format(plugins_dir, attribute['name'])
                )

            except ImportError:
                print('Failed to load the {0} attribute.'.format(
                        attribute['name']
                    )
                )

def establish_database_connection(config):
    """
    Attempt to establish a connection to the specified database. Exit with an
    error message on failure.

    Args:
        config: dict
            Settings for the database connection.
    """
    try:
        config['autocommit'] = True
        connection = mysql.connector.connect(**config)
        connection.connect()
        return connection
    except mysql.connector.Error:
        print('\rUnable to establish connection to database.')
        sys.exit(1)

def process_configuration(config_file):
    """
    Load and validate the given configuration file.

    Args:
        config_file: File
            File object with the configuration contents.

    Returns:
        Validated dictionary with configuration parameters.
    """
    global tokens

    try:
        config = json.load(config_file)
    except:
        print('Your config file is malformed, or does not exist!')
        sys.exit(1)
    finally:
        tokens = config['options'].get('github_tokens', [])
        return config

def process_plugins(plugins_dir):
    manifest = plugins_dir + '/manifest.json'
    try:
        with open(manifest, 'r') as f:
            plugins = json.load(f)

            attributes = plugins['attributes']
            threshold = plugins['threshold']

            for attribute in attributes:
                exit_if_dep_missing(attribute)

            return (threshold, attributes)

    except:
        print('{} is malformed, or does not exist!'.format(manifest))
        sys.exit(1)

def exit_if_dep_missing(attribute):
    if not attribute['enabled'] or 'dependencies' not in attribute:
        dependencies = []
    else:
        dependencies = attribute['dependencies']

    for dependency in dependencies:
        if not distutils.spawn.find_executable(dependency):
            print('Missing dependency for attribute {}: {}'
                .format(attribute['name'], dependency)
            )
            sys.exit(1)

def process_key_string(attributes, key_string):
    """Reprocesses the attribute configuration based on a keystring

    The keystring is a single string interpreted as a list of individual
    initials. Each initial indicates a particular attribute, as defined in
    config.json.

    If no keystring is provided, the attribute configuration is left untouched.
    If a keystring is provided, then first, all attributes are disabled. Only
    those attribtues referred by their initials in the keystring are
    re-enabled. Further, if the attribute initial in the keystring is
    capitalized, then that attribute's data will be persisted.
    """
    if key_string is None:
        return

    initials = dict()  # Attribute initial => attribute dict
    for attribute in attributes:
        initial = attribute['initial'].lower()
        initials[initial] = attribute
        attribute['enabled'] = False  # Disable all attributes first
        attribute['persistResult'] = False

    if key_string == "help":
        for key_low, attribute in initials.items():
            print(key_low + " " + attribute['name'])
        sys.exit()

    for key_raw in key_string:
        key_low = key_raw.lower()
        attribute = initials[key_low]
        # Enable attribute since its initial is in the keystring
        attribute['enabled'] = True
        # Persist results if the raw key is capitalized
        attribute['persistResult'] = (not key_raw == key_low)

def get_persist_attrs(attributes):
    persist_attrs = []
    for attribute in attributes:
        if attribute['persistResult']:
            persist_attrs.append(attribute['name'])

    return (len(persist_attrs) > 0, persist_attrs)
<|endoftext|>


<commit_before>#!/usr/bin/env python
'''
Iterate over all cgroup mountpoints and output global cgroup statistics
'''

import os
import sys
import stat
import pwd
import time
import psutil

from collections import defaultdict

try:
    import curses
except ImportError:
    print >> sys.stderr, "Curse is not available on this system. Exiting."
    sys.exit(0)

HIDE_EMPTY_CGROUP = True
UPDATE_INTERVAL = 1.0 # seconds
CGROUP_MOUNTPOINTS={}
CGROUPS_SUBSYS=[]

CGROUPS = {
    'prev': {},
    'cur': {},
}

# TODO:
# - curse list
# - select refresh rate
# - select sort column
# - visual CPU/memory usage
# - block-io
# - auto-color
# - adapt name / commands to underlying container system
# - hiereachical view

## Utils

def to_human(num, suffix='B'):
    num = int(num)
    for unit in ['','K','M','G','T','P','E','Z']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Y', suffix)

## Grab system facts

# Get all cgroup subsystems avalaible on this system
with open("/proc/cgroups") as f:
    cgroups = f.read().strip()

for cgroup in cgroups.split('\n'):
    if cgroup[0] == '#': continue
    CGROUPS_SUBSYS.append(cgroup.split()[0])

# Match cgroup mountpoints to susbsytems. Always take the first matching
with open("/proc/mounts") as f:
    mounts = f.read().strip()

for mount in mounts.split('\n'):
    mount = mount.split(' ')

    if mount[2] != "cgroup":
        continue

    for arg in mount[3].split(','):
        if arg in CGROUPS_SUBSYS and arg not in CGROUP_MOUNTPOINTS:
            CGROUP_MOUNTPOINTS[arg] = mount[1]

## Tools

class Cgroup(object):
    def __init__(self, path, base_path):
        self.path = path
        self.base_path = base_path

    @property
    def name(self):
        return self.path[len(self.base_path):] or '/'

    @property
    def owner(self):
        path = os.path.join(self.base_path, self.path, 'tasks')
        uid = os.stat(path).st_uid
        try:
            return pwd.getpwuid(uid).pw_name
        except:
            return uid

    def _coerce(self, value):
        try:
            return int(value)
        except:
            pass

        try:
            return float(value)
        except:
            pass

        return value

    def __getitem__(self, name):
        path = os.path.join(self.base_path, self.path, name)

        with open(path) as f:
            content = f.read().strip()

        if name == 'tasks' or '\n' in content:
            content = content.split('\n')

            if ' ' in content[0]:
                content = dict((l.split(' ') for l in content))
                for k, v in content.iteritems():
                    content[k] = self._coerce(v)
            else:
                content = [self._coerce(v) for v in content]

        else:
            content = self._coerce(content)

        return content

def cgroups(base_path):
    '''
    Generator of cgroups under path ``name``
    '''
    for cgroup_path, dirs, files in os.walk(base_path):
        yield Cgroup(cgroup_path, base_path)

## Grab cgroup data

def collect(measures):

    # Collect global data
    if 'cpuacct' in CGROUP_MOUNTPOINTS:
        # list all "folders" under mountpoint
        for cgroup in cgroups(CGROUP_MOUNTPOINTS['cpuacct']):
            # Collect tasks
            measures['data'][cgroup.name]['tasks'] = cgroup['tasks']

            # Collect user
            measures['data'][cgroup.name]['owner'] = cgroup.owner

    # Collect memory statistics
    if 'memory' in CGROUP_MOUNTPOINTS:
        # list all "folders" under mountpoint
        for cgroup in cgroups(CGROUP_MOUNTPOINTS['memory']):
            measures['data'][cgroup.name]['memory.usage_in_bytes'] = cgroup['memory.usage_in_bytes']
            measures['data'][cgroup.name]['memory.max_usage_in_bytes'] = cgroup['memory.usage_in_bytes']
            measures['data'][cgroup.name]['memory.limit_in_bytes'] = min(int(cgroup['memory.limit_in_bytes']), measures['global']['total_memory'])

    # Collect CPU statistics
    if 'cpuacct' in CGROUP_MOUNTPOINTS:
        # list all "folders" under mountpoint
        for cgroup in cgroups(CGROUP_MOUNTPOINTS['cpuacct']):
            # Collect CPU stats
            prev = measures['data'][cgroup.name].get('cpuacct.stat', None)
            measures['data'][cgroup.name]['cpuacct.stat'] = cgroup['cpuacct.stat']

            # Collect CPU increase on run > 1
            if prev is None:
                measures['data'][cgroup.name]['cpuacct.stat.diff'] = {'user':0, 'system':0}
            else:
                for key, value in measures['data'][cgroup.name]['cpuacct.stat'].iteritems():
                    measures['data'][cgroup.name]['cpuacct.stat.diff'][key] = value - prev[key]

def display(scr, measures, sort_key):
    # Time
    prev_time = measures['global'].get('time', -1)
    cur_time = time.time()
    time_delta = cur_time - prev_time
    measures['global']['time'] = cur_time

    # Sort
    results = sorted(measures['data'].iteritems(), key=lambda cgroup: int(cgroup[1].get(sort_key, 0)), reverse=True)

    # Display statistics: Find the biggest user
    curses.endwin()
    scr.addstr(0, 0, '                memory                  cpu')
    scr.addstr(1, 0, 'owner      proc current         peak    system user cgroup')
    LINE_TMPL = "{:10s} {:4d} {:15s} {:7s} {: >5.1%} {: >5.1%} {}"
    cpu_to_percent = measures['global']['scheduler_frequency'] * measures['global']['total_cpu'] * time_delta

    lineno = 2
    for cgroup, data in results:
        cpu_usage = data.get('cpuacct.stat.diff', {})
        line = (
            data.get('owner', 'nobody'),
            len(data['tasks']),
            "%s/%s" % (to_human(data.get('memory.usage_in_bytes', 0)), to_human(data.get('memory.limit_in_bytes', measures['global']['total_memory']))),
            to_human(data.get('memory.max_usage_in_bytes', 0)),
            cpu_usage.get('system', 0) / cpu_to_percent,
            cpu_usage.get('user', 0) / cpu_to_percent,
            cgroup,
        )

        scr.addstr(lineno, 0, LINE_TMPL.format(*line))
        lineno += 1

    stdscr.refresh()

if __name__ == "__main__":
    # Initialization, global system data
    measures = {
        'data': defaultdict(dict),
        'global': {
            'total_cpu': psutil.cpu_count(),
            'total_memory': psutil.virtual_memory().total,
            'scheduler_frequency': os.sysconf('SC_CLK_TCK'),
        }
    }

    # Curse initialization
    stdscr = curses.initscr()
    curses.noecho()    # do not echo text
    curses.cbreak()    # do not wait for "enter"
    curses.curs_set(0) # hide cursor
    stdscr.keypad(1)   # parse keypad controll sequences

    try:
        while True:
            collect(measures)
            display(stdscr, measures, 'memory.usage_in_bytes')
            time.sleep(1.0)
    except KeyboardInterrupt:
        pass
    finally:
        curses.nocbreak()
        stdscr.keypad(0)
        curses.echo()
        curses.endwin()

<commit_msg>use refresh rate global parameter<commit_after>#!/usr/bin/env python
'''
Iterate over all cgroup mountpoints and output global cgroup statistics
'''

import os
import sys
import stat
import pwd
import time
import psutil

from collections import defaultdict

try:
    import curses
except ImportError:
    print >> sys.stderr, "Curse is not available on this system. Exiting."
    sys.exit(0)

HIDE_EMPTY_CGROUP = True
UPDATE_INTERVAL = 1.0 # seconds
CGROUP_MOUNTPOINTS={}
CGROUPS_SUBSYS=[]

CGROUPS = {
    'prev': {},
    'cur': {},
}

# TODO:
# - curse list
# - select refresh rate
# - select sort column
# - visual CPU/memory usage
# - block-io
# - auto-color
# - adapt name / commands to underlying container system
# - hiereachical view

## Utils

def to_human(num, suffix='B'):
    num = int(num)
    for unit in ['','K','M','G','T','P','E','Z']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Y', suffix)

## Grab system facts

# Get all cgroup subsystems avalaible on this system
with open("/proc/cgroups") as f:
    cgroups = f.read().strip()

for cgroup in cgroups.split('\n'):
    if cgroup[0] == '#': continue
    CGROUPS_SUBSYS.append(cgroup.split()[0])

# Match cgroup mountpoints to susbsytems. Always take the first matching
with open("/proc/mounts") as f:
    mounts = f.read().strip()

for mount in mounts.split('\n'):
    mount = mount.split(' ')

    if mount[2] != "cgroup":
        continue

    for arg in mount[3].split(','):
        if arg in CGROUPS_SUBSYS and arg not in CGROUP_MOUNTPOINTS:
            CGROUP_MOUNTPOINTS[arg] = mount[1]

## Tools

class Cgroup(object):
    def __init__(self, path, base_path):
        self.path = path
        self.base_path = base_path

    @property
    def name(self):
        return self.path[len(self.base_path):] or '/'

    @property
    def owner(self):
        path = os.path.join(self.base_path, self.path, 'tasks')
        uid = os.stat(path).st_uid
        try:
            return pwd.getpwuid(uid).pw_name
        except:
            return uid

    def _coerce(self, value):
        try:
            return int(value)
        except:
            pass

        try:
            return float(value)
        except:
            pass

        return value

    def __getitem__(self, name):
        path = os.path.join(self.base_path, self.path, name)

        with open(path) as f:
            content = f.read().strip()

        if name == 'tasks' or '\n' in content:
            content = content.split('\n')

            if ' ' in content[0]:
                content = dict((l.split(' ') for l in content))
                for k, v in content.iteritems():
                    content[k] = self._coerce(v)
            else:
                content = [self._coerce(v) for v in content]

        else:
            content = self._coerce(content)

        return content

def cgroups(base_path):
    '''
    Generator of cgroups under path ``name``
    '''
    for cgroup_path, dirs, files in os.walk(base_path):
        yield Cgroup(cgroup_path, base_path)

## Grab cgroup data

def collect(measures):

    # Collect global data
    if 'cpuacct' in CGROUP_MOUNTPOINTS:
        # list all "folders" under mountpoint
        for cgroup in cgroups(CGROUP_MOUNTPOINTS['cpuacct']):
            # Collect tasks
            measures['data'][cgroup.name]['tasks'] = cgroup['tasks']

            # Collect user
            measures['data'][cgroup.name]['owner'] = cgroup.owner

    # Collect memory statistics
    if 'memory' in CGROUP_MOUNTPOINTS:
        # list all "folders" under mountpoint
        for cgroup in cgroups(CGROUP_MOUNTPOINTS['memory']):
            measures['data'][cgroup.name]['memory.usage_in_bytes'] = cgroup['memory.usage_in_bytes']
            measures['data'][cgroup.name]['memory.max_usage_in_bytes'] = cgroup['memory.usage_in_bytes']
            measures['data'][cgroup.name]['memory.limit_in_bytes'] = min(int(cgroup['memory.limit_in_bytes']), measures['global']['total_memory'])

    # Collect CPU statistics
    if 'cpuacct' in CGROUP_MOUNTPOINTS:
        # list all "folders" under mountpoint
        for cgroup in cgroups(CGROUP_MOUNTPOINTS['cpuacct']):
            # Collect CPU stats
            prev = measures['data'][cgroup.name].get('cpuacct.stat', None)
            measures['data'][cgroup.name]['cpuacct.stat'] = cgroup['cpuacct.stat']

            # Collect CPU increase on run > 1
            if prev is None:
                measures['data'][cgroup.name]['cpuacct.stat.diff'] = {'user':0, 'system':0}
            else:
                for key, value in measures['data'][cgroup.name]['cpuacct.stat'].iteritems():
                    measures['data'][cgroup.name]['cpuacct.stat.diff'][key] = value - prev[key]

def display(scr, measures, sort_key):
    # Time
    prev_time = measures['global'].get('time', -1)
    cur_time = time.time()
    time_delta = cur_time - prev_time
    measures['global']['time'] = cur_time

    # Sort
    results = sorted(measures['data'].iteritems(), key=lambda cgroup: int(cgroup[1].get(sort_key, 0)), reverse=True)

    # Display statistics: Find the biggest user
    curses.endwin()
    scr.addstr(0, 0, '                memory                  cpu')
    scr.addstr(1, 0, 'owner      proc current         peak    system user cgroup')
    LINE_TMPL = "{:10s} {:4d} {:15s} {:7s} {: >5.1%} {: >5.1%} {}"
    cpu_to_percent = measures['global']['scheduler_frequency'] * measures['global']['total_cpu'] * time_delta

    lineno = 2
    for cgroup, data in results:
        cpu_usage = data.get('cpuacct.stat.diff', {})
        line = (
            data.get('owner', 'nobody'),
            len(data['tasks']),
            "%s/%s" % (to_human(data.get('memory.usage_in_bytes', 0)), to_human(data.get('memory.limit_in_bytes', measures['global']['total_memory']))),
            to_human(data.get('memory.max_usage_in_bytes', 0)),
            cpu_usage.get('system', 0) / cpu_to_percent,
            cpu_usage.get('user', 0) / cpu_to_percent,
            cgroup,
        )

        scr.addstr(lineno, 0, LINE_TMPL.format(*line))
        lineno += 1

    stdscr.refresh()

if __name__ == "__main__":
    # Initialization, global system data
    measures = {
        'data': defaultdict(dict),
        'global': {
            'total_cpu': psutil.cpu_count(),
            'total_memory': psutil.virtual_memory().total,
            'scheduler_frequency': os.sysconf('SC_CLK_TCK'),
        }
    }

    # Curse initialization
    stdscr = curses.initscr()
    curses.noecho()    # do not echo text
    curses.cbreak()    # do not wait for "enter"
    curses.curs_set(0) # hide cursor
    stdscr.keypad(1)   # parse keypad controll sequences

    try:
        while True:
            collect(measures)
            display(stdscr, measures, 'memory.usage_in_bytes')
            time.sleep(UPDATE_INTERVAL)
    except KeyboardInterrupt:
        pass
    finally:
        curses.nocbreak()
        stdscr.keypad(0)
        curses.echo()
        curses.endwin()

<|endoftext|>


<commit_before>import ast, argparse, statistics, textwrap, csv, json, os
from PIL import Image, ImageDraw, ImageFont
from voteConverter import convert
from booksonaGen import make_book

#fonts
arial = ImageFont.truetype('./resources/arial.ttf',20)
bigArial =  ImageFont.truetype('./resources/arial.ttf',30)
smallArial = ImageFont.truetype('./resources/arial.ttf',13)

def parse_args():
	parser = argparse.ArgumentParser()
	parser.add_argument('input')
	parser.add_argument("-e", "--perc_elim", nargs='?', const=5, default=5)
	parser.add_argument("-t", "--num_gold", nargs='?', const=5, default=5)
	args = parser.parse_args()
	path = args.input

	convert(path)
	prompt = open('./{}/prompt.txt'.format(path),'r').read().split('\n')[0]

	twowers = []
	responses = []
	boosts = {}

	with open('./{}/responses.csv'.format(path),'r') as csvfile:#read responses
		reader = csv.reader(csvfile)
		for row in reader:
			twowers.append(row[0])
			responses.append(row[1])

			try:
				boosts[row[0]] = int(row[2])
			except:
				pass

	scores = [[response,[]]for response in responses]
	votes = json.load(open('./{}/votes.json'.format(path),'r'))

	indiv_twowers = list(set(twowers))#data about the data (metadata?)
	twower_count = len(indiv_twowers)

	top_number = int(args.num_gold) #chart coloring ranges
	elim_number=0
	if int(args.perc_elim) < 0:
		elim_number = -int(args.perc_elim)
	else:
		elim_number = round(int(args.perc_elim)*len(indiv_twowers)/100)

	return (path, prompt, twowers, responses, scores, votes, indiv_twowers, twower_count, top_number, elim_number, boosts)

def wrap_text(text,chars):
	lines = textwrap.wrap(text,chars)
	new_text='\n'.join(lines)
	return new_text

def draw_header(prompt, base, drawer, responses):
	prompt = wrap_text(prompt,100)
	header_height = drawer.textsize(prompt,bigArial)[1]+35
	base = Image.new('RGBA',(1368,header_height+int(67/2*len(responses))),color=(255,255,255,255))
	drawer = ImageDraw.Draw(base)
	base.paste(Image.open('./resources/header.png'),(0,header_height-40))
	drawer.text((15,0),prompt,font=bigArial, fill=(0,0,0,255))

	return (prompt, base, drawer, header_height)

def remove_dups(seq):
    seen = set()
    seen_add = seen.add
    return [x for x in seq if not (x in seen or seen_add(x))]

def process_votes(votes, scores, twowers, boosts):
	for vote in votes:#maps votes to responses
		vote = remove_dups(vote)
		try:
			percentage = 100
			for resp_num in vote:
				scores[resp_num][1].append(percentage)
				percentage -= 11
		except Exception:
			pass

	for scoredata in scores:
		scoredata = calc_stats(scoredata,scores.index(scoredata),twowers,boosts)

	mergeSort(scores)#sorts from best to worst. Mergesort for best worst case
	return scores

def calc_stats(scoredata,resp_ind,twowers,boosts):#calculate stats, dm if you want more
	try:
		scoredata.append(statistics.mean(scoredata[1]))
	except Exception:
		print('\"{}\" by {} was not voted for'.format(scoredata[0],twowers[resp_ind]))
		scoredata.append(0)
	try:
		scoredata[2] += boosts[twowers[resp_ind]]
		if scoredata[2]>100:
			scoredata[2]=100
	except:
		pass

	try:
		scoredata.append(statistics.stdev(scoredata[1]))
	except Exception:
		scoredata.append(0)

	scoredata.append(len(scoredata[1]))#number of votes
	scoredata[1]= twowers[resp_ind]#twower name
	scoredata[0],scoredata[1]=scoredata[1],scoredata[0]#rearranges list in order on chart
	return scoredata

def draw_rankings(scores, top_number, elim_number,twower_count,base,drawer,header_height,indiv_twowers,boosts):#this one is a bit of a mess
	backgroundCol=0
	addBackground=0
	ranking=1

	for i in range(len(scores)):
		twower, response, mean, standev, voteCount = scores[i][0], scores[i][1], scores[i][2], scores[i][3], scores[i][4]

		if ranking == (top_number+1):#change background depending on ranking
			backgroundCol = 1
			addBackground = 0
		elif ranking == (twower_count-elim_number+1) and twower in indiv_twowers:
			backgroundCol = 2
			addBackground = 0

		if (addBackground % 2) ==0:#only needs extra stuff added every two twowers
			if backgroundCol==0:
				base.paste(Image.open('./resources/top.png'),(0,int(67/2*i)+header_height))
			elif backgroundCol==1:
				base.paste(Image.open('./resources/normal.png'),(0,int(67/2*i)+header_height))
			elif backgroundCol==2:
				base.paste(Image.open('./resources/eliminated.png'),(0,int(67/2*i)+header_height))

		if not os.path.isfile('./booksonas/'+twower+'.png'):
			make_book(twower,'./booksonas/')

		try:#attempt to add booksona
			booksona = Image.open('./booksonas/'+twower+'.png')
			booksona.thumbnail((32,32),Image.BICUBIC)
			base.paste(booksona,(333,int(67/2*i)+header_height),booksona)
		except:
			pass

		if twower in indiv_twowers:#handles multiple submissions
			indiv_twowers.remove(twower)
			if ranking % 10 == 1:
				rankingString = str(ranking)+'st'
			elif ranking % 10 == 2:
				rankingString = str(ranking)+'nd'
			elif ranking % 10 == 3:
				rankingString = str(ranking)+'rd'
			else:
				rankingString = str(ranking)+'th'

			drawer.text((15,int(67/2*i+7)+header_height),rankingString,font=arial,fill=(0,0,0,255))
			ranking += 1

		if drawer.textsize(twower,arial)[0] > 255: #draws twower name
			drawer.text((320-drawer.textsize(twower,smallArial)[0],int(67/2*i+7)+header_height),
				twower,font=smallArial,fill=(0,0,0,255))

		else:
			drawer.text((320-drawer.textsize(twower,arial)[0],int(67/2*i+7)+header_height),
				twower,font=arial,fill=(0,0,0,255))

		if drawer.textsize(response,arial)[0] > 618: #draws responses
			responseLines = textwrap.wrap(response,90)
			response = ''
			for s in responseLines:
				response += (s+'\n')

			drawer.text((378,int(67/2*i)+header_height),
				response,font=smallArial,fill=(0,0,0,255))
		else:
			drawer.text((378,int(67/2*i+7)+header_height),
				response,font=arial,fill=(0,0,0,255))

		draw_stats(drawer,twower,mean,standev,boosts,voteCount,header_height,i)

		addBackground += 1

	return base

def draw_stats(drawer,twower,mean,standev,boosts,voteCount,header_height,rank):
	try:
		mean_str = str(mean-boosts[twower])[:4]
		mean_str += '(+{})'.format(boosts[twower])
		mean_str += '%'

		drawer.text((998,int(67/2*rank+7)+header_height),
			mean_str,font=arial,fill=(0,0,0,255))
	except:
		drawer.text((998,int(67/2*rank+7)+header_height),
			str(mean)[:5]+'%',font=arial,fill=(0,0,0,255))

	drawer.text((1164,int(67/2*rank+7)+header_height),
		str(standev)[:5]+'%',font=arial,fill=(0,0,0,255))

	drawer.text((1309-drawer.textsize(str(voteCount),arial)[0]/2,
		int(67/2*rank+7)+header_height),str(voteCount),
		font=arial,fill=(0,0,0,255))

def mergeSort(alist):
    if len(alist)>1:
        mid = len(alist)//2
        lefthalf = alist[:mid]
        righthalf = alist[mid:]

        mergeSort(lefthalf)
        mergeSort(righthalf)

        i=0
        j=0
        k=0
        while i < len(lefthalf) and j < len(righthalf):
            if lefthalf[i][2] > righthalf[j][2]:
                alist[k]=lefthalf[i]
                i=i+1
            else:
                alist[k]=righthalf[j]
                j=j+1
            k=k+1

        while i < len(lefthalf):
            alist[k]=lefthalf[i]
            i=i+1
            k=k+1

        while j < len(righthalf):
            alist[k]=righthalf[j]
            j=j+1
            k=k+1

def main():
	path, prompt, twowers, responses, scores, votes, indiv_twowers, twower_count, top_number, elim_number, boosts = parse_args()

	base = Image.new('RGBA',(1368,1368),color=(255,255,255))
	drawer = ImageDraw.Draw(base)

	prompt, base, drawer, header_height = draw_header(prompt, base, drawer, responses)
	scores = process_votes(votes, scores, twowers, boosts)
	draw_rankings(scores,top_number,elim_number,twower_count,base,drawer,header_height,indiv_twowers, boosts)

	base.save('./'+path+'/results.png')

if __name__=='__main__':
	main()
<commit_msg>Bit more cleanup<commit_after>import ast, argparse, statistics, textwrap, csv, json, os
from PIL import Image, ImageDraw, ImageFont
from voteConverter import convert
from booksonaGen import make_book

#fonts
arial = ImageFont.truetype('./resources/arial.ttf',20)
bigArial =  ImageFont.truetype('./resources/arial.ttf',30)
smallArial = ImageFont.truetype('./resources/arial.ttf',13)

def parse_args():
	parser = argparse.ArgumentParser()
	parser.add_argument('input')
	parser.add_argument("-e", "--perc_elim", nargs='?', const=5, default=5)
	parser.add_argument("-t", "--num_gold", nargs='?', const=5, default=5)
	args = parser.parse_args()
	path = args.input

	convert(path)
	prompt = open('./{}/prompt.txt'.format(path),'r').read().split('\n')[0]

	twowers = []
	responses = []
	boosts = {}

	with open('./{}/responses.csv'.format(path),'r') as csvfile:#read responses
		reader = csv.reader(csvfile)
		for row in reader:
			twowers.append(row[0])
			responses.append(row[1])

			try:
				boosts[row[0]] = int(row[2])
			except:
				pass

	scores = [[response,[]]for response in responses]
	votes = json.load(open('./{}/votes.json'.format(path),'r'))

	indiv_twowers = list(set(twowers))#data about the data (metadata?)
	twower_count = len(indiv_twowers)

	top_number = int(args.num_gold) #chart coloring ranges
	elim_number=0
	if int(args.perc_elim) < 0:
		elim_number = -int(args.perc_elim)
	else:
		elim_number = round(int(args.perc_elim)*len(indiv_twowers)/100)

	return (path, prompt, twowers, responses, scores, votes, indiv_twowers, twower_count, top_number, elim_number, boosts)

def wrap_text(text,chars):
	lines = textwrap.wrap(text,chars)
	new_text='\n'.join(lines)
	return new_text

def draw_header(prompt, base, drawer, responses):
	prompt = wrap_text(prompt,100)
	header_height = drawer.textsize(prompt,bigArial)[1]+35
	base = Image.new('RGBA',(1368,header_height+int(67/2*len(responses))),color=(255,255,255,255))
	drawer = ImageDraw.Draw(base)
	base.paste(Image.open('./resources/header.png'),(0,header_height-40))
	drawer.text((15,0),prompt,font=bigArial, fill=(0,0,0,255))

	return (prompt, base, drawer, header_height)

def remove_dups(seq):
    seen = set()
    seen_add = seen.add
    return [x for x in seq if not (x in seen or seen_add(x))]

def process_votes(votes, scores, twowers, boosts):
	for vote in votes:#maps votes to responses
		vote = remove_dups(vote)
		try:
			percentage = 100
			for resp_num in vote:
				scores[resp_num][1].append(percentage)
				percentage -= 11
		except Exception:
			pass

	for scoredata in scores:
		scoredata = calc_stats(scoredata,scores.index(scoredata),twowers,boosts)

	mergeSort(scores)#sorts from best to worst. Mergesort for best worst case
	return scores

def calc_stats(scoredata,resp_ind,twowers,boosts):#calculate stats, dm if you want more
	try:
		scoredata.append(statistics.mean(scoredata[1]))
		'''
		if twowers[resp_ind].startswith('hanss314'):
			scoredata[2]=1000
		'''
	except Exception:
		print('\"{}\" by {} was not voted for'.format(scoredata[0],twowers[resp_ind]))
		scoredata.append(0)
	try:
		scoredata[2] += boosts[twowers[resp_ind]]
		if scoredata[2]>100:
			scoredata[2]=100
	except:
		pass

	try:
		scoredata.append(statistics.stdev(scoredata[1]))
	except Exception:
		scoredata.append(0)

	scoredata.append(len(scoredata[1]))#number of votes
	scoredata[1]= twowers[resp_ind]#twower name
	scoredata[0],scoredata[1]=scoredata[1],scoredata[0]#rearranges list in order on chart
	return scoredata

def draw_rankings(scores, top_number, elim_number,twower_count,base,drawer,header_height,indiv_twowers,boosts):#this one is a bit of a mess
	backgroundCol=0
	addBackground=0
	ranking=1

	for i in range(len(scores)):
		twower, response, mean, standev, voteCount = scores[i][0], scores[i][1], scores[i][2], scores[i][3], scores[i][4]

		if ranking == (top_number+1):#change background depending on ranking
			backgroundCol = 1
			addBackground = 0
		elif ranking == (twower_count-elim_number+1) and twower in indiv_twowers:
			backgroundCol = 2
			addBackground = 0

		if (addBackground % 2) ==0:#only needs extra stuff added every two twowers
			if backgroundCol==0:
				base.paste(Image.open('./resources/top.png'),(0,int(67/2*i)+header_height))
			elif backgroundCol==1:
				base.paste(Image.open('./resources/normal.png'),(0,int(67/2*i)+header_height))
			elif backgroundCol==2:
				base.paste(Image.open('./resources/eliminated.png'),(0,int(67/2*i)+header_height))

		if not os.path.isfile('./booksonas/'+twower+'.png'):
			make_book(twower,'./booksonas/')

		try:#attempt to add booksona
			booksona = Image.open('./booksonas/'+twower+'.png')
			booksona.thumbnail((32,32),Image.BICUBIC)
			base.paste(booksona,(333,int(67/2*i)+header_height),booksona)
		except:
			pass

		if twower in indiv_twowers:#handles multiple submissions
			indiv_twowers.remove(twower)
			if ranking % 10 == 1:
				rankingString = str(ranking)+'st'
			elif ranking % 10 == 2:
				rankingString = str(ranking)+'nd'
			elif ranking % 10 == 3:
				rankingString = str(ranking)+'rd'
			else:
				rankingString = str(ranking)+'th'

			drawer.text((15,int(67/2*i+7)+header_height),rankingString,font=arial,fill=(0,0,0,255))
			ranking += 1

		if drawer.textsize(twower,arial)[0] > 255: #draws twower name
			drawer.text((320-drawer.textsize(twower,smallArial)[0],int(67/2*i+7)+header_height),
				twower,font=smallArial,fill=(0,0,0,255))

		else:
			drawer.text((320-drawer.textsize(twower,arial)[0],int(67/2*i+7)+header_height),
				twower,font=arial,fill=(0,0,0,255))

		if drawer.textsize(response,arial)[0] > 618: #draws responses
			responseLines = textwrap.wrap(response,90)
			response = ''
			for s in responseLines:
				response += (s+'\n')

			drawer.text((378,int(67/2*i)+header_height),
				response,font=smallArial,fill=(0,0,0,255))
		else:
			drawer.text((378,int(67/2*i+7)+header_height),
				response,font=arial,fill=(0,0,0,255))

		draw_stats(drawer,twower,mean,standev,boosts,voteCount,header_height,i)

		addBackground += 1

	return base

def draw_stats(drawer,twower,mean,standev,boosts,voteCount,header_height,rank):
	try:
		mean_str = str(mean-boosts[twower])[:4]
		mean_str += '(+{})'.format(boosts[twower])
		mean_str += '%'

		drawer.text((998,int(67/2*rank+7)+header_height),
			mean_str,font=arial,fill=(0,0,0,255))
	except:
		drawer.text((998,int(67/2*rank+7)+header_height),
			str(mean)[:5]+'%',font=arial,fill=(0,0,0,255))

	drawer.text((1164,int(67/2*rank+7)+header_height),
		str(standev)[:5]+'%',font=arial,fill=(0,0,0,255))

	drawer.text((1309-drawer.textsize(str(voteCount),arial)[0]/2,
		int(67/2*rank+7)+header_height),str(voteCount),
		font=arial,fill=(0,0,0,255))

def mergeSort(alist):
    if len(alist)>1:
        mid = len(alist)//2
        lefthalf = alist[:mid]
        righthalf = alist[mid:]

        mergeSort(lefthalf)
        mergeSort(righthalf)

        i=0
        j=0
        k=0
        while i < len(lefthalf) and j < len(righthalf):
            if lefthalf[i][2] > righthalf[j][2]:
                alist[k]=lefthalf[i]
                i=i+1
            else:
                alist[k]=righthalf[j]
                j=j+1
            k=k+1

        while i < len(lefthalf):
            alist[k]=lefthalf[i]
            i=i+1
            k=k+1

        while j < len(righthalf):
            alist[k]=righthalf[j]
            j=j+1
            k=k+1

def main():
	path, prompt, twowers, responses, scores, votes, indiv_twowers, twower_count, top_number, elim_number, boosts = parse_args()

	base = Image.new('RGBA',(1368,1368),color=(255,255,255))
	drawer = ImageDraw.Draw(base)

	prompt, base, drawer, header_height = draw_header(prompt, base, drawer, responses)
	scores = process_votes(votes, scores, twowers, boosts)
	draw_rankings(scores,top_number,elim_number,twower_count,base,drawer,header_height,indiv_twowers, boosts)

	base.save('./'+path+'/results.png')

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


<commit_before>"""" Function definitions """
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import urllib
import threading
import multiprocessing
import requests
from GoogleScraper import scrape_with_config, GoogleSearchError

def get_immediate_subdirectories(a_dir):
    """ Get only the immediate subfolders """
    return [name for name in os.listdir(a_dir)
            if os.path.isdir(
                os.path.join(
                    a_dir, name))]

drive = "C:\\Users\\nnikh\\Google Drive"
author = os.path.join(drive,"nikhilatphyzok")
projectpath = os.path.join(author,"scraping")
procpathS = get_immediate_subdirectories(projectpath)

def uprint(*objects, sep=' ', end='\n', file=sys.stdout):
    """This is just a print wrapper"""
    enc = file.encoding
    if enc == 'UTF-8':
        print(*objects, sep=sep, end=end, file=file)
    else:
        fup = lambda obj: str(obj).encode(
            enc, errors='replace').decode(enc)
        print(
            *map(fup, objects), sep=sep, end=end, file=file)

def chapterscraper(aprocpath):
    """ Reads the contents of a chapter
    and calls phrase operations"""

    kwfile = os.path.join(
        projectpath, aprocpath, 'phraselist.txt')
    with open(
        kwfile, mode='r', encoding='utf-8') as g:
        chapterphrases = list(
            set(g.read().splitlines()))
    uprint("Read all phrases for {}\n\n".
           format(aprocpath))

    for phrase in chapterphrases:
        phrase = phrase.strip()
    phraseprocs = [multiprocessing.Process(
        target=phrasescraper(
            phrase, aprocpath))
                   for phrase in chapterphrases]

    if __name__ == '__main__':
        for phraseproc in phraseprocs:
            phraseproc.daemon = False
            phraseproc.start()
            uprint("Running phrase operations")
        for phraseproc in phraseprocs:
            phraseproc.join()
        uprint('Phrase operations for {}'
               .format(aprocpath.split('\\')[-1]))

def phrasescraper(aphrase, aprocpath):
    """ Gets images for a phrase
    and writes to the phrase folder """
    config = {
        'keyword': aphrase,
        'database_name' : 'redox',
        'print_results' : 'summarize',
        'log_level' : 'INFO',
        'num_results_per_page' : 30,
        'num_pages_for_keyword' : 1,
        'num_workers' : 1,
        'maximum_workers' : 2
        }

    try:
        search = scrape_with_config(config)
    except GoogleSearchError as e:
        uprint(e)

    image_urls = []
    for serp in search.serps:
        image_urls.extend(
            [link.link for link in serp.links])

    num_threads = 30
    phraseimages = os.path.join(
        projectpath, aprocpath, 'images', aphrase)
    threads = [FetchResource(
        phraseimages, []) for i in range(
            num_threads)]
    while image_urls:
        for t in threads:
            try:
                t.furls.append(
                    image_urls.pop())
            except IndexError:
                break

    threads = [t for t in threads if t.furls]
    b = 0
    for t in threads:
        t.start()
    for t in threads:
        t.join()
        b += 1
    uprint('downloaded {} images for {}'.format(b, aphrase))

    fileiter = (os.path.join(
        root, f) for root, _, files in os.walk(
            aprocpath +'\\images\\' + aphrase)
                for f in files)
    smallfileiter = (f for f in fileiter
                     if os.path.getsize(f) < 300 * 800)
    for small in smallfileiter:
        os.remove(small)
        print("\nRemoved small images from {}"
              .format(aphrase))

class FetchResource(threading.Thread):
    """ Gets the content of a url """

    def __init__(self, target, furls):
        super().__init__()
        self.target = target
        self.furls = furls

    def run(self):
        if not os.path.exists(
                self.target.strip()):
            os.makedirs(
                self.target.strip())
        for furl in self.furls:
            furl = urllib.parse.unquote(furl)
            name = furl.split('/')[-1]
            with open(
                os.path.join(
                    self.target, name), 'wb') as fname:
                content = requests.get(furl).content
                fname.write(content)

def main():
    """ Reads target folder
    and calls chapter operations """
    chapterprocs = [multiprocessing.Process(
        target=chapterscraper(
            procpath))
                    for procpath in procpathS]
    if __name__ == '__main__':
        for chapterproc in chapterprocs:
            chapterproc.start()
        for chapterproc in chapterprocs:
            chapterproc.join()
    uprint("Finished processing everything")

if __name__ == '__main__':
<commit_msg>Update ox.py<commit_after>"""" Function definitions """
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import urllib
import threading
import multiprocessing
import requests
from GoogleScraper import scrape_with_config, GoogleSearchError

def get_immediate_subdirectories(a_dir):
    """ Get only the immediate subfolders """
    return [name for name in os.listdir(a_dir)
            if os.path.isdir(
                os.path.join(
                    a_dir, name))]

drive = "C:\\Users\\nnikh\\Google Drive"
author = os.path.join(drive,"nikhilatphyzok")
projectpath = os.path.join(author,"automation", "scraping")
procpathS = get_immediate_subdirectories(projectpath)

def uprint(*objects, sep=' ', end='\n', file=sys.stdout):
    """This is just a print wrapper"""
    enc = file.encoding
    if enc == 'UTF-8':
        print(*objects, sep=sep, end=end, file=file)
    else:
        fup = lambda obj: str(obj).encode(
            enc, errors='replace').decode(enc)
        print(
            *map(fup, objects), sep=sep, end=end, file=file)

def chapterscraper(aprocpath):
    """ Reads the contents of a chapter
    and calls phrase operations"""

    kwfile = os.path.join(
        projectpath, aprocpath, 'phraselist.txt')
    with open(
        kwfile, mode='r', encoding='utf-8') as g:
        chapterphrases = list(
            set(g.read().splitlines()))
    uprint("Read all phrases for {}\n\n".
           format(aprocpath))

    for phrase in chapterphrases:
        phrase = phrase.strip()
    phraseprocs = [multiprocessing.Process(
        target=phrasescraper(
            phrase, aprocpath))
                   for phrase in chapterphrases]

    if __name__ == '__main__':
        for phraseproc in phraseprocs:
            phraseproc.daemon = False
            phraseproc.start()
            uprint("Running phrase operations")
        for phraseproc in phraseprocs:
            phraseproc.join()
        uprint('Phrase operations for {}'
               .format(aprocpath.split('\\')[-1]))

def phrasescraper(aphrase, aprocpath):
    """ Gets images for a phrase
    and writes to the phrase folder """
    config = {
        'keyword': aphrase,
        'database_name' : 'redox',
        'print_results' : 'summarize',
        'log_level' : 'INFO',
        'num_results_per_page' : 30,
        'num_pages_for_keyword' : 1,
        'num_workers' : 1,
        'maximum_workers' : 2
        }

    try:
        search = scrape_with_config(config)
    except GoogleSearchError as e:
        uprint(e)

    image_urls = []
    for serp in search.serps:
        image_urls.extend(
            [link.link for link in serp.links])

    num_threads = 30
    phraseimages = os.path.join(
        projectpath, aprocpath, 'images', aphrase)
    threads = [FetchResource(
        phraseimages, []) for i in range(
            num_threads)]
    while image_urls:
        for t in threads:
            try:
                t.furls.append(
                    image_urls.pop())
            except IndexError:
                break

    threads = [t for t in threads if t.furls]
    b = 0
    for t in threads:
        t.start()
    for t in threads:
        t.join()
        b += 1
    uprint('downloaded {} images for {}'.format(b, aphrase))

    fileiter = (os.path.join(
        root, f) for root, _, files in os.walk(
            aprocpath +'\\images\\' + aphrase)
                for f in files)
    smallfileiter = (f for f in fileiter
                     if os.path.getsize(f) < 300 * 800)
    for small in smallfileiter:
        os.remove(small)
        print("\nRemoved small images from {}"
              .format(aphrase))

class FetchResource(threading.Thread):
    """ Gets the content of a url """

    def __init__(self, target, furls):
        super().__init__()
        self.target = target
        self.furls = furls

    def run(self):
        if not os.path.exists(
                self.target.strip()):
            os.makedirs(
                self.target.strip())
        for furl in self.furls:
            furl = urllib.parse.unquote(furl)
            name = furl.split('/')[-1]
            with open(
                os.path.join(
                    self.target, name), 'wb') as fname:
                content = requests.get(furl).content
                fname.write(content)

def main():
    """ Reads target folder
    and calls chapter operations """
    chapterprocs = [multiprocessing.Process(
        target=chapterscraper(
            procpath))
                    for procpath in procpathS]
    if __name__ == '__main__':
        for chapterproc in chapterprocs:
            chapterproc.start()
        for chapterproc in chapterprocs:
            chapterproc.join()
    uprint("Finished processing everything")

if __name__ == '__main__':
<|endoftext|>


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

class Buffer(object):
    """
    Represents unicode text using a gap buffer. Can be initialized with any
    iterable item, as long as that item yields unicode characters during
    iteration.
    """

    def __init__(self, initial_content=None):

        # points to first empty space after left text
        self.__gap_start = 0

        # points to first taken space at beginning of right text
        self.__gap_end = 0

        # where the gap should be moved to when it needs to be moved
        self.__cursor = 0

        # whether the cursor has been moved without the gap being moved
        self.__cursor_dirty = False

        # the minimum amount of space to create when resizing the gap
        self.__min_gap_len = 10

        # the end of content in the buffer
        self.__content_end = 0

        # NOTE: we set the content end individually later so we can allow
        # initial_content objects that have no __len__ method.

        # initialize the buffer we use to hold text data in, a unicode array
        if initial_content is None:
            # insert placeholder content if no initial content was specified
            self.__buf = array.array("u", u" " * 10)
            self.__content_end = 0
        else:
            # insert the initial content and set the content end pointer for it
            self.__buf = array.array("u", initial_content)
            self.__content_end = len(self.__buf)

    @property
    def cursor(self):
        """Get the current position of the cursor."""
        return self.__cursor

    @cursor.setter
    def cursor(self, position):
        """Move the cursor within the buffer."""

        # NOTE: does not move the gap! when moved, the cursor position is marked
        # as 'dirty' so that when text changes happen, the gap can be moved
        # accordingly elsewhere.

        # store the former cursor position
        old_cursor = self.__cursor

        # move the cursor
        self.__cursor = position

        # constrain the cursor to the buffer's virtual boundaries
        if self.__cursor < 0:
            self.__cursor = 0
        elif self.__cursor > len(self):
            self.__cursor = len(self)

        # mark the cursor as 'dirty' if it has moved
        self.__cursor_dirty = self.__cursor != old_cursor

    @property
    def __gap_len(self):
        """Get the length of the current gap."""
        return self.__gap_end - self.__gap_start

    @property
    def __left_len(self):
        """Get the length of the text to the left of the gap's start."""
        return self.__gap_start

    @property
    def __right_len(self):
        """Get the length of the text to the right of the gap's end."""
        return self.__content_end - self.__gap_end

    #
    # NOTE: together, __len__ and __getitem__ make for iteration!
    #

    def __len__(self):
        """Get the length of the buffer's text."""
        return self.__left_len + self.__right_len

    def __getitem__(self, index):
        """Get the character or slice at some index."""

        # get a slice rather than an index if necessary
        if isinstance(index, slice):
            # unpack 'indices()' into xrange, then pass generator to new Buffer
            return u''.join(self[i] for i in xrange(*index.indices(len(self))))
        else:
            # constrain index bounds
            if index >= len(self):
                raise IndexError(self.__class__.__name__ + " index out of range")

            # if before the gap, access buffer directly, else account for gap
            i = index if index < self.__gap_start else index + self.__gap_len

            # access the buffer
            return self.__buf[i]

    def __setitem__(self, index, value):
        """Set the character or slice at some index."""

        # set a slice rather than an index if necessary
        if isinstance(index, slice):
            # if we can get the length of the value, attempt a direct replace
            slice_len = abs(index.stop - index.start)
            if hasattr(value, "__len__") and len(value) == slice_len:
                # do a direct replacement
                for i in xrange(*index.indices(len(self))):
                    self[i] = value[i]
            else:
                # store the cursor position
                old_cursor = self.cursor

                # delete old slice and insert the new one
                self.cursor = index.start
                self.delete(index.stop - index.start)
                self.insert(value)

                # restore the cursor position while accounting for size change
                self.cursor = min(old_cursor, len(self))
        else:
            if index >= len(self):
                raise IndexError(self.__class__.__name__ + " index out of range")

            i = index if index < self.__gap_start else index + self.__gap_len
            self.__buf[i] = value

        """
        Insert text before the cursor's current position. Moves the cursor to
        the last character of the inserted text.

        'text' must be return unicode characters during iteration, or a
        TypeError will be raised.
        """

        # ensure the gap is large enough for the text being inserted
        self.__resize_gap(len(text))

        # move the gap to the cursor
        self.__move_gap()

        # insert text at left gap edge while moving the gap start forward
        for c in text:
            self.__buf[self.__gap_start] = c
            self.__gap_start += 1

        # move the cursor to match the new gap start and mark it as 'clean'
        self.__cursor = self.__gap_start
        self.__cursor_dirty = False

        return self

    def delete(self, length):
        """
        Remove text starting with the character at the cursor's current position
        and ending after 'length' characters have been removed.

        If 'length' is larger than the size of the buffer's remaining content,
        the text from the cursor to the end of the buffer is removed.
        """

        if length < 0:
            raise ValueError("length must be greater than or equal to 0")

        # don't do anything if nothing is to be deleted
        if length != 0:
            # move the gap to the cursor
            self.__move_gap()

            # increase the size of the gap to consume the deleted characters
            self.__gap_end += length

        return self

    def debug_view(self):
        """
        Get a debug view of the buffer's contents and internal values as a
        unicode string.
        """

        # write special values into gap
        for i in xrange(self.__gap_start, self.__gap_end):
            self.__buf[i] = u"_"

        # write special values into blank area
        for i in xrange(self.__content_end, len(self.__buf)):
            self.__buf[i] = u"#"

        # our desired display characters and their positions
        chars = [
            ("s", self.__gap_start),
            ("e", self.__gap_end),
            ("^", self.__cursor),
            ("$", self.__content_end)
        ]

        # find the left-most value we'll be displaying
        max_pos = max(map(lambda t: t[1], chars))

        # track all the rows we'll need
        rows = []

        # add the first row
        rows.append([None] * (max_pos + 1))

        # insert all the characters into rows
        for char, pos in chars:
            # try all the rows in turn until an empty slot is found
            for i in xrange(len(chars)):
                # add more space if we need it
                if len(rows) == i:
                    rows.append([None] * (max_pos + 1))

                # fill the slot if it was empty, then move on
                if rows[i][pos] is None:
                    rows[i][pos] = char
                    break

        # build the final string
        s = [u"'", unicode(self), u"', ", unicode(len(self)), u"\n",
                self.__buf.tounicode(), u"\n"]

        for row in rows:
            for c in row:
                s.append(c if c is not None else u" ")
            s.append(u'\n')

        return u''.join(s)

    def __resize_gap(self, target_size):
        """Ensure that the gap is at least as large as some target."""

        if self.__gap_len < target_size:

            # calculate size increase of the gap
            gap_delta = target_size + self.__min_gap_len - self.__gap_len

            # make room for the current content and the new gap
            self.__resize_buf(len(self) + gap_delta)

            # shift the right content down to make room for the new gap
            for i in reversed(xrange(self.__gap_end, self.__content_end)):
                self.__buf[i + gap_delta] = self.__buf[i]

            # move the gap and content end pointers forward
            self.__gap_end += gap_delta
            self.__content_end += gap_delta

        return self

    def __resize_buf(self, target_size):
        """Ensure that the buffer is at least as large as some target."""

        # double the size while the buffer is shorter than the target length
        while len(self.__buf) < target_size:
            self.__buf.extend(self.__buf)

        return self

    def __move_gap(self):
        """
        Move the gap to the current cursor position if the cursor is dirty.
        """

        if self.__cursor_dirty:
            # move the gap left as far as necessary
            while self.__gap_start > self.__cursor:
                # slide the gap to the left
                self.__gap_start -= 1
                self.__gap_end -= 1

                # copy the gap's former preceding character to the gap's old
                # final slot.
                self.__buf[self.__gap_end] = self.__buf[self.__gap_start]

            # move the gap right as far as necessary
            while self.__gap_start < self.__cursor:
                # copy the gap's following character to the gap's first slot.
                self.__buf[self.__gap_start] = self.__buf[self.__gap_end]

                # slide the gap to the right
                self.__gap_start += 1
                self.__gap_end += 1

            # the gap start now matches the cursor, so mark it 'clean'
            self.__cursor_dirty = False

        return self

    def __unicode__(self):
        """Get the buffer's entire text."""
        return self[:]

    def __str__(self):
        return unicode(self)

    def __repr__(self):
        return unicode(self.__class__.__name__ + "(" + repr(self[:]) + ")")

if __name__ == "__main__":
    b = Buffer()
    print b.debug_view()

    b.insert(u"Hello!")
    print b.debug_view()

    b.cursor -= 1
    print b.debug_view()

    b.insert(u", world")
    print b.debug_view()

    b.cursor -= 5
    print b.debug_view()

    b.delete(5)
    print b.debug_view()

    b.insert(u"pants")
    print b.debug_view()

    b.cursor = 0
    print b.debug_view()

    b.insert(u"Whoa! ")
    print b.debug_view()
<commit_msg>Add 'position' param to insert/delete<commit_after>#!/usr/bin/env python
import array

class Buffer(object):
    """
    Represents unicode text using a gap buffer. Can be initialized with any
    iterable item, as long as that item yields unicode characters during
    iteration.
    """

    def __init__(self, initial_content=None):

        # points to first empty space after left text
        self.__gap_start = 0

        # points to first taken space at beginning of right text
        self.__gap_end = 0

        # where the gap should be moved to when it needs to be moved
        self.__cursor = 0

        # whether the cursor has been moved without the gap being moved
        self.__cursor_dirty = False

        # the minimum amount of space to create when resizing the gap
        self.__min_gap_len = 10

        # the end of content in the buffer
        self.__content_end = 0

        # NOTE: we set the content end individually later so we can allow
        # initial_content objects that have no __len__ method.

        # initialize the buffer we use to hold text data in, a unicode array
        if initial_content is None:
            # insert placeholder content if no initial content was specified
            self.__buf = array.array("u", u" " * 10)
            self.__content_end = 0
        else:
            # insert the initial content and set the content end pointer for it
            self.__buf = array.array("u", initial_content)
            self.__content_end = len(self.__buf)

    @property
    def cursor(self):
        """Get the current position of the cursor."""
        return self.__cursor

    @cursor.setter
    def cursor(self, position):
        """Move the cursor within the buffer."""

        # NOTE: does not move the gap! when moved, the cursor position is marked
        # as 'dirty' so that when text changes happen, the gap can be moved
        # accordingly elsewhere.

        # store the former cursor position
        old_cursor = self.__cursor

        # move the cursor
        self.__cursor = position

        # constrain the cursor to the buffer's virtual boundaries
        if self.__cursor < 0:
            self.__cursor = 0
        elif self.__cursor > len(self):
            self.__cursor = len(self)

        # mark the cursor as 'dirty' if it has moved
        self.__cursor_dirty = self.__cursor != old_cursor

    @property
    def __gap_len(self):
        """Get the length of the current gap."""
        return self.__gap_end - self.__gap_start

    @property
    def __left_len(self):
        """Get the length of the text to the left of the gap's start."""
        return self.__gap_start

    @property
    def __right_len(self):
        """Get the length of the text to the right of the gap's end."""
        return self.__content_end - self.__gap_end

    #
    # NOTE: together, __len__ and __getitem__ make for iteration!
    #

    def __len__(self):
        """Get the length of the buffer's text."""
        return self.__left_len + self.__right_len

    def __getitem__(self, index):
        """Get the character or slice at some index."""

        # get a slice rather than an index if necessary
        if isinstance(index, slice):
            # unpack 'indices()' into xrange, then pass generator to new Buffer
            return u''.join(self[i] for i in xrange(*index.indices(len(self))))
        else:
            # constrain index bounds
            if index >= len(self):
                raise IndexError(self.__class__.__name__ + " index out of range")

            # if before the gap, access buffer directly, else account for gap
            i = index if index < self.__gap_start else index + self.__gap_len

            # access the buffer
            return self.__buf[i]

    def __setitem__(self, index, value):
        """Set the character or slice at some index."""

        # set a slice rather than an index if necessary
        if isinstance(index, slice):
            # if we can get the length of the value, attempt a direct replace
            slice_len = abs(index.stop - index.start)
            if hasattr(value, "__len__") and len(value) == slice_len:
                # do a direct replacement
                for i in xrange(*index.indices(len(self))):
                    self[i] = value[i]
            else:
                # store the cursor position
                old_cursor = self.cursor

                # delete old slice and insert the new one
                self.cursor = index.start
                self.delete(index.stop - index.start)
                self.insert(value)

                # restore the cursor position while accounting for size change
                self.cursor = min(old_cursor, len(self))
        else:
            if index >= len(self):
                raise IndexError(self.__class__.__name__ + " index out of range")

            i = index if index < self.__gap_start else index + self.__gap_len
            self.__buf[i] = value

    def insert(self, text, position=None):
        """
        Insert text before the cursor's current position.

        If 'position' is specified, moves the cursor to the given position
        before inserting the text.

        Moves the cursor to the last character of the inserted text.

        'text' must be return unicode characters during iteration, or a
        TypeError will be raised.
        """

        # first move the cursor to the position if a position was specified
        if position is not None:
            self.cursor = position

        # move the gap to the cursor
        self.__move_gap()

        # ensure the gap is large enough for the text being inserted
        self.__resize_gap(len(text))

        # insert text at left gap edge while moving the gap start forward
        for c in text:
            self.__buf[self.__gap_start] = c
            self.__gap_start += 1

        # move the cursor to match the new gap start and mark it as 'clean'
        self.__cursor = self.__gap_start
        self.__cursor_dirty = False

        return self

    def delete(self, length, position=None):
        """
        Remove text starting with the character at the cursor's current position
        and ending after 'length' characters have been removed.

        If 'position' is specified, moves the cursor to the given position
        before deleting the text.

        If 'length' is larger than the size of the buffer's remaining content,
        the text from the cursor to the end of the buffer is removed.
        """

        if position is not None:
            self.cursor = position

        if length < 0:
            raise ValueError("length must be greater than or equal to 0")

        # don't do anything if nothing is to be deleted
        if length != 0:
            # move the gap to the cursor
            self.__move_gap()

            # increase the size of the gap to consume the deleted characters,
            # but only up to the size of the internal buffer.
            self.__gap_end += min(length, len(self.__buf))

        return self

    def debug_view(self):
        """
        Get a debug view of the buffer's contents and internal values as a
        unicode string.
        """

        # write special values into gap
        for i in xrange(self.__gap_start, self.__gap_end):
            self.__buf[i] = u"_"

        # write special values into blank area
        for i in xrange(self.__content_end, len(self.__buf)):
            self.__buf[i] = u"#"

        # our desired display characters and their positions
        chars = [
            ("s", self.__gap_start),
            ("e", self.__gap_end),
            ("^", self.__cursor),
            ("$", self.__content_end)
        ]

        # find the left-most value we'll be displaying
        max_pos = max(map(lambda t: t[1], chars))

        # track all the rows we'll need
        rows = []

        # add the first row
        rows.append([None] * (max_pos + 1))

        # insert all the characters into rows
        for char, pos in chars:
            # try all the rows in turn until an empty slot is found
            for i in xrange(len(chars)):
                # add more space if we need it
                if len(rows) == i:
                    rows.append([None] * (max_pos + 1))

                # fill the slot if it was empty, then move on
                if rows[i][pos] is None:
                    rows[i][pos] = char
                    break

        # build the final string
        s = [u"'", unicode(self), u"', ", unicode(len(self)), u"\n",
                self.__buf.tounicode(), u"\n"]

        for row in rows:
            for c in row:
                s.append(c if c is not None else u" ")
            s.append(u'\n')

        return u''.join(s)

    def __resize_gap(self, target_size):
        """Ensure that the gap is at least as large as some target."""

        if self.__gap_len < target_size:

            # calculate size increase of the gap
            gap_delta = target_size + self.__min_gap_len - self.__gap_len

            # make room for the current content and the new gap
            self.__resize_buf(len(self) + gap_delta)

            # shift the right content down to make room for the new gap
            for i in reversed(xrange(self.__gap_end, self.__content_end)):
                self.__buf[i + gap_delta] = self.__buf[i]

            # move the gap and content end pointers forward
            self.__gap_end += gap_delta
            self.__content_end += gap_delta

        return self

    def __resize_buf(self, target_size):
        """Ensure that the buffer is at least as large as some target."""

        # double the size while the buffer is shorter than the target length
        while len(self.__buf) < target_size:
            self.__buf.extend(self.__buf)

        return self

    def __move_gap(self):
        """
        Move the gap to the current cursor position if the cursor is dirty.
        """

        if self.__cursor_dirty:
            # move the gap left as far as necessary
            while self.__gap_start > self.__cursor:
                # slide the gap to the left
                self.__gap_start -= 1
                self.__gap_end -= 1

                # copy the gap's former preceding character to the gap's old
                # final slot.
                self.__buf[self.__gap_end] = self.__buf[self.__gap_start]

            # move the gap right as far as necessary
            while self.__gap_start < self.__cursor:
                # copy the gap's following character to the gap's first slot.
                self.__buf[self.__gap_start] = self.__buf[self.__gap_end]

                # slide the gap to the right
                self.__gap_start += 1
                self.__gap_end += 1

            # the gap start now matches the cursor, so mark it 'clean'
            self.__cursor_dirty = False

        return self

    def __unicode__(self):
        """Get the buffer's entire text."""
        return self[:]

    def __str__(self):
        return unicode(self)

    def __repr__(self):
        return unicode(self.__class__.__name__ + "(" + repr(self[:]) + ")")

if __name__ == "__main__":
    b = Buffer()
    print b.debug_view()

    b.insert(u"Hello!")
    print b.debug_view()

    b.cursor -= 1
    print b.debug_view()

    b.insert(u", world")
    print b.debug_view()

    b.cursor -= 5
    print b.debug_view()

    b.delete(5)
    print b.debug_view()

    b.insert(u"pants")
    print b.debug_view()

    b.cursor = 0
    print b.debug_view()

    b.insert(u"Whoa! ")
    print b.debug_view()
<|endoftext|>