<commit_before>from argparse import ArgumentParser
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor, wait
from datetime import datetime, timezone
from math import log
from os import path, walk
import re
import sys
import yaml
from yaml.nodes import MappingNode

import digestive
from digestive.entropy import Entropy
from digestive.ewf import EWFSource, supported_exts as ewf_supported_formats
from digestive.hash import MD5, SHA1, SHA256, SHA512, sha3_enabled, SHA3256, SHA3512
from digestive.io import Source

# binary suffixes for byte sizes
_sizes = ['bytes', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB']
# the corresponding suffixes
_suffixes = [size[0].lower() for size in _sizes]
# multipliers for byte size arguments taken from _suffixes
_multipliers = {suffix: 1 << (i * 10) for i, suffix in enumerate(_suffixes)}

def file_size(size, template='{value:.4g} {unit}'):
    """
    Converts a byte size into a human-readable format using binary suffixes.

    :param size: The value to be formatted.
    :param template: A format string with two parameters: value and unit.
    :return: A human-readable file size.
    """
    order = int(log(size, 2) // 10) if size else 0
    if order >= len(_sizes):
        # exceeding ludicrous file sizes, default to bytes
        return '{} bytes'.format(size)
    return template.format(value=size / (1 << (order * 10)), unit=_sizes[order])

def num_bytes(size):
    """
    Converts a human-readable byte size into an int, stripping an optional binary suffix.

    :param size: The size to be parsed.
    :return: An int.
    """
    match = re.match(r'(?P<size>\d+)(?P<suffix>[{}])?$'.format(''.join(_suffixes)), size, re.IGNORECASE)
    if match:
        # one optional group, no suffix is 'b' for bytes
        size, suffix = match.groups('b')
        return int(size) * _multipliers.get(suffix.lower())
    else:
        raise TypeError(size)

def parse_arguments(arguments=None):
    """
    Parses commandline arguments defining both program options and input and output.

    :param arguments: The arguments to parse, or None. Arguments will be read from sys.argv if None.
    :return: An argparse.Namespace object.
    """
    parser = ArgumentParser(description='run multiple digests on files')
    # hash digest sinks
    parser.add_argument('-m', '--md5', action='append_const', dest='sinks', const=MD5,
                        help='calculate MD5 hash')
    parser.add_argument('-1', '--sha1', action='append_const', dest='sinks', const=SHA1,
                        help='calculate SHA-1 hash')
    parser.add_argument('-2', '--sha256', action='append_const', dest='sinks', const=SHA256,
                        help='calculate SHA-256 hash')
    parser.add_argument('-5', '--sha512', action='append_const', dest='sinks', const=SHA512,
                        help='calculate SHA-512 hash')
    hashes = [MD5, SHA1, SHA256, SHA512]
    hashes_help = 'calculate MD5, SHA-1, SHA-256 and SHA-512 hashes (equivalent to -m125)'

    if sha3_enabled:
        parser.add_argument('-3', '--sha3-256', action='append_const', dest='sinks', const=SHA3256,
                            help='calculate SHA3-256 hash')
        parser.add_argument('--sha3-512', action='append_const', dest='sinks', const=SHA3512,
                            help='calculate SHA3-512 hash')
        hashes.append(SHA3256)
        hashes_help = 'calculate MD5, SHA-1, SHA-256, SHA-512 and SHA3-256 hashes (equivalent to -m1253)'

    # convenience switch to include all hashes
    parser.add_argument('--hashes', action='store_const', dest='sinks', const=hashes,
                        help=hashes_help)
    # entropy sink
    parser.add_argument('-e', '--entropy', action='append_const', dest='sinks', const=Entropy,
                        help='calculate binary entropy')
    # misc options
    parser.add_argument('-j', '--jobs', type=int, metavar='JOBS',
                        help='use up to %(metavar)s threads to process digests (defaults to the number of digests)')
    parser.add_argument('-b', '--block-size', type=num_bytes, metavar='BYTES', default='1M',
                        help='read data in chunks of %(metavar)s at a time (defaults to 1M)')
    # TODO: specifying --format ewf on files that don't match supported exts will raise ValueError
    parser.add_argument('-f', '--format', choices=('auto', 'raw', 'ewf'), default='auto',
                        help='specify source format (defaults to auto)')
    # TODO: add choice (and implement) throughput / speed, eta
    parser.add_argument('-p', '--progress', choices=('bytes',), default='bytes',
                        help='show progress information (defaults to bytes)')
    parser.add_argument('-P', '--no-progress', action='store_false', dest='progress',
                        help='disable progress output (always disabled for piped output)')
    # TODO: -t, --time
    parser.add_argument('-r', '--recursive', action='store_true',
                        help='process sources recursively')
    parser.add_argument('-o', '--output',
                        help='write yaml-encoded output to file')
    # positional arguments: sources
    parser.add_argument('sources', metavar='FILE', nargs='+',
                        help='input files')

    arguments = parser.parse_args(arguments)
    process_arguments(arguments, parser)

    return arguments

def process_arguments(arguments, parser):
    """
    Post-processes parsed arguments.

    :param arguments: The arguments to be processed (an argparse.Namespace object).
    :param parser: The parser used to parse the arguments (used for error reporting).
    :return: The passed arguments, adjusted where needed.
    """
    if not arguments.sinks:
        parser.error('at least one sink is required')

    arguments.jobs = arguments.jobs if arguments.jobs else len(arguments.sinks)

def output_to_file(output):
    """
    Collector of dicts to be written to output file as explicit YAML
    documents.

    :param output: The file name to write to (or None).
    :return: A collecting generator (ignoring input if output was None).
    """
    if output:
        # create a generator within a text-io context manager for output
        with open(output, 'w') as stream:
            info = yield
            # use no explicit document start for the info leader
            yaml.dump(info, stream=stream)

            while True:
                # receive source name and sink results
                value = yield
                # dump value to output, creating an explicit document start
                yaml.dump(value, stream=stream, explicit_start=True)
    else:
        while True:
            # do nothing with any value received
            _ = yield  # noqa: variable _ is assigned to explicitly to make clear this is a collecting yield

def process_source(executor, source, sinks, block_size=1 << 20, progress=None):
    """
    Processes a data source, feeding chunks of at most block_size to each sink in parallel.

    :param executor: The executor to submit execution jobs to.
    :param source: The data source to read from.
    :param sinks: The sink instances to process data chunks with.
    :param block_size: The maximum chunk size to read.
    :param progress: An object to call add(int) on for every read block.
    """
    generator = source.blocks(block_size)
    block = next(generator, False)
    while block:
        if callable(progress):
            progress(len(block))
        futures = [executor.submit(sink.process, block) for sink in sinks]
        block = next(generator, False)
        wait(futures)

def files(sources, recurse=False, followlinks=False):
    """
    Generates paths to files.

    :param sources: The base sources passed as arguments.
    :param recurse: Whether to recurse into directories.
    :param followlinks: Whether to follow symbolic links.
    :return: A generator of sources based on the provided arguments.
    """
    if recurse:
        for source in sources:
            # use walk to recurse into each source…
            for base, _, names in walk(source, followlinks=followlinks):
                # …and yield all the files within them
                # TODO: some intelligence might be nice here, like removing .E02 when .E01 is found
                yield from (path.join(base, name) for name in names)
    else:
        # simply use sources
        yield from sources

def get_source(file, source_type='auto'):
    if source_type == 'raw':
        return Source(file)
    elif source_type == 'ewf':
        return EWFSource(file)
    elif source_type == 'auto':
        basename, ext = path.splitext(file)
        source_type = 'ewf' if ext.lower() in ewf_supported_formats else 'raw'
        return get_source(file, source_type=source_type)
    else:
        raise ValueError('unknown source type: {}'.format(source_type))

def main(arguments=None):
    """
    Runs digestive.

    :param arguments: Commandline arguments, passed to parse_arguments.
    """

    class Progress:
        """
        FIXME: I'm a bit of a hack...
        """

        def __init__(self, end):
            self.value = 0
            self.end = end

        def __call__(self, amount):
            self.value += amount
            if self.end:
                # progress only makes sense if end > 0
                print_progress(self.value, self.end)

    def print_progress(current, end):
        print('\r  {percent:>4.0%} [{bar:<20}] ({value})'.format(
            percent=(current / end),
            bar=('»' * int((20 * current / end))),  # TODO: will » work everywhere?
            value=file_size(current, '{value:>8.3f} {unit}')
        ), end='')

    # add an order-retaining representer for OrderedDict (default impl calls sorted())
    yaml.add_representer(OrderedDict, lambda dumper, data: MappingNode(
        'tag:yaml.org,2002:map',
        # create a represented list of tuples, in order or data.items()
        [(dumper.represent_data(key), dumper.represent_data(value)) for key, value in data.items()]
    ))

    arguments = parse_arguments(arguments)
    # create the output generator
    output = output_to_file(arguments.output)
    # initialize output (moves it to the first occurrence of yield)
    next(output)
    info = OrderedDict()
    info['digestive'] = str(digestive.__version__)
    info['started'] = datetime.now(tz=timezone.utc)
    output.send(info)

    with ThreadPoolExecutor(arguments.jobs) as executor:
        # TODO: globs like tests/files/file.* includes both file.E01 and file.E02
        # TODO: only file.E01 will get treated as ewf, file.E02 should be removed from sources
        for file in files(arguments.sources, arguments.recursive):
            with get_source(file, arguments.format) as source:
                # instantiate sinks from requested types
                sinks = [sink() for sink in arguments.sinks]
                # flush initial status line to force it to show in something like | less
                print('{} ({})'.format(source, file_size(len(source))), flush=True)

                if arguments.progress and sys.stdout.isatty():
                    # stdout should support carriage returns, engage progress information!
                    process_source(executor, source, sinks, arguments.block_size, progress=Progress(end=len(source)))
                    # use terminal escape to clear line and \r return cursor to start of line
                    print('\033[2K\r', end='')
                else:
                    # omit progress altogether
                    process_source(executor, source, sinks, arguments.block_size)

                results = OrderedDict((sink.name, sink.result()) for sink in sinks)
                for name, result in results.items():
                    print('  {:<12} {}'.format(name, result))

                # create meta data leader
                # TODO: using kwargs here would be nice, but that destroys order :( (see PEP-468)
                info = OrderedDict((
                    ('source', file),
                    ('size', len(source)),
                    ('completed', datetime.now(tz=timezone.utc))
                ))
                # add results
                info.update(results)
                # send info to the output collector
                output.send(info)

    # close the output collector, which in turn closes the output stream
    output.close()

if __name__ == '__main__':
    main()
<commit_msg>Include explicit doc start for info leader<commit_after>from argparse import ArgumentParser
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor, wait
from datetime import datetime, timezone
from math import log
from os import path, walk
import re
import sys
import yaml
from yaml.nodes import MappingNode

import digestive
from digestive.entropy import Entropy
from digestive.ewf import EWFSource, supported_exts as ewf_supported_formats
from digestive.hash import MD5, SHA1, SHA256, SHA512, sha3_enabled, SHA3256, SHA3512
from digestive.io import Source

# binary suffixes for byte sizes
_sizes = ['bytes', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB']
# the corresponding suffixes
_suffixes = [size[0].lower() for size in _sizes]
# multipliers for byte size arguments taken from _suffixes
_multipliers = {suffix: 1 << (i * 10) for i, suffix in enumerate(_suffixes)}

def file_size(size, template='{value:.4g} {unit}'):
    """
    Converts a byte size into a human-readable format using binary suffixes.

    :param size: The value to be formatted.
    :param template: A format string with two parameters: value and unit.
    :return: A human-readable file size.
    """
    order = int(log(size, 2) // 10) if size else 0
    if order >= len(_sizes):
        # exceeding ludicrous file sizes, default to bytes
        return '{} bytes'.format(size)
    return template.format(value=size / (1 << (order * 10)), unit=_sizes[order])

def num_bytes(size):
    """
    Converts a human-readable byte size into an int, stripping an optional binary suffix.

    :param size: The size to be parsed.
    :return: An int.
    """
    match = re.match(r'(?P<size>\d+)(?P<suffix>[{}])?$'.format(''.join(_suffixes)), size, re.IGNORECASE)
    if match:
        # one optional group, no suffix is 'b' for bytes
        size, suffix = match.groups('b')
        return int(size) * _multipliers.get(suffix.lower())
    else:
        raise TypeError(size)

def parse_arguments(arguments=None):
    """
    Parses commandline arguments defining both program options and input and output.

    :param arguments: The arguments to parse, or None. Arguments will be read from sys.argv if None.
    :return: An argparse.Namespace object.
    """
    parser = ArgumentParser(description='run multiple digests on files')
    # hash digest sinks
    parser.add_argument('-m', '--md5', action='append_const', dest='sinks', const=MD5,
                        help='calculate MD5 hash')
    parser.add_argument('-1', '--sha1', action='append_const', dest='sinks', const=SHA1,
                        help='calculate SHA-1 hash')
    parser.add_argument('-2', '--sha256', action='append_const', dest='sinks', const=SHA256,
                        help='calculate SHA-256 hash')
    parser.add_argument('-5', '--sha512', action='append_const', dest='sinks', const=SHA512,
                        help='calculate SHA-512 hash')
    hashes = [MD5, SHA1, SHA256, SHA512]
    hashes_help = 'calculate MD5, SHA-1, SHA-256 and SHA-512 hashes (equivalent to -m125)'

    if sha3_enabled:
        parser.add_argument('-3', '--sha3-256', action='append_const', dest='sinks', const=SHA3256,
                            help='calculate SHA3-256 hash')
        parser.add_argument('--sha3-512', action='append_const', dest='sinks', const=SHA3512,
                            help='calculate SHA3-512 hash')
        hashes.append(SHA3256)
        hashes_help = 'calculate MD5, SHA-1, SHA-256, SHA-512 and SHA3-256 hashes (equivalent to -m1253)'

    # convenience switch to include all hashes
    parser.add_argument('--hashes', action='store_const', dest='sinks', const=hashes,
                        help=hashes_help)
    # entropy sink
    parser.add_argument('-e', '--entropy', action='append_const', dest='sinks', const=Entropy,
                        help='calculate binary entropy')
    # misc options
    parser.add_argument('-j', '--jobs', type=int, metavar='JOBS',
                        help='use up to %(metavar)s threads to process digests (defaults to the number of digests)')
    parser.add_argument('-b', '--block-size', type=num_bytes, metavar='BYTES', default='1M',
                        help='read data in chunks of %(metavar)s at a time (defaults to 1M)')
    # TODO: specifying --format ewf on files that don't match supported exts will raise ValueError
    parser.add_argument('-f', '--format', choices=('auto', 'raw', 'ewf'), default='auto',
                        help='specify source format (defaults to auto)')
    # TODO: add choice (and implement) throughput / speed, eta
    parser.add_argument('-p', '--progress', choices=('bytes',), default='bytes',
                        help='show progress information (defaults to bytes)')
    parser.add_argument('-P', '--no-progress', action='store_false', dest='progress',
                        help='disable progress output (always disabled for piped output)')
    # TODO: -t, --time
    parser.add_argument('-r', '--recursive', action='store_true',
                        help='process sources recursively')
    parser.add_argument('-o', '--output',
                        help='write yaml-encoded output to file')
    # positional arguments: sources
    parser.add_argument('sources', metavar='FILE', nargs='+',
                        help='input files')

    arguments = parser.parse_args(arguments)
    process_arguments(arguments, parser)

    return arguments

def process_arguments(arguments, parser):
    """
    Post-processes parsed arguments.

    :param arguments: The arguments to be processed (an argparse.Namespace object).
    :param parser: The parser used to parse the arguments (used for error reporting).
    :return: The passed arguments, adjusted where needed.
    """
    if not arguments.sinks:
        parser.error('at least one sink is required')

    arguments.jobs = arguments.jobs if arguments.jobs else len(arguments.sinks)

def output_to_file(output):
    """
    Collector of dicts to be written to output file as explicit YAML
    documents.

    :param output: The file name to write to (or None).
    :return: A collecting generator (ignoring input if output was None).
    """
    if output:
        # create a generator within a text-io context manager for output
        with open(output, 'w') as stream:
            while True:
                # receive source name and sink results
                value = yield
                # dump value to output, creating an explicit document start
                yaml.dump(value, stream=stream, explicit_start=True)
    else:
        while True:
            # do nothing with any value received
            _ = yield  # noqa: variable _ is assigned to explicitly to make clear this is a collecting yield

def process_source(executor, source, sinks, block_size=1 << 20, progress=None):
    """
    Processes a data source, feeding chunks of at most block_size to each sink in parallel.

    :param executor: The executor to submit execution jobs to.
    :param source: The data source to read from.
    :param sinks: The sink instances to process data chunks with.
    :param block_size: The maximum chunk size to read.
    :param progress: An object to call add(int) on for every read block.
    """
    generator = source.blocks(block_size)
    block = next(generator, False)
    while block:
        if callable(progress):
            progress(len(block))
        futures = [executor.submit(sink.process, block) for sink in sinks]
        block = next(generator, False)
        wait(futures)

def files(sources, recurse=False, followlinks=False):
    """
    Generates paths to files.

    :param sources: The base sources passed as arguments.
    :param recurse: Whether to recurse into directories.
    :param followlinks: Whether to follow symbolic links.
    :return: A generator of sources based on the provided arguments.
    """
    if recurse:
        for source in sources:
            # use walk to recurse into each source…
            for base, _, names in walk(source, followlinks=followlinks):
                # …and yield all the files within them
                # TODO: some intelligence might be nice here, like removing .E02 when .E01 is found
                yield from (path.join(base, name) for name in names)
    else:
        # simply use sources
        yield from sources

def get_source(file, source_type='auto'):
    if source_type == 'raw':
        return Source(file)
    elif source_type == 'ewf':
        return EWFSource(file)
    elif source_type == 'auto':
        basename, ext = path.splitext(file)
        source_type = 'ewf' if ext.lower() in ewf_supported_formats else 'raw'
        return get_source(file, source_type=source_type)
    else:
        raise ValueError('unknown source type: {}'.format(source_type))

def main(arguments=None):
    """
    Runs digestive.

    :param arguments: Commandline arguments, passed to parse_arguments.
    """

    class Progress:
        """
        FIXME: I'm a bit of a hack...
        """

        def __init__(self, end):
            self.value = 0
            self.end = end

        def __call__(self, amount):
            self.value += amount
            if self.end:
                # progress only makes sense if end > 0
                print_progress(self.value, self.end)

    def print_progress(current, end):
        print('\r  {percent:>4.0%} [{bar:<20}] ({value})'.format(
            percent=(current / end),
            bar=('»' * int((20 * current / end))),  # TODO: will » work everywhere?
            value=file_size(current, '{value:>8.3f} {unit}')
        ), end='')

    # add an order-retaining representer for OrderedDict (default impl calls sorted())
    yaml.add_representer(OrderedDict, lambda dumper, data: MappingNode(
        'tag:yaml.org,2002:map',
        # create a represented list of tuples, in order or data.items()
        [(dumper.represent_data(key), dumper.represent_data(value)) for key, value in data.items()]
    ))

    arguments = parse_arguments(arguments)
    # create the output generator
    output = output_to_file(arguments.output)
    # initialize output (moves it to the first occurrence of yield)
    next(output)
    info = OrderedDict()
    info['digestive'] = str(digestive.__version__)
    info['started'] = datetime.now(tz=timezone.utc)
    output.send(info)

    with ThreadPoolExecutor(arguments.jobs) as executor:
        # TODO: globs like tests/files/file.* includes both file.E01 and file.E02
        # TODO: only file.E01 will get treated as ewf, file.E02 should be removed from sources
        for file in files(arguments.sources, arguments.recursive):
            with get_source(file, arguments.format) as source:
                # instantiate sinks from requested types
                sinks = [sink() for sink in arguments.sinks]
                # flush initial status line to force it to show in something like | less
                print('{} ({})'.format(source, file_size(len(source))), flush=True)

                if arguments.progress and sys.stdout.isatty():
                    # stdout should support carriage returns, engage progress information!
                    process_source(executor, source, sinks, arguments.block_size, progress=Progress(end=len(source)))
                    # use terminal escape to clear line and \r return cursor to start of line
                    print('\033[2K\r', end='')
                else:
                    # omit progress altogether
                    process_source(executor, source, sinks, arguments.block_size)

                results = OrderedDict((sink.name, sink.result()) for sink in sinks)
                for name, result in results.items():
                    print('  {:<12} {}'.format(name, result))

                # create meta data leader
                # TODO: using kwargs here would be nice, but that destroys order :( (see PEP-468)
                info = OrderedDict((
                    ('source', file),
                    ('size', len(source)),
                    ('completed', datetime.now(tz=timezone.utc))
                ))
                # add results
                info.update(results)
                # send info to the output collector
                output.send(info)

    # close the output collector, which in turn closes the output stream
    output.close()

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


<commit_before>import discord, threading, sys
from plugins import Plugin
from config import version
import logging

log = logging.getLogger("discord")

class DiscoBot(discord.Client):

	def connect(self, email, password):
		log.debug("Email: " + email)
		self.login(email, password)

		if not self.is_logged_in:
			print "Login failed."
			log.error("Login failed.")
			sys.exit(1)

		self.t = threading.Thread(target=self.run)
		self.t.daemon = True
		self.t.start()

		log.info("Thread 'DiscoBot' started.")

		while self.is_logged_in:
			self.defcon()

	def defcon(self):
		print "command(s): (exit|e) (leave|l) (join <inv>)"
		cmd = raw_input("> ")
		if cmd == "e" or cmd == "exit":
			log.info("Command: " + cmd)
			log.info("Exiting.")
			self.logout()
			sys.exit(1)
		elif cmd == "leave" or cmd == "l":
			for s in self.servers:
				print "Leaving server: " + s.name
				log.warning("Leaving server: " + s.name)
				self.leave_server(s)
		elif cmd.split(" ")[0] == "join":
			invite = cmd.split(" ")[1]
			print "Joining server."
			log.info("Joining server.")
			self.accept_invite(invite)

	def on_ready(self):
		print "Connected!"

		log.info(self.user.name + " successfully connected!")
		log.debug("Username: " + self.user.name)
		log.debug("ID: " + self.user.id)

		log.info("Plugins loaded(%i):" % len(Plugin.plugins))
		for plugin in Plugin.plugins:
			log.info("..." + plugin.title)

		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_ready"):
				plugin.on_ready(self)

	def on_message(self, message):
		if message.content.startswith("!list"):
			commands = []
			for plugin in Plugin.plugins:
				if hasattr(plugin, "commands"):
					commands += plugin.commands
			if len(commands) > 0:
				cmd = "\n".join(commands)
				self.send_message(message.channel, "I'm at your service.\n\n%s" % cmd)
		elif message.content.startswith("!plugins"):
			plugs = ""
			for plugin in Plugin.plugins:
				if hasattr(plugin, "title"):
					plugs += plugin.title + "\n"
			self.send_message(message.channel, "Here be my plugins, master.\n\n" + plugs)
		elif message.content.startswith("!help"):
			credit = "DiscoBot v%s by Syntox <https://github.com/Syntox32/DiscoBot>" % version
			self.send_message(message.channel, credit + "\n\nType `!list` to show all commands")

		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_message"):
				plugin.on_message(self, message)

	def on_message_delete(self, message):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_message_delete"):
				plugin.on_message_delete(self, message)

	def on_message_edit(self, before, after):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_message_edit"):
				plugin.on_message_edit(self, before,after)

	def on_status(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_status"):
				plugin.on_status(self, member)

	def on_channel_delete(self, channel):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_channel_delete"):
				plugin.on_channel_delete(self, channel)

	def on_channel_create(self, channel):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_channel_create"):
				plugin.on_channel_create(self, channel)

	def on_channel_update(self, channel):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_channel_update"):
				plugin.on_channel_update(self, channel)

	def on_member_join(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_member_join"):
				plugin.on_member_join(self, member)

	def on_member_remove(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_member_remove"):
				plugin.on_member_remove(self, member)

	def on_member_update(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_member_update"):
				plugin.on_member_update(self, member)

	def on_server_join(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_join"):
				plugin.on_server_join(self, server)

	def on_server_remove(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_remove"):
				plugin.on_server_remove(self, server)

	def on_server_role_create(self, server, role):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_role_create"):
				plugin.on_server_role_create(self, server,role)

	def on_server_role_delete(self, server, role):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_role_delete"):
				plugin.on_server_role_delete(self, server,role)

	def on_server_role_update(self, role):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_role_update"):
				plugin.on_server_role_update(self, role)

	def on_server_available(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_available"):
				plugin.on_server_available(self, server)

	def on_server_unavailable(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_unavailable"):
				plugin.on_server_unavailable(self, server)

	def on_voice_state_update(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_voice_state_update"):
				plugin.on_voice_state_update(self, member)<commit_msg>improved !list output<commit_after>import discord, threading, sys
from plugins import Plugin
from config import version
import logging

log = logging.getLogger("discord")

class DiscoBot(discord.Client):

	def connect(self, email, password):
		log.debug("Email: " + email)
		self.login(email, password)

		if not self.is_logged_in:
			print "Login failed."
			log.error("Login failed.")
			sys.exit(1)

		self.t = threading.Thread(target=self.run)
		self.t.daemon = True
		self.t.start()

		log.info("Thread 'DiscoBot' started.")

		while self.is_logged_in:
			self.defcon()

	def defcon(self):
		print "command(s): (exit|e) (leave|l) (join <inv>)"
		cmd = raw_input("> ")
		if cmd == "e" or cmd == "exit":
			log.info("Command: " + cmd)
			log.info("Exiting.")
			self.logout()
			sys.exit(1)
		elif cmd == "leave" or cmd == "l":
			for s in self.servers:
				print "Leaving server: " + s.name
				log.warning("Leaving server: " + s.name)
				self.leave_server(s)
		elif cmd.split(" ")[0] == "join":
			invite = cmd.split(" ")[1]
			print "Joining server."
			log.info("Joining server.")
			self.accept_invite(invite)

	def on_ready(self):
		print "Connected!"

		log.info(self.user.name + " successfully connected!")
		log.debug("Username: " + self.user.name)
		log.debug("ID: " + self.user.id)

		log.info("Plugins loaded(%i):" % len(Plugin.plugins))
		for plugin in Plugin.plugins:
			log.info("..." + plugin.title)

		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_ready"):
				plugin.on_ready(self)

	def on_message(self, message):
		if message.content.startswith("!list"):
			commands = []
			for plugin in Plugin.plugins:
				if hasattr(plugin, "commands"):
					for c in plugin.commands:
						s = c.split("::")
						cmd = s[0].strip()
						desc = ""
						if len(s) > 1:
							desc = " *%s*" % s[1].strip()
						commands.append("`%s`%s" % (cmd, desc))
			if len(commands) > 0:
				cmd = "\n".join(commands)
				self.send_message(message.channel, "I'm at your service.\n\n%s" % cmd)

		elif message.content.startswith("!plugins"):
			plugs = ""
			for plugin in Plugin.plugins:
				if hasattr(plugin, "title"):
					plugs += plugin.title + "\n"
			self.send_message(message.channel, "Here be my plugins, master.\n\n" + plugs)

		elif message.content.startswith("!help"):
			credit = "DiscoBot v%s by Syntox <https://github.com/Syntox32/DiscoBot>" % version
			self.send_message(message.channel, credit + "\n\nType `!list` to show all commands")

		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_message"):
				plugin.on_message(self, message)

	def on_message_delete(self, message):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_message_delete"):
				plugin.on_message_delete(self, message)

	def on_message_edit(self, before, after):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_message_edit"):
				plugin.on_message_edit(self, before,after)

	def on_status(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_status"):
				plugin.on_status(self, member)

	def on_channel_delete(self, channel):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_channel_delete"):
				plugin.on_channel_delete(self, channel)

	def on_channel_create(self, channel):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_channel_create"):
				plugin.on_channel_create(self, channel)

	def on_channel_update(self, channel):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_channel_update"):
				plugin.on_channel_update(self, channel)

	def on_member_join(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_member_join"):
				plugin.on_member_join(self, member)

	def on_member_remove(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_member_remove"):
				plugin.on_member_remove(self, member)

	def on_member_update(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_member_update"):
				plugin.on_member_update(self, member)

	def on_server_join(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_join"):
				plugin.on_server_join(self, server)

	def on_server_remove(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_remove"):
				plugin.on_server_remove(self, server)

	def on_server_role_create(self, server, role):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_role_create"):
				plugin.on_server_role_create(self, server,role)

	def on_server_role_delete(self, server, role):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_role_delete"):
				plugin.on_server_role_delete(self, server,role)

	def on_server_role_update(self, role):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_role_update"):
				plugin.on_server_role_update(self, role)

	def on_server_available(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_available"):
				plugin.on_server_available(self, server)

	def on_server_unavailable(self, server):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_server_unavailable"):
				plugin.on_server_unavailable(self, server)

	def on_voice_state_update(self, member):
		for plugin in Plugin.plugins:
			if hasattr(plugin, "on_voice_state_update"):
				plugin.on_voice_state_update(self, member)<|endoftext|>


<commit_before>#! /usr/bin/env python
import Levenshtein
import abjad
import discograph
import mongoengine
import re
from flask import (
    Flask,
    abort,
    jsonify,
    redirect,
    render_template,
    )
#from werkzeug.contrib.cache import MemcachedCache

app = Flask(__name__)
#cache = MemcachedCache(['127.0.0.1:11211'])
#cache.clear()
#cache = None
mongoengine.connect('discograph')

@app.route('/')
def route__index():
    return render_template(
        'index.html',
        artist=None,
        title='discoGraph',
        )

@app.route('/random')
def route__random():
    artist = discograph.models.SQLArtist.get_random()
    artist_id = artist.id
    return redirect('/artist/{}'.format(artist_id), code=302)

@app.route('/artist/<int:artist_id>', methods=['GET'])
def route__artist_id(artist_id):
    query = discograph.SQLArtist.select()
    query = query.where(discograph.SQLArtist.id == artist_id)
    result = list(query)
    if not result:
        abort(404)
    artist = result[0]
    return render_template(
        'index.html',
        key='artist-{}'.format(artist.id),
        title='discoGraph: {}'.format(artist.name),
        )

@app.route('/api/artist/network/<int:artist_id>', methods=['GET'])
def route__api__cluster(artist_id):
    try:
        artist = discograph.models.Artist.objects.get(discogs_id=artist_id)
    except:
        abort(404)
    role_names = [
        'Alias',
        'Member Of',
        #'Producer',
        #'Remix',
        #'Guitar',
        #'Bass Guitar',
        #'Rhythm Guitar',
        #'Electric Guitar',
        #'Lead Guitar',
        #'Drums',
        #'Vocals',
        #'Lead Vocals',
        #'Backing Vocals',
        ]
    relation_grapher = discograph.RelationGrapher(
        entities=[artist],
        degree=12,
        max_nodes=100,
        role_names=role_names,
        )
    with abjad.systemtools.Timer(exit_message='Network query time:'):
        data = relation_grapher.get_network()
    return jsonify(data)

urlify_pattern = re.compile(r"\s+", re.MULTILINE)

@app.route('/api/search/<search_string>', methods=['GET'])
def route__api__search(search_string):
    data = discograph.models.Artist.search_text(search_string, limit=50)
    data.sort(key=lambda x: Levenshtein.distance(x['name'], search_string))
    print(search_string)
    for datum in data:
        datum['key'] = 'artist-{}'.format(datum['_id'])
        del(datum['_id'])
        print(datum)
    data = {'results': data}
    return jsonify(data)

if __name__ == '__main__':
    app.run(debug=True)<commit_msg>Integrated SQLite full-text search into Flask app.<commit_after>#! /usr/bin/env python
import Levenshtein
import abjad
import discograph
import mongoengine
import re
from flask import (
    Flask,
    abort,
    jsonify,
    redirect,
    render_template,
    )
#from werkzeug.contrib.cache import MemcachedCache

app = Flask(__name__)
#cache = MemcachedCache(['127.0.0.1:11211'])
#cache.clear()
#cache = None
mongoengine.connect('discograph')

@app.route('/')
def route__index():
    return render_template(
        'index.html',
        artist=None,
        title='discoGraph',
        )

@app.route('/random')
def route__random():
    artist = discograph.models.SQLArtist.get_random()
    artist_id = artist.id
    return redirect('/artist/{}'.format(artist_id), code=302)

@app.route('/artist/<int:artist_id>', methods=['GET'])
def route__artist_id(artist_id):
    query = discograph.SQLArtist.select()
    query = query.where(discograph.SQLArtist.id == artist_id)
    result = list(query)
    if not result:
        abort(404)
    artist = result[0]
    return render_template(
        'index.html',
        key='artist-{}'.format(artist.id),
        title='discoGraph: {}'.format(artist.name),
        )

@app.route('/api/artist/network/<int:artist_id>', methods=['GET'])
def route__api__cluster(artist_id):
    try:
        artist = discograph.models.Artist.objects.get(discogs_id=artist_id)
    except:
        abort(404)
    role_names = [
        'Alias',
        'Member Of',
        #'Producer',
        #'Remix',
        #'Guitar',
        #'Bass Guitar',
        #'Rhythm Guitar',
        #'Electric Guitar',
        #'Lead Guitar',
        #'Drums',
        #'Vocals',
        #'Lead Vocals',
        #'Backing Vocals',
        ]
    relation_grapher = discograph.RelationGrapher(
        entities=[artist],
        degree=12,
        max_nodes=100,
        role_names=role_names,
        )
    with abjad.systemtools.Timer(exit_message='Network query time:'):
        data = relation_grapher.get_network()
    return jsonify(data)

urlify_pattern = re.compile(r"\s+", re.MULTILINE)

@app.route('/api/search/<search_string>', methods=['GET'])
def route__api__search(search_string):
    print(search_string)
    result = discograph.SQLFTSArtist.search_bm25(search_string)
    data = []
    for sql_fts_artist in result:
        datum = dict(
            key='artist-{}'.format(sql_fts_artist.id),
            name=sql_fts_artist.name,
            )
        data.append(datum)
        print('    {}'.format(datum))
    data = {'results': tuple(data)}
    return jsonify(data)

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


<commit_before>'''
This script calculates the gross emissions in tonnes CO2e/ha for every loss pixel.
The properties of each pixel determine the appropriate emissions equation, the constants for the equation, and the
carbon pool values that go into the equation.
Unlike all other flux model components, this one uses C++ to quickly iterate through every pixel in each tile.
Before running the model, the C++ script must be compiled.
From carbon-budget/emissions/, do c++ ./cpp_util/calc_gross_emissions_biomass_soil.cpp -o ./cpp_util/calc_gross_emissions_biomass_soil.exe -lgdal
calc_emissions_v3_biomass_soil.exe should appear in the directory.
Run mp_calculate_gross_emissions.py by typing python mp_calculate_gross_emissions.py -p [POOL_OPTION]. The Python script will call the
compiled C++ code as needed.
The other C++ scripts (equations.cpp and flu_val.cpp) do not need to be compiled.
The --pools-to-use argument specifies whether to calculate gross emissions from biomass+soil or just from soil.
The --model-type argument specifies whether the model run is a sensitivity analysis or standard run.
Emissions from each driver (including loss that had no driver assigned) gets its own tile, as does all emissions combined.
Emissions from all drivers is also output as emissions due to CO2 only and emissions due to other GHG (CH4 and N2O).
The other output shows which branch of the decision tree that determines the emissions equation applies to each pixel.
These codes are summarized in carbon-budget/emissions/node_codes.txt
'''

import multiprocessing
import argparse
import os
import calculate_gross_emissions
from functools import partial
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

def main ():

    # One argument for the script: whether only emissions from biomass (soil_only) is being calculated or emissions from biomass and soil (biomass_soil)
    parser = argparse.ArgumentParser(description='Calculate gross emissions')
    parser.add_argument('--pools-to-use', '-p', required=True,
                        help='Options are soil_only or biomass_soil. Former only considers emissions from soil. Latter considers emissions from biomass and soil.')
    parser.add_argument('--model-type', '-t', required=True,
                        help='{}'.format(cn.model_type_arg_help))
    args = parser.parse_args()
    sensit_type = args.model_type
    pools = args.pools_to_use
    # Checks whether the sensitivity analysis argument is valid
    uu.check_sensit_type(sensit_type)

    # Checks the validity of the pools argument
    if (pools not in ['soil_only', 'biomass_soil']):
        raise Exception('Invalid pool input. Please choose soil_only or biomass_soil.')

    tile_id_list = uu.tile_list(cn.AGC_emis_year_dir)
    # tile_id_list = ['00N_110E', '30N_080W', '40N_050E', '50N_100E', '80N_020E'] # test tiles
    # tile_id_list = ['20N_010E'] # test tiles
    # tile_id_list = ['00N_110E', '80N_020E', '30N_080W', '00N_020E'] # test tiles: no mangrove or planted forest, mangrove only, planted forest only, mangrove and planted forest
    print tile_id_list
    print "There are {} tiles to process".format(str(len(tile_id_list)))

    # Checks if the correct c++ script has been compiled for the pool option selected
    if pools == 'biomass_soil':
        if os.path.exists('./cpp_util/calc_gross_emissions_biomass_soil.exe'):
            print "C++ for biomass+soil already compiled."

            # Output file directories for biomass+soil. Must be in same order as output pattern directories.
            output_dir_list = [cn.gross_emis_commod_biomass_soil_dir,
                               cn.gross_emis_shifting_ag_biomass_soil_dir,
                               cn.gross_emis_forestry_biomass_soil_dir,
                               cn.gross_emis_wildfire_biomass_soil_dir,
                               cn.gross_emis_urban_biomass_soil_dir,
                               cn.gross_emis_no_driver_biomass_soil_dir,
                               cn.gross_emis_all_gases_all_drivers_biomass_soil_dir,
                               cn.gross_emis_co2_only_all_drivers_biomass_soil_dir,
                               cn.gross_emis_non_co2_all_drivers_biomass_soil_dir,
                               cn.gross_emis_nodes_biomass_soil_dir]

            output_pattern_list = [cn.pattern_gross_emis_commod_biomass_soil,
                                   cn.pattern_gross_emis_shifting_ag_biomass_soil,
                                   cn.pattern_gross_emis_forestry_biomass_soil,
                                   cn.pattern_gross_emis_wildfire_biomass_soil,
                                   cn.pattern_gross_emis_urban_biomass_soil,
                                   cn.pattern_gross_emis_no_driver_biomass_soil,
                                   cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil,
                                   cn.pattern_gross_emis_co2_only_all_drivers_biomass_soil,
                                   cn.pattern_gross_emis_non_co2_all_drivers_biomass_soil,
                                   cn.pattern_gross_emis_nodes_biomass_soil]
        else:
            raise Exception('Must compile biomass+soil C++...')

    elif pools == 'soil_only':
        if os.path.exists('./cpp_util/calc_gross_emissions_soil_only.exe'):
            print "C++ for soil_only already compiled."

            # Output file directories for soil_only. Must be in same order as output pattern directories.
            output_dir_list = [cn.gross_emis_commod_soil_only_dir,
                               cn.gross_emis_shifting_ag_soil_only_dir,
                               cn.gross_emis_forestry_soil_only_dir,
                               cn.gross_emis_wildfire_soil_only_dir,
                               cn.gross_emis_urban_soil_only_dir,
                               cn.gross_emis_no_driver_soil_only_dir,
                               cn.gross_emis_all_gases_all_drivers_soil_only_dir,
                               cn.gross_emis_co2_only_all_drivers_soil_only_dir,
                               cn.gross_emis_non_co2_all_drivers_soil_only_dir,
                               cn.gross_emis_nodes_soil_only_dir]

            output_pattern_list = [cn.pattern_gross_emis_commod_soil_only,
                                   cn.pattern_gross_emis_shifting_ag_soil_only,
                                   cn.pattern_gross_emis_forestry_soil_only,
                                   cn.pattern_gross_emis_wildfire_soil_only,
                                   cn.pattern_gross_emis_urban_soil_only,
                                   cn.pattern_gross_emis_no_driver_soil_only,
                                   cn.pattern_gross_emis_all_gases_all_drivers_soil_only,
                                   cn.pattern_gross_emis_co2_only_all_drivers_soil_only,
                                   cn.pattern_gross_emis_non_co2_all_drivers_soil_only,
                                   cn.pattern_gross_emis_nodes_soil_only]
        else:
            raise Exception('Must compile soil_only C++...')

    else:
        raise Exception('Pool option not valid')

    download_dict = {
        cn.AGC_emis_year_dir: [cn.pattern_AGC_emis_year, 'true'],
        cn.BGC_emis_year_dir: [cn.pattern_BGC_emis_year, 'true'],
        cn.deadwood_emis_year_2000_dir: [cn.pattern_deadwood_emis_year_2000, 'true'],
        cn.litter_emis_year_2000_dir: [cn.pattern_litter_emis_year_2000, 'true'],
        cn.soil_C_emis_year_2000_dir: [cn.pattern_soil_C_emis_year_2000, 'true'],
        cn.peat_mask_dir: [cn.pattern_peat_mask, 'false'],
        cn.ifl_primary_processed_dir: [cn.pattern_ifl_primary, 'false'],
        cn.planted_forest_type_unmasked_dir: [cn.pattern_planted_forest_type_unmasked, 'false'],
        cn.drivers_processed_dir: [cn.pattern_drivers, 'false'],
        cn.climate_zone_processed_dir: [cn.pattern_climate_zone, 'false'],
        cn.bor_tem_trop_processed_dir: [cn.pattern_bor_tem_trop_processed, 'false'],
        cn.burn_year_dir: [cn.pattern_burn_year, 'false'],
        cn.plant_pre_2000_processed_dir: [cn.pattern_plant_pre_2000, 'false'],
        cn.loss_dir: ['', 'false']
    }

    for key, values in download_dict.iteritems():
        dir = key
        pattern = values[0]
        sensit_use = values[1]
        # uu.s3_flexible_download(dir, pattern, '.', sensit_type, sensit_use, tile_id_list)

    # If the model run isn't the standard one, the output directory and file names are changed
    if sensit_type != 'std':
        print "Changing output directory and file name pattern based on sensitivity analysis"
        output_dir_list = uu.alter_dirs(sensit_type, output_dir_list)
        output_pattern_list = uu.alter_patterns(sensit_type, output_pattern_list)

    print "Removing loss pixels from plantations that existed in Indonesia and Malaysia before 2000..."
    # # Pixels that were in plantations that existed before 2000 should not be included in gross emissions.
    # # Pre-2000 plantations have not previously been masked, so that is done here.
    # # There are only 8 tiles to process, so count/2 will cover all of them in one go.
    # count = multiprocessing.cpu_count()
    # pool = multiprocessing.Pool(count/2)
    # pool.map(calculate_gross_emissions.mask_pre_2000_plant, tile_id_list)

    # For single processor use
    for tile in tile_id_list:
          calculate_gross_emissions.mask_pre_2000_plant(tile)

    # The C++ code expects a plantations tile for every input 10x10.
    # However, not all Hansen tiles have plantations.
    # This function creates "dummy" plantation tiles for all Hansen tiles that do not have plantations.
    # That way, the C++ script gets all the necessary input files
    folder = 'cpp_util/'

    # All of the inputs that need to have dummy tiles made in order to match the tile list of the carbon pools
    pattern_list = [cn.pattern_planted_forest_type_unmasked, cn.pattern_peat_mask, cn.pattern_ifl_primary,
                    cn.pattern_drivers, cn.pattern_bor_tem_trop_processed]

    for pattern in pattern_list:
        count = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(count-10)
        pool.map(partial(uu.make_blank_tile, pattern=pattern, folder=folder), tile_id_list)
        pool.close()
        pool.join()

    # # For single processor use
    # for pattern in pattern_list:
    #     for tile in tile_id_list:
    #         uu.make_blank_tile(tile, pattern, folder)

    # Calculates gross emissions for each tile
    # count/4 uses about 390 GB on a r4.16xlarge spot machine.
    # processes=18 uses about 440 GB on an r4.16xlarge spot machine.
    count = multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=18)
    pool.map(partial(calculate_gross_emissions.calc_emissions, pools=pools), tile_id_list)

    # # For single processor use
    # for tile in tile_id_list:
    #       calculate_gross_emissions.calc_emissions(tile, pools)

    # Uploads emissions to appropriate directory for the carbon pools chosen
    for i in range(0, len(output_dir_list)):
        uu.upload_final_set(output_dir_list[i], output_pattern_list[i])

if __name__ == '__main__':
    main()<commit_msg>Running gross emissions for biomass+soil and soil only for model v1.1.2.<commit_after>'''
This script calculates the gross emissions in tonnes CO2e/ha for every loss pixel.
The properties of each pixel determine the appropriate emissions equation, the constants for the equation, and the
carbon pool values that go into the equation.
Unlike all other flux model components, this one uses C++ to quickly iterate through every pixel in each tile.
Before running the model, the C++ script must be compiled.
From carbon-budget/emissions/, do c++ ./cpp_util/calc_gross_emissions_biomass_soil.cpp -o ./cpp_util/calc_gross_emissions_biomass_soil.exe -lgdal
calc_emissions_v3_biomass_soil.exe should appear in the directory.
Run mp_calculate_gross_emissions.py by typing python mp_calculate_gross_emissions.py -p [POOL_OPTION]. The Python script will call the
compiled C++ code as needed.
The other C++ scripts (equations.cpp and flu_val.cpp) do not need to be compiled.
The --pools-to-use argument specifies whether to calculate gross emissions from biomass+soil or just from soil.
The --model-type argument specifies whether the model run is a sensitivity analysis or standard run.
Emissions from each driver (including loss that had no driver assigned) gets its own tile, as does all emissions combined.
Emissions from all drivers is also output as emissions due to CO2 only and emissions due to other GHG (CH4 and N2O).
The other output shows which branch of the decision tree that determines the emissions equation applies to each pixel.
These codes are summarized in carbon-budget/emissions/node_codes.txt
'''

import multiprocessing
import argparse
import os
import calculate_gross_emissions
from functools import partial
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

def main ():

    # One argument for the script: whether only emissions from biomass (soil_only) is being calculated or emissions from biomass and soil (biomass_soil)
    parser = argparse.ArgumentParser(description='Calculate gross emissions')
    parser.add_argument('--pools-to-use', '-p', required=True,
                        help='Options are soil_only or biomass_soil. Former only considers emissions from soil. Latter considers emissions from biomass and soil.')
    parser.add_argument('--model-type', '-t', required=True,
                        help='{}'.format(cn.model_type_arg_help))
    args = parser.parse_args()
    sensit_type = args.model_type
    pools = args.pools_to_use
    # Checks whether the sensitivity analysis argument is valid
    uu.check_sensit_type(sensit_type)

    # Checks the validity of the pools argument
    if (pools not in ['soil_only', 'biomass_soil']):
        raise Exception('Invalid pool input. Please choose soil_only or biomass_soil.')

    tile_id_list = uu.tile_list(cn.AGC_emis_year_dir)
    # tile_id_list = ['00N_110E', '30N_080W', '40N_050E', '50N_100E', '80N_020E'] # test tiles
    # tile_id_list = ['20N_010E'] # test tiles
    # tile_id_list = ['00N_110E', '80N_020E', '30N_080W', '00N_020E'] # test tiles: no mangrove or planted forest, mangrove only, planted forest only, mangrove and planted forest
    print tile_id_list
    print "There are {} tiles to process".format(str(len(tile_id_list)))

    # Checks if the correct c++ script has been compiled for the pool option selected
    if pools == 'biomass_soil':
        if os.path.exists('./cpp_util/calc_gross_emissions_biomass_soil.exe'):
            print "C++ for biomass+soil already compiled."

            # Output file directories for biomass+soil. Must be in same order as output pattern directories.
            output_dir_list = [cn.gross_emis_commod_biomass_soil_dir,
                               cn.gross_emis_shifting_ag_biomass_soil_dir,
                               cn.gross_emis_forestry_biomass_soil_dir,
                               cn.gross_emis_wildfire_biomass_soil_dir,
                               cn.gross_emis_urban_biomass_soil_dir,
                               cn.gross_emis_no_driver_biomass_soil_dir,
                               cn.gross_emis_all_gases_all_drivers_biomass_soil_dir,
                               cn.gross_emis_co2_only_all_drivers_biomass_soil_dir,
                               cn.gross_emis_non_co2_all_drivers_biomass_soil_dir,
                               cn.gross_emis_nodes_biomass_soil_dir]

            output_pattern_list = [cn.pattern_gross_emis_commod_biomass_soil,
                                   cn.pattern_gross_emis_shifting_ag_biomass_soil,
                                   cn.pattern_gross_emis_forestry_biomass_soil,
                                   cn.pattern_gross_emis_wildfire_biomass_soil,
                                   cn.pattern_gross_emis_urban_biomass_soil,
                                   cn.pattern_gross_emis_no_driver_biomass_soil,
                                   cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil,
                                   cn.pattern_gross_emis_co2_only_all_drivers_biomass_soil,
                                   cn.pattern_gross_emis_non_co2_all_drivers_biomass_soil,
                                   cn.pattern_gross_emis_nodes_biomass_soil]
        else:
            raise Exception('Must compile biomass+soil C++...')

    elif pools == 'soil_only':
        if os.path.exists('./cpp_util/calc_gross_emissions_soil_only.exe'):
            print "C++ for soil_only already compiled."

            # Output file directories for soil_only. Must be in same order as output pattern directories.
            output_dir_list = [cn.gross_emis_commod_soil_only_dir,
                               cn.gross_emis_shifting_ag_soil_only_dir,
                               cn.gross_emis_forestry_soil_only_dir,
                               cn.gross_emis_wildfire_soil_only_dir,
                               cn.gross_emis_urban_soil_only_dir,
                               cn.gross_emis_no_driver_soil_only_dir,
                               cn.gross_emis_all_gases_all_drivers_soil_only_dir,
                               cn.gross_emis_co2_only_all_drivers_soil_only_dir,
                               cn.gross_emis_non_co2_all_drivers_soil_only_dir,
                               cn.gross_emis_nodes_soil_only_dir]

            output_pattern_list = [cn.pattern_gross_emis_commod_soil_only,
                                   cn.pattern_gross_emis_shifting_ag_soil_only,
                                   cn.pattern_gross_emis_forestry_soil_only,
                                   cn.pattern_gross_emis_wildfire_soil_only,
                                   cn.pattern_gross_emis_urban_soil_only,
                                   cn.pattern_gross_emis_no_driver_soil_only,
                                   cn.pattern_gross_emis_all_gases_all_drivers_soil_only,
                                   cn.pattern_gross_emis_co2_only_all_drivers_soil_only,
                                   cn.pattern_gross_emis_non_co2_all_drivers_soil_only,
                                   cn.pattern_gross_emis_nodes_soil_only]
        else:
            raise Exception('Must compile soil_only C++...')

    else:
        raise Exception('Pool option not valid')

    download_dict = {
        cn.AGC_emis_year_dir: [cn.pattern_AGC_emis_year, 'true'],
        cn.BGC_emis_year_dir: [cn.pattern_BGC_emis_year, 'true'],
        cn.deadwood_emis_year_2000_dir: [cn.pattern_deadwood_emis_year_2000, 'true'],
        cn.litter_emis_year_2000_dir: [cn.pattern_litter_emis_year_2000, 'true'],
        cn.soil_C_emis_year_2000_dir: [cn.pattern_soil_C_emis_year_2000, 'true'],
        cn.peat_mask_dir: [cn.pattern_peat_mask, 'false'],
        cn.ifl_primary_processed_dir: [cn.pattern_ifl_primary, 'false'],
        cn.planted_forest_type_unmasked_dir: [cn.pattern_planted_forest_type_unmasked, 'false'],
        cn.drivers_processed_dir: [cn.pattern_drivers, 'false'],
        cn.climate_zone_processed_dir: [cn.pattern_climate_zone, 'false'],
        cn.bor_tem_trop_processed_dir: [cn.pattern_bor_tem_trop_processed, 'false'],
        cn.burn_year_dir: [cn.pattern_burn_year, 'false'],
        cn.plant_pre_2000_processed_dir: [cn.pattern_plant_pre_2000, 'false'],
        cn.loss_dir: ['', 'false']
    }

    for key, values in download_dict.iteritems():
        dir = key
        pattern = values[0]
        sensit_use = values[1]
        uu.s3_flexible_download(dir, pattern, './cpp_util/', sensit_type, sensit_use, tile_id_list)

    # If the model run isn't the standard one, the output directory and file names are changed
    if sensit_type != 'std':
        print "Changing output directory and file name pattern based on sensitivity analysis"
        output_dir_list = uu.alter_dirs(sensit_type, output_dir_list)
        output_pattern_list = uu.alter_patterns(sensit_type, output_pattern_list)

    print "Removing loss pixels from plantations that existed in Indonesia and Malaysia before 2000..."
    # # Pixels that were in plantations that existed before 2000 should not be included in gross emissions.
    # # Pre-2000 plantations have not previously been masked, so that is done here.
    # # There are only 8 tiles to process, so count/2 will cover all of them in one go.
    # count = multiprocessing.cpu_count()
    # pool = multiprocessing.Pool(count/2)
    # pool.map(calculate_gross_emissions.mask_pre_2000_plant, tile_id_list)

    # For single processor use
    for tile in tile_id_list:
          calculate_gross_emissions.mask_pre_2000_plant(tile)

    # The C++ code expects a plantations tile for every input 10x10.
    # However, not all Hansen tiles have plantations.
    # This function creates "dummy" plantation tiles for all Hansen tiles that do not have plantations.
    # That way, the C++ script gets all the necessary input files
    folder = 'cpp_util/'

    # All of the inputs that need to have dummy tiles made in order to match the tile list of the carbon pools
    pattern_list = [cn.pattern_planted_forest_type_unmasked, cn.pattern_peat_mask, cn.pattern_ifl_primary,
                    cn.pattern_drivers, cn.pattern_bor_tem_trop_processed]

    for pattern in pattern_list:
        count = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(count-10)
        pool.map(partial(uu.make_blank_tile, pattern=pattern, folder=folder), tile_id_list)
        pool.close()
        pool.join()

    # # For single processor use
    # for pattern in pattern_list:
    #     for tile in tile_id_list:
    #         uu.make_blank_tile(tile, pattern, folder)

    # Calculates gross emissions for each tile
    # count/4 uses about 390 GB on a r4.16xlarge spot machine.
    # processes=18 uses about 440 GB on an r4.16xlarge spot machine.
    count = multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=18)
    pool.map(partial(calculate_gross_emissions.calc_emissions, pools=pools), tile_id_list)

    # # For single processor use
    # for tile in tile_id_list:
    #       calculate_gross_emissions.calc_emissions(tile, pools)

    # Uploads emissions to appropriate directory for the carbon pools chosen
    for i in range(0, len(output_dir_list)):
        uu.upload_final_set(output_dir_list[i], output_pattern_list[i])

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


<commit_before>""" import and bulid a bunch of example programs """

import os
from os.path import dirname, abspath, basename
import os.path
import sys
import subprocess
import json

ROOT = abspath(dirname(dirname(dirname(dirname(__file__)))))
sys.path.insert(0, ROOT)

from tools.build_api import get_mbed_official_release
from tools.targets import TARGET_MAP

EXAMPLES = json.load(open(os.path.join(os.path.dirname(__file__),
                                       "examples.json")))

def print_stuff(name, lst):
    if list:
        print("#"*80)
        print("# {} example combinations".format(name))
        print("#")
    for thing in lst:
        print(thing)

def main():
    """Entry point"""
    failures = []
    sucesses = []
    for example, build_features in EXAMPLES.iteritems():
        subprocess.call(["mbed-cli", "import", example])
        os.chdir(basename(example))
        for target, toolchains in [(target, toolchains) for target, toolchains
                                   in get_mbed_official_release("5") if
                                   all(feature in TARGET_MAP[target].features
                                       for feature in build_features)]:
            for toolchain in toolchains:
                proc = subprocess.Popen(["mbed-cli", "compile", "-t",
                                         toolchain, "-m", target])
                proc.wait()
                example_name = "{} {} {}".format(basename(example), target,
                                                 toolchain)
                if proc.returncode:
                    failures.append(example_name)
                else:
                    sucesses.append(example_name)
        os.chdir("..")
    print_stuff("Passed", sucesses)
    print_stuff("Failed", failures)
    return len(failures)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Allow command-line filtering of toolchains<commit_after>""" import and bulid a bunch of example programs """

from argparse import ArgumentParser
import os
from os.path import dirname, abspath, basename
import os.path
import sys
import subprocess
import json

ROOT = abspath(dirname(dirname(dirname(dirname(__file__)))))
sys.path.insert(0, ROOT)

from tools.build_api import get_mbed_official_release
from tools.targets import TARGET_MAP
from tools.utils import argparse_force_uppercase_type

EXAMPLES = json.load(open(os.path.join(os.path.dirname(__file__),
                                       "examples.json")))

def print_stuff(name, lst):
    if lst:
        print("#"*80)
        print("# {} example combinations".format(name))
        print("#")
    for thing in lst:
        print(thing)

SUPPORTED_TOOLCHAINS = ["ARM", "IAR", "GCC_ARM"]

def target_cross_toolchain(required_features, allowed_toolchains):
    """Generate pairs of target and toolchains

    Args:
    required_features - the features that must be in the features array of a
                        target
    allowed_toolchains - a list of all possible toolchains
    """
    for target, toolchains in get_mbed_official_release("5"):
        for toolchain in toolchains:
            if (toolchain in allowed_toolchains and
                all(feature in TARGET_MAP[target].features
                    for feature in required_features)):
                yield target, toolchain

def main():
    """Entry point"""
    parser = ArgumentParser()
    parser.add_argument(
        "toolchains", nargs="*", default=SUPPORTED_TOOLCHAINS,
        type=argparse_force_uppercase_type(SUPPORTED_TOOLCHAINS,
                                           "toolchain"))
    args = parser.parse_args()

    failures = []
    sucesses = []
    for example, build_features in EXAMPLES.iteritems():
        subprocess.call(["mbed-cli", "import", example])
        os.chdir(basename(example))
        for target, toolchain in target_cross_toolchain(build_features,
                                                        args.toolchains):
            proc = subprocess.Popen(["mbed-cli", "compile", "-t",
                                     toolchain, "-m", target])
            proc.wait()
            example_name = "{} {} {}".format(basename(example), target,
                                             toolchain)
            if proc.returncode:
                failures.append(example_name)
            else:
                sucesses.append(example_name)
        os.chdir("..")

    print_stuff("Passed", sucesses)
    print_stuff("Failed", failures)
    return len(failures)

if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>


<commit_before>import numpy as np
import pandas as pd
import matplotlib as mpl
mpl.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)

plt.figure(figsize=snakemake.config["plots"]["figsize"])

pmfs = []
for i, (mean, pmf, known_counts) in enumerate(zip(
    snakemake.params.means,
    snakemake.input.pmfs,
    snakemake.input.known_counts)):
    print(mean)

    pmf = pd.read_table(pmf, index_col=[0, 1])
    known_counts = pd.read_table(known_counts, index_col=[0, 1])
    codebook = "mhd4" if snakemake.wildcards.dist == "4" else "mhd2"
    known_counts = known_counts[known_counts[codebook]]["count"]
    common = pmf.index.intersection(known_counts.index)

    known_counts = known_counts.loc[common]
    pmf = pmf.loc[common]

    pmf["expr"] -= known_counts
    for _, _pmf in pmf.groupby(level=[0, 1]):
        _pmf.reset_index(drop=True, inplace=True)
        pmfs.append(pd.Series(data=_pmf["prob"], index=_pmf["expr"]))
    break

pmfs = pd.concat(pmfs, axis=1).fillna(0)
sns.heatmap(pmfs)
plt.ylabel("predicted - truth")
plt.xlabel("PMF")

plt.savefig(snakemake.output[0], bbox_inches="tight")
<commit_msg>Plotting.<commit_after>import numpy as np
import pandas as pd
import matplotlib as mpl
mpl.use("agg")
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(style="ticks", palette="colorblind", context=snakemake.wildcards.context)

plt.figure(figsize=snakemake.config["plots"]["figsize"])

pmfs = []
for i, (mean, pmf, known_counts) in enumerate(zip(
    snakemake.params.means,
    snakemake.input.pmfs,
    snakemake.input.known_counts)):
    print(mean)

    pmf = pd.read_table(pmf, index_col=[0, 1])
    known_counts = pd.read_table(known_counts, index_col=[0, 1])
    codebook = "mhd4" if snakemake.wildcards.dist == "4" else "mhd2"
    known_counts = known_counts[known_counts[codebook]]["count"]
    common = pmf.index.intersection(known_counts.index)

    known_counts = known_counts.loc[common]
    pmf = pmf.loc[common]

    pmf["expr"] -= known_counts
    pmf.set_index("expr", append=True, inplace=True)
    pmf = pmf.unstack(level=[0, 1])
    pmf = np.exp(pmf.sample(10, axis=1))
    pmf.columns = np.arange(pmf.shape[1])
    pmfs.append(pmf)

print("plotting")
pmfs = pd.concat(pmfs, axis=1)
pmfs = pmfs.reindex(np.arange(-10, 11))
pmfs.fillna(0, inplace=True)
pmfs.sort_index(ascending=False, inplace=True)
sns.heatmap(pmfs, cbar=False, cmap="Greys", xticklabels=False, yticklabels=10)

plt.ylabel("predicted - truth")
plt.xlabel("PMF")

plt.savefig(snakemake.output[0], bbox_inches="tight")
<|endoftext|>


<commit_before>from django.conf.urls import patterns, include, url

# Uncomment the next two lines to enable the admin:
from django.contrib import admin
admin.autodiscover()

urlpatterns = patterns('',

    # admin documentation
    url(r'^admin/doc/', include('django.contrib.admindocs.urls')),

    # admin
    url(r'^admin/', include(admin.site.urls)),

    url(r'^/', include('courses.urls'))
)
<commit_msg>Fixed urls.py<commit_after>from django.conf.urls import *

# Uncomment the next two lines to enable the admin:
from django.contrib import admin
admin.autodiscover()

urlpatterns = patterns('',

    # admin documentation
    url(r'^admin/doc/', include('django.contrib.admindocs.urls')),

    # admin
    url(r'^admin/', include(admin.site.urls)),

    url(r'^', include('courses.urls'))
)
<|endoftext|>


<commit_before>import os
import socket
import threading
import Queue
import time
import struct
import logging

from consts import *
from utils import *

logger = logging.getLogger(__name__)

# mfsmaster need to been patched with dcache
ENABLE_DCACHE = False

class StatInfo:
    def __init__(self, totalspace, availspace, trashspace,
                reservedspace, inodes):
        self.totalspace    = totalspace
        self.availspace    = availspace
        self.trashspace    = trashspace
        self.reservedspace = reservedspace
        self.inodes        = inodes

class Chunk:
    def __init__(self, id, length, version, csdata):
        self.id = id
        self.length = length
        self.version = version
        self.addrs = self._parse(csdata)

    def _parse(self, csdata):
        return [(socket.inet_ntoa(csdata[i:i+4]),
                    unpack("H", csdata[i+4:i+6])[0])
                for i in range(len(csdata))[::6]]

    def __repr__(self):
        return "<Chunk(%d, %d, %d)>" % (self.id, self.version, self.length)

def try_again(f):
    def _(self, *a, **kw):
        for i in range(3):
            try:
                return f(self, *a, **kw)
            except IOError, e:
                self.close()
                logger.warning("mfs master connection: %s", e)
                time.sleep(2**i*0.1)
        else:
            raise
    return _

def spawn(target, *args, **kw):
    t = threading.Thread(target=target, name=target.__name__, args=args, kwargs=kw)
    t.daemon = True
    t.start()
    return t

class MasterConn:
    def __init__(self, host='mfsmaster', port=9421):
        self.host = host
        self.port = port
        self.uid = os.getuid()
        self.gid = os.getgid()
        self.sessionid = 0
        self.conn = None
        self.packetid = 0
        self.fail_count = 0
        self.dcache = {}
        self.dstat = {}

        self.lock = threading.RLock()
        self.reply = Queue.Queue()
        self.is_ready = False
        spawn(self.heartbeat)
        spawn(self.recv_thread)

    def heartbeat(self):
        while True:
            try:
                self.nop()
            except Exception, e:
                self.close()
            time.sleep(2)

    def connect(self):
        if self.conn is not None:
            return

        for _ in range(10):
            try:
                self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.conn.connect((self.host, self.port))
                break
            except socket.error, e:
                self.conn = None
                #self.next_try = time.time() + 1.5 ** self.fail_count
                self.fail_count += 1
                time.sleep(1.5 ** self.fail_count)

        if not self.conn:
            raise IOError("mfsmaster not availbale")

        regbuf = pack(CUTOMA_FUSE_REGISTER, FUSE_REGISTER_BLOB_NOACL,
                      self.sessionid, VERSION)
        self.send(regbuf)
        recv = self.recv(8)
        cmd, i = unpack("II", recv)
        if cmd != MATOCU_FUSE_REGISTER:
            raise Exception("got incorrect answer from mfsmaster %s" % cmd)

        if i not in (1, 4):
            raise Exception("got incorrect size from mfsmaster")

        data = self.recv(i)
        if i == 1:
            code, = unpack("B", data)
            if code != 0:
                raise Exception("mfsmaster register error: "
                        + mfs_strerror(code))
        if self.sessionid == 0:
            self.sessionid, = unpack("I", data)

        self.is_ready = True

    def close(self):
        with self.lock:
            if self.conn:
                self.conn.close()
                self.conn = None
                self.dcache.clear()
                self.is_ready = False

    def send(self, buf):
        with self.lock:
            conn = self.conn
        if not conn:
            raise IOError("not connected")
        n = conn.send(buf)
        while n < len(buf):
            sent = conn.send(buf[n:])
            if not sent:
                self.close()
                raise IOError("write to master failed")
            n += sent

    def nop(self):
        with self.lock:
            self.connect()
            msg = pack(ANTOAN_NOP, 0)
            self.send(msg)

    def recv(self, n):
        with self.lock:
            conn = self.conn
        if not conn:
            raise IOError("not connected")
        r = conn.recv(n)
        while len(r) < n:
            rr = conn.recv(n - len(r))
            if not rr:
                self.close()
                raise IOError("unexpected error: need %d" % (n-len(r)))
            r += rr
        return r

    def recv_cmd(self):
        d = self.recv(12)
        cmd, size = unpack("II", d)
        data = self.recv(size-4) if size > 4 else ''
        while cmd in (ANTOAN_NOP, MATOCU_FUSE_NOTIFY_ATTR, MATOCU_FUSE_NOTIFY_DIR):
            if cmd == ANTOAN_NOP:
                pass
            elif cmd == MATOCU_FUSE_NOTIFY_ATTR:
                while len(data) >= 43:
                    parent, inode = unpack("II", data)
                    attr = data[8:43]
                    if parent in self.dcache:
                        cache = self.dcache[parent]
                        for name in cache:
                            if cache[name].inode == inode:
                                cache[name] = attrToFileInfo(inode, attr)
                                break
                    data = data[43:]
            elif cmd == MATOCU_FUSE_NOTIFY_DIR:
                while len(data) >= 4:
                    inode, = unpack("I", data)
                    if inode in self.dcache:
                        del self.dcache[inode]
                        with self.lock:
                            self.send(pack(CUTOMA_FUSE_DIR_REMOVED, 0, inode))
                    data = data[4:]
            d = self.recv(12)
            cmd, size = unpack("II", d)
            data = self.recv(size-4) if size > 4 else ''
        return d, data

    def recv_thread(self):
        while True:
            with self.lock:
                if not self.is_ready:
                    time.sleep(0.01)
                    continue
            try:
                r = self.recv_cmd()
                self.reply.put(r)
            except IOError, e:
                self.reply.put(e)

    @try_again
    def sendAndReceive(self, cmd, *args):
        #print 'sendAndReceive', cmd, args
        self.packetid += 1
        msg = pack(cmd, self.packetid, *args)
        with self.lock:
            self.connect()
            while not self.reply.empty():
                self.reply.get_nowait()
            self.send(msg)
        r = self.reply.get()
        if isinstance(r, Exception):
            raise r
        h, d = r
        rcmd, size, pid = unpack("III", h)
        if rcmd != cmd+1 or pid != self.packetid or size <= 4:
            self.close()
            raise Exception("incorrect answer (%s!=%s, %s!=%s, %d<=4",
                rcmd, cmd+1, pid, self.packetid, size)
        if len(d) == 1 and ord(d[0]) != 0:
            raise Error(ord(d[0]))
        return d

    def statfs(self):
        ans = self.sendAndReceive(CUTOMA_FUSE_STATFS)
        return StatInfo(*unpack("QQQQI", ans))

#    def access(self, inode, modemask):
#        return self.sendAndReceive(CUTOMA_FUSE_ACCESS, inode,
#            self.uid, self.gid, uint8(modemask))
#
    def lookup(self, parent, name):
        if ENABLE_DCACHE:
            cache = self.dcache.get(parent)
            if cache is None and self.dstat.get(parent, 0) > 1:
                cache = self.getdirplus(parent)
            if cache is not None:
                return cache.get(name), None

            self.dstat[parent] = self.dstat.get(parent, 0) + 1

        ans = self.sendAndReceive(CUTOMA_FUSE_LOOKUP, parent,
                uint8(len(name)), name, 0, 0)
        if len(ans) == 1:
            return None, ""
        if len(ans) != 39:
            raise Exception("bad length")
        inode, = unpack("I", ans)
        return attrToFileInfo(inode, ans[4:]), None

    def getattr(self, inode):
        ans = self.sendAndReceive(CUTOMA_FUSE_GETATTR, inode,
                self.uid, self.gid)
        return attrToFileInfo(inode, ans)

    def readlink(self, inode):
        ans = self.sendAndReceive(CUTOMA_FUSE_READLINK, inode)
        length, = unpack("I", ans)
        if length+4 != len(ans):
            raise Exception("invalid length")
        return ans[4:-1]

    def getdir(self, inode):
        "return: {name: (inode,type)}"
        ans = self.sendAndReceive(CUTOMA_FUSE_GETDIR, inode,
                self.uid, self.gid)
        p = 0
        names = {}
        while p < len(ans):
            length, = unpack("B", ans[p:p+1])
            p += 1
            if length + p + 5 > len(ans):
                break
            name = ans[p:p+length]
            p += length
            inode, type = unpack("IB", ans)
            names[name] = (inode, type)
            p += 5
        return names

    def getdirplus(self, inode):
        "return {name: FileInfo()}"
        if ENABLE_DCACHE:
            infos = self.dcache.get(inode)
            if infos is not None:
                return infos

        flag = GETDIR_FLAG_WITHATTR
        if ENABLE_DCACHE:
            flag |= GETDIR_FLAG_DIRCACHE
        ans = self.sendAndReceive(CUTOMA_FUSE_GETDIR, inode,
                self.uid, self.gid, uint8(flag))
        p = 0
        infos = {}
        while p < len(ans):
            length, = unpack("B", ans[p:p+1])
            p += 1
            name = ans[p:p+length]
            p += length
            i, = unpack("I", ans[p:p+4])
            attr = ans[p+4:p+39]
            infos[name] = attrToFileInfo(i, attr, name)
            p += 39
        if ENABLE_DCACHE:
            self.dcache[inode] = infos
        return infos

    def opencheck(self, inode, flag=1):
        ans = self.sendAndReceive(CUTOMA_FUSE_OPEN, inode,
                self.uid, self.gid, uint8(flag))
        return ans

    def readchunk(self, inode, index):
        ans = self.sendAndReceive(CUTOMA_FUSE_READ_CHUNK, inode, index)
        n = len(ans)
        if n < 20 or (n-20)%6 != 0:
            raise Exception("read chunk: invalid length: %s" % n)
        length, id, version = unpack("QQI", ans)
        return Chunk(id, length, version, ans[20:])

def test():
    m = MasterConn("mfsmaster")
    m.connect()
    m.close()
    #print m.get_attr(1)
    while True:
        print m.getdir(1)
        print m.getdirplus(1)
        time.sleep(60)
    info, err = m.lookup(1, "test.csv")
    print info, err
    #print m.opencheck(info.inode)
    chunks = m.readchunk(info.inode, 0)
    print chunks, chunks.addrs

    for i in range(1000):
        info, err = m.lookup(1, "test.csv")
        chunks = m.readchunk(info.inode, 0)
        print i,err, chunks
        time.sleep(10)

    m.close()

if __name__ == '__main__':
    test()
<commit_msg>add append operation<commit_after>import os
import socket
import threading
import Queue
import time
import struct
import logging

from consts import *
from utils import *

logger = logging.getLogger(__name__)

# mfsmaster need to been patched with dcache
ENABLE_DCACHE = False

class StatInfo:

    def __init__(self, totalspace, availspace, trashspace,
                 reservedspace, inodes):
        self.totalspace = totalspace
        self.availspace = availspace
        self.trashspace = trashspace
        self.reservedspace = reservedspace
        self.inodes = inodes

class Chunk:

    def __init__(self, id, length, version, csdata):
        self.id = id
        self.length = length
        self.version = version
        self.addrs = self._parse(csdata)

    def _parse(self, csdata):
        return [(socket.inet_ntoa(csdata[i:i + 4]),
                 unpack("H", csdata[i + 4:i + 6])[0])
                for i in range(len(csdata))[::6]]

    def __repr__(self):
        return "<Chunk(%d, %d, %d)>" % (self.id, self.version, self.length)

def try_again(f):
    def _(self, *a, **kw):
        for i in range(3):
            try:
                return f(self, *a, **kw)
            except IOError as e:
                self.close()
                logger.warning("mfs master connection: %s", e)
                time.sleep(2**i * 0.1)
        else:
            raise
    return _

def spawn(target, *args, **kw):
    t = threading.Thread(
        target=target,
        name=target.__name__,
        args=args,
        kwargs=kw)
    t.daemon = True
    t.start()
    return t

class MasterConn:

    def __init__(self, host='mfsmaster', port=9421):
        self.host = host
        self.port = port
        self.uid = os.getuid()
        self.gid = os.getgid()
        self.sessionid = 0
        self.conn = None
        self.packetid = 0
        self.fail_count = 0
        self.dcache = {}
        self.dstat = {}

        self.lock = threading.RLock()
        self.reply = Queue.Queue()
        self.is_ready = False
        spawn(self.heartbeat)
        spawn(self.recv_thread)

    def heartbeat(self):
        while True:
            try:
                self.nop()
            except Exception as e:
                self.close()
            time.sleep(2)

    def connect(self):
        if self.conn is not None:
            return

        for _ in range(10):
            try:
                self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.conn.connect((self.host, self.port))
                break
            except socket.error as e:
                self.conn = None
                #self.next_try = time.time() + 1.5 ** self.fail_count
                self.fail_count += 1
                time.sleep(1.5 ** self.fail_count)

        if not self.conn:
            raise IOError("mfsmaster not availbale")

        regbuf = pack(CUTOMA_FUSE_REGISTER, FUSE_REGISTER_BLOB_NOACL,
                      self.sessionid, VERSION)
        self.send(regbuf)
        recv = self.recv(8)
        cmd, i = unpack("II", recv)
        if cmd != MATOCU_FUSE_REGISTER:
            raise Exception("got incorrect answer from mfsmaster %s" % cmd)

        if i not in (1, 4):
            raise Exception("got incorrect size from mfsmaster")

        data = self.recv(i)
        if i == 1:
            code, = unpack("B", data)
            if code != 0:
                raise Exception("mfsmaster register error: "
                                + mfs_strerror(code))
        if self.sessionid == 0:
            self.sessionid, = unpack("I", data)

        self.is_ready = True

    def close(self):
        with self.lock:
            if self.conn:
                self.conn.close()
                self.conn = None
                self.dcache.clear()
                self.is_ready = False

    def send(self, buf):
        with self.lock:
            conn = self.conn
        if not conn:
            raise IOError("not connected")
        n = conn.send(buf)
        while n < len(buf):
            sent = conn.send(buf[n:])
            if not sent:
                self.close()
                raise IOError("write to master failed")
            n += sent

    def nop(self):
        with self.lock:
            self.connect()
            msg = pack(ANTOAN_NOP, 0)
            self.send(msg)

    def recv(self, n):
        with self.lock:
            conn = self.conn
        if not conn:
            raise IOError("not connected")
        r = conn.recv(n)
        while len(r) < n:
            rr = conn.recv(n - len(r))
            if not rr:
                self.close()
                raise IOError("unexpected error: need %d" % (n - len(r)))
            r += rr
        return r

    def recv_cmd(self):
        d = self.recv(12)
        cmd, size = unpack("II", d)
        data = self.recv(size - 4) if size > 4 else ''
        while cmd in (ANTOAN_NOP, MATOCU_FUSE_NOTIFY_ATTR,
                      MATOCU_FUSE_NOTIFY_DIR):
            if cmd == ANTOAN_NOP:
                pass
            elif cmd == MATOCU_FUSE_NOTIFY_ATTR:
                while len(data) >= 43:
                    parent, inode = unpack("II", data)
                    attr = data[8:43]
                    if parent in self.dcache:
                        cache = self.dcache[parent]
                        for name in cache:
                            if cache[name].inode == inode:
                                cache[name] = attrToFileInfo(inode, attr)
                                break
                    data = data[43:]
            elif cmd == MATOCU_FUSE_NOTIFY_DIR:
                while len(data) >= 4:
                    inode, = unpack("I", data)
                    if inode in self.dcache:
                        del self.dcache[inode]
                        with self.lock:
                            self.send(pack(CUTOMA_FUSE_DIR_REMOVED, 0, inode))
                    data = data[4:]
            d = self.recv(12)
            cmd, size = unpack("II", d)
            data = self.recv(size - 4) if size > 4 else ''
        return d, data

    def recv_thread(self):
        while True:
            with self.lock:
                if not self.is_ready:
                    time.sleep(0.01)
                    continue
            try:
                r = self.recv_cmd()
                self.reply.put(r)
            except IOError as e:
                self.reply.put(e)

    @try_again
    def sendAndReceive(self, cmd, *args):
        # print 'sendAndReceive', cmd, args
        self.packetid += 1
        msg = pack(cmd, self.packetid, *args)
        with self.lock:
            self.connect()
            while not self.reply.empty():
                self.reply.get_nowait()
            self.send(msg)
        r = self.reply.get()
        if isinstance(r, Exception):
            raise r
        h, d = r
        rcmd, size, pid = unpack("III", h)
        if rcmd != cmd + 1 or pid != self.packetid or size <= 4:
            self.close()
            raise Exception("incorrect answer (%s!=%s, %s!=%s, %d<=4",
                            rcmd, cmd + 1, pid, self.packetid, size)
        if len(d) == 1 and ord(d[0]) != 0:
            raise Error(ord(d[0]))
        return d

    def statfs(self):
        ans = self.sendAndReceive(CUTOMA_FUSE_STATFS)
        return StatInfo(*unpack("QQQQI", ans))

#    def access(self, inode, modemask):
#        return self.sendAndReceive(CUTOMA_FUSE_ACCESS, inode,
#            self.uid, self.gid, uint8(modemask))
#
    def lookup(self, parent, name):
        if ENABLE_DCACHE:
            cache = self.dcache.get(parent)
            if cache is None and self.dstat.get(parent, 0) > 1:
                cache = self.getdirplus(parent)
            if cache is not None:
                return cache.get(name), None

            self.dstat[parent] = self.dstat.get(parent, 0) + 1

        ans = self.sendAndReceive(CUTOMA_FUSE_LOOKUP, parent,
                                  uint8(len(name)), name, 0, 0)
        if len(ans) == 1:
            return None, ""
        if len(ans) != 39:
            raise Exception("bad length")
        inode, = unpack("I", ans)
        return attrToFileInfo(inode, ans[4:]), None

    def getattr(self, inode):
        ans = self.sendAndReceive(CUTOMA_FUSE_GETATTR, inode,
                                  self.uid, self.gid)
        return attrToFileInfo(inode, ans)

    def readlink(self, inode):
        ans = self.sendAndReceive(CUTOMA_FUSE_READLINK, inode)
        length, = unpack("I", ans)
        if length + 4 != len(ans):
            raise Exception("invalid length")
        return ans[4:-1]

    def getdir(self, inode):
        "return: {name: (inode,type)}"
        ans = self.sendAndReceive(CUTOMA_FUSE_GETDIR, inode,
                                  self.uid, self.gid)
        p = 0
        names = {}
        while p < len(ans):
            length, = unpack("B", ans[p:p + 1])
            p += 1
            if length + p + 5 > len(ans):
                break
            name = ans[p:p + length]
            p += length
            inode, type = unpack("IB", ans)
            names[name] = (inode, type)
            p += 5
        return names

    def getdirplus(self, inode):
        "return {name: FileInfo()}"
        if ENABLE_DCACHE:
            infos = self.dcache.get(inode)
            if infos is not None:
                return infos

        flag = GETDIR_FLAG_WITHATTR
        if ENABLE_DCACHE:
            flag |= GETDIR_FLAG_DIRCACHE
        ans = self.sendAndReceive(CUTOMA_FUSE_GETDIR, inode,
                                  self.uid, self.gid, uint8(flag))
        p = 0
        infos = {}
        while p < len(ans):
            length, = unpack("B", ans[p:p + 1])
            p += 1
            name = ans[p:p + length]
            p += length
            i, = unpack("I", ans[p:p + 4])
            attr = ans[p + 4:p + 39]
            infos[name] = attrToFileInfo(i, attr, name)
            p += 39
        if ENABLE_DCACHE:
            self.dcache[inode] = infos
        return infos

    def opencheck(self, inode, flag=1):
        ans = self.sendAndReceive(CUTOMA_FUSE_OPEN, inode,
                                  self.uid, self.gid, uint8(flag))
        return ans

    def readchunk(self, inode, index):
        ans = self.sendAndReceive(CUTOMA_FUSE_READ_CHUNK, inode, index)
        n = len(ans)
        if n < 20 or (n - 20) % 6 != 0:
            raise Exception("read chunk: invalid length: %s" % n)
        length, id, version = unpack("QQI", ans)
        return Chunk(id, length, version, ans[20:])

    def append(self, inode_dst, inode_src):
        self.sendAndReceive(CUTOMA_FUSE_APPEND, inode_dst, inode_src,
                            self.uid, self.gid)

def test():
    m = MasterConn("mfsmaster")
    m.connect()
    m.close()
    # print m.get_attr(1)
    while True:
        print m.getdir(1)
        print m.getdirplus(1)
        time.sleep(60)
    info, err = m.lookup(1, "test.csv")
    print info, err
    # print m.opencheck(info.inode)
    chunks = m.readchunk(info.inode, 0)
    print chunks, chunks.addrs

    for i in range(1000):
        info, err = m.lookup(1, "test.csv")
        chunks = m.readchunk(info.inode, 0)
        print i, err, chunks
        time.sleep(10)

    m.close()

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