<commit_before>import pkg_resources
pkg_resources.declare_namespace(__name__)

from lazy import lazy
from mr.awsome.config import Config
from mr.awsome import template
import logging
import argparse
import os
import sys


# shutup pyflakes
__all__ = [template.__name__]


log = logging.getLogger('mr.awsome')


class AWS(object):
    def __init__(self, configpath=None):
        log.setLevel(logging.INFO)
        ch = logging.StreamHandler()
        ch.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
        log.addHandler(ch)
        if configpath is None:
            configpath = 'etc/aws.conf'
        if os.path.isdir(configpath):
            configpath = os.path.join(configpath, 'aws.conf')
        self.configfile = configpath

    @lazy
    def config(self):
        configpath = os.path.abspath(self.configfile)
        if not os.path.exists(configpath):
            log.error("Config '%s' doesn't exist." % configpath)
            sys.exit(1)
        config = Config(configpath, bbb_config=True)
        config.parse()
        return config

    @lazy
    def masters(self):
        masters = []
        for plugin in self.config['plugin'].values():
            masters.extend(plugin['module'].get_masters(self.config))
        return masters

    @lazy
    def known_hosts(self):
        return os.path.join(self.config.path, 'known_hosts')

    def get_masters(self, command):
        masters = []
        for master in self.masters:
            if getattr(master, command, None) is not None:
                masters.append(master)
        return masters

    @lazy
    def instances(self):
        instances = {}
        for master in self.masters:
            for instance_id in master.instances:
                if instance_id in instances:
                    log.error("Instance and server names must be unique, '%s' is already defined." % instance_id)
                    sys.exit(1)
                instances[instance_id] = master.instances[instance_id]
        return instances

    def get_instances(self, command):
        instances = {}
        for instance_id in self.instances:
            instance = self.instances[instance_id]
            if getattr(instance, command, None) is not None:
                instances[instance_id] = instance
        return instances

    def cmd_status(self, argv, help):
        """Prints status"""
        parser = argparse.ArgumentParser(
            prog="aws status",
            description=help,
        )
        instances = self.get_instances(command='status')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        args = parser.parse_args(argv)
        server = instances[args.server[0]]
        server.status()

    def cmd_stop(self, argv, help):
        """Stops the instance"""
        parser = argparse.ArgumentParser(
            prog="aws stop",
            description=help,
        )
        instances = self.get_instances(command='stop')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        args = parser.parse_args(argv)
        server = instances[args.server[0]]
        server.stop()

    def cmd_terminate(self, argv, help):
        """Terminates the instance"""
        parser = argparse.ArgumentParser(
            prog="aws terminate",
            description=help,
        )
        instances = self.get_instances(command='terminate')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        args = parser.parse_args(argv)
        server = instances[args.server[0]]
        server.terminate()

    def _parse_overrides(self, options):
        overrides = dict()
        if options.overrides is not None:
            for override in options.overrides:
                if '=' not in override:
                    log.error("Invalid format for override '%s', should be NAME=VALUE." % override)
                    sys.exit(1)
                key, value = override.split('=', 1)
                key = key.strip()
                value = value.strip()
                if key == '':
                    log.error("Empty key for override '%s'." % override)
                    sys.exit(1)
                overrides[key] = value
        return overrides

    def cmd_start(self, argv, help):
        """Starts the instance"""
        parser = argparse.ArgumentParser(
            prog="aws start",
            description=help,
        )
        instances = self.get_instances(command='start')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        parser.add_argument("-o", "--override", nargs="*", type=str,
                            dest="overrides", metavar="OVERRIDE",
                            help="Option to override in server config for startup script (name=value).")
        args = parser.parse_args(argv)
        overrides = self._parse_overrides(args)
        overrides['servers'] = self.instances
        server = instances[args.server[0]]
        instance = server.start(overrides)
        if instance is None:
            return
        server.status()

    def cmd_debug(self, argv, help):
        """Prints some debug info for this script"""
        parser = argparse.ArgumentParser(
            prog="aws debug",
            description=help,
        )
        instances = self.get_instances(command='startup_script')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        parser.add_argument("-v", "--verbose", dest="verbose",
                            action="store_true", help="Print more info and output the startup script")
        parser.add_argument("-i", "--interactive", dest="interactive",
                            action="store_true", help="Creates a connection and drops you into pdb")
        parser.add_argument("-r", "--raw", dest="raw",
                            action="store_true", help="Outputs the raw possibly compressed startup script")
        parser.add_argument("-o", "--override", nargs="*", type=str,
                            dest="overrides", metavar="OVERRIDE",
                            help="Option to override server config for startup script (name=value).")
        args = parser.parse_args(argv)
        overrides = self._parse_overrides(args)
        overrides['servers'] = self.instances
        server = instances[args.server[0]]
        startup_script = server.startup_script(overrides=overrides, debug=True)
        max_size = getattr(server, 'max_startup_script_size', 16 * 1024)
        log.info("Length of startup script: %s/%s", len(startup_script['raw']), max_size)
        if args.verbose:
            if 'startup_script' in server.config:
                if startup_script['original'] == startup_script['raw']:
                    log.info("Startup script:")
                elif args.raw:
                    log.info("Compressed startup script:")
                else:
                    log.info("Uncompressed startup script:")
            else:
                log.info("No startup script specified")
        if args.raw:
            print startup_script['raw'],
        elif args.verbose:
            print startup_script['original'],
        if args.interactive:  # pragma: no cover
            import readline
            conn = server.conn
            instance = server.instance
            conn, instance  # shutup pyflakes
            readline.parse_and_bind('tab: complete')
            local = locals()
            try:
                import rlcompleter
                readline.set_completer(rlcompleter.Completer(local).complete)
            except ImportError:
                pass
            __import__("code").interact(local=local)

    def cmd_do(self, argv, help):
        """Do stuff on the cluster (using fabric)"""
        parser = argparse.ArgumentParser(
            prog="aws do",
            description=help,
            add_help=False,
        )
        instances = self.get_instances(command='init_ssh_key')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance or server from the config.",
                            choices=list(instances))
        parser.add_argument("...", nargs=argparse.REMAINDER,
                            help="Fabric options")
        parser.parse_args(argv[:1])<commit_msg>Add option to get console output for an instance.<commit_after>import pkg_resources
pkg_resources.declare_namespace(__name__)

from lazy import lazy
from mr.awsome.config import Config
from mr.awsome import template
import logging
import argparse
import os
import sys


# shutup pyflakes
__all__ = [template.__name__]


log = logging.getLogger('mr.awsome')


class AWS(object):
    def __init__(self, configpath=None):
        log.setLevel(logging.INFO)
        ch = logging.StreamHandler()
        ch.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
        log.addHandler(ch)
        if configpath is None:
            configpath = 'etc/aws.conf'
        if os.path.isdir(configpath):
            configpath = os.path.join(configpath, 'aws.conf')
        self.configfile = configpath

    @lazy
    def config(self):
        configpath = os.path.abspath(self.configfile)
        if not os.path.exists(configpath):
            log.error("Config '%s' doesn't exist." % configpath)
            sys.exit(1)
        config = Config(configpath, bbb_config=True)
        config.parse()
        return config

    @lazy
    def masters(self):
        masters = []
        for plugin in self.config['plugin'].values():
            masters.extend(plugin['module'].get_masters(self.config))
        return masters

    @lazy
    def known_hosts(self):
        return os.path.join(self.config.path, 'known_hosts')

    def get_masters(self, command):
        masters = []
        for master in self.masters:
            if getattr(master, command, None) is not None:
                masters.append(master)
        return masters

    @lazy
    def instances(self):
        instances = {}
        for master in self.masters:
            for instance_id in master.instances:
                if instance_id in instances:
                    log.error("Instance and server names must be unique, '%s' is already defined." % instance_id)
                    sys.exit(1)
                instances[instance_id] = master.instances[instance_id]
        return instances

    def get_instances(self, command):
        instances = {}
        for instance_id in self.instances:
            instance = self.instances[instance_id]
            if getattr(instance, command, None) is not None:
                instances[instance_id] = instance
        return instances

    def cmd_status(self, argv, help):
        """Prints status"""
        parser = argparse.ArgumentParser(
            prog="aws status",
            description=help,
        )
        instances = self.get_instances(command='status')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        args = parser.parse_args(argv)
        server = instances[args.server[0]]
        server.status()

    def cmd_stop(self, argv, help):
        """Stops the instance"""
        parser = argparse.ArgumentParser(
            prog="aws stop",
            description=help,
        )
        instances = self.get_instances(command='stop')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        args = parser.parse_args(argv)
        server = instances[args.server[0]]
        server.stop()

    def cmd_terminate(self, argv, help):
        """Terminates the instance"""
        parser = argparse.ArgumentParser(
            prog="aws terminate",
            description=help,
        )
        instances = self.get_instances(command='terminate')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        args = parser.parse_args(argv)
        server = instances[args.server[0]]
        server.terminate()

    def _parse_overrides(self, options):
        overrides = dict()
        if options.overrides is not None:
            for override in options.overrides:
                if '=' not in override:
                    log.error("Invalid format for override '%s', should be NAME=VALUE." % override)
                    sys.exit(1)
                key, value = override.split('=', 1)
                key = key.strip()
                value = value.strip()
                if key == '':
                    log.error("Empty key for override '%s'." % override)
                    sys.exit(1)
                overrides[key] = value
        return overrides

    def cmd_start(self, argv, help):
        """Starts the instance"""
        parser = argparse.ArgumentParser(
            prog="aws start",
            description=help,
        )
        instances = self.get_instances(command='start')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        parser.add_argument("-o", "--override", nargs="*", type=str,
                            dest="overrides", metavar="OVERRIDE",
                            help="Option to override in server config for startup script (name=value).")
        args = parser.parse_args(argv)
        overrides = self._parse_overrides(args)
        overrides['servers'] = self.instances
        server = instances[args.server[0]]
        instance = server.start(overrides)
        if instance is None:
            return
        server.status()

    def cmd_debug(self, argv, help):
        """Prints some debug info for this script"""
        parser = argparse.ArgumentParser(
            prog="aws debug",
            description=help,
        )
        instances = self.get_instances(command='startup_script')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance from the config.",
                            choices=list(instances))
        parser.add_argument("-v", "--verbose", dest="verbose",
                            action="store_true", help="Print more info and output the startup script")
        parser.add_argument("-c", "--console-output", dest="console_output",
                            action="store_true", help="Prints the console output of the instance if available")
        parser.add_argument("-i", "--interactive", dest="interactive",
                            action="store_true", help="Creates a connection and drops you into pdb")
        parser.add_argument("-r", "--raw", dest="raw",
                            action="store_true", help="Outputs the raw possibly compressed startup script")
        parser.add_argument("-o", "--override", nargs="*", type=str,
                            dest="overrides", metavar="OVERRIDE",
                            help="Option to override server config for startup script (name=value).")
        args = parser.parse_args(argv)
        overrides = self._parse_overrides(args)
        overrides['servers'] = self.instances
        server = instances[args.server[0]]
        startup_script = server.startup_script(overrides=overrides, debug=True)
        max_size = getattr(server, 'max_startup_script_size', 16 * 1024)
        log.info("Length of startup script: %s/%s", len(startup_script['raw']), max_size)
        if args.verbose:
            if 'startup_script' in server.config:
                if startup_script['original'] == startup_script['raw']:
                    log.info("Startup script:")
                elif args.raw:
                    log.info("Compressed startup script:")
                else:
                    log.info("Uncompressed startup script:")
            else:
                log.info("No startup script specified")
        if args.raw:
            print startup_script['raw'],
        elif args.verbose:
            print startup_script['original'],
        if args.console_output:
            print server.instance.get_console_output().output
        if args.interactive:  # pragma: no cover
            import readline
            conn = server.conn
            instance = server.instance
            conn, instance  # shutup pyflakes
            readline.parse_and_bind('tab: complete')
            local = locals()
            try:
                import rlcompleter
                readline.set_completer(rlcompleter.Completer(local).complete)
            except ImportError:
                pass
            __import__("code").interact(local=local)

    def cmd_do(self, argv, help):
        """Do stuff on the cluster (using fabric)"""
        parser = argparse.ArgumentParser(
            prog="aws do",
            description=help,
            add_help=False,
        )
        instances = self.get_instances(command='init_ssh_key')
        parser.add_argument("server", nargs=1,
                            metavar="instance",
                            help="Name of the instance or server from the config.",
                            choices=list(instances))
        parser.add_argument("...", nargs=argparse.REMAINDER,
                            help="Fabric options")
        parser.parse_args(argv[:1])
{'fraction_non_alphanumeric': 0.10508089389102321, 'fraction_numerical': 0.002489042800714247, 'mean_word_length': 1.8173780487804878, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6161916', 'n_tokens_mistral': 4819, 'n_tokens_neox': 4759, 'n_words': 1136}