<commit_before>################################################################################
#
# Copyright (c) 2012 ObjectLabs Corporation
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
################################################################################

import pymongo
import sys
import time
from utils import pretty_json
from analyzer import QueryAnalyzer, ReportAggregation
from parsers import LogParser, ProfileParser
from datetime import datetime
from datetime import timedelta
import traceback
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

################################################################################
# Configuration
################################################################################

IGNORE_DBS = ['local', 'admin', 'config']
IGNORE_COLLECTIONS = [u'system.namespaces',
                      u'system.profile',
                      u'system.users',
                      u'system.indexes']
WATCH_INTERVAL_SECONDS = 3.0
WATCH_DISPLAY_REFRESH_SECONDS = 30.0
DEFAULT_PROFILE_LEVEL = pymongo.SLOW_ONLY

################################################################################
# Dex
#   Uses a QueryAnalyzer (with included LogParser) to analyze a MongoDB
#   query or logfile
################################################################################
class Dex:

    ############################################################################
    def __init__(self, db_uri, verbose, namespaces_list, slowms, check_indexes, timeout):
        self._check_indexes = check_indexes
        self._query_analyzer = QueryAnalyzer(check_indexes)
        self._db_uri = db_uri
        self._slowms = slowms
        self._verbose = verbose
        self._requested_namespaces = self._validate_namespaces(namespaces_list)
        self._recommendation_cache = []
        self._report = ReportAggregation()
        self._start_time = None
        self._timeout_time = None
        self._timeout = timeout
        self._run_stats = self._get_initial_run_stats()
        self._first_line = True

    ############################################################################
    def generate_query_report(self, db_uri, query, db_name, collection_name):
        """Analyzes a single query"""
        return self._query_analyzer.generate_query_report(db_uri,
                                                          query,
                                                          db_name,
                                                          collection_name)

    ############################################################################
    def _process_query(self, input, parser):
        self._run_stats['linesRead'] += 1

        line_time = parser.get_line_time(input)

        if line_time is not None:
            if ((self._run_stats['timeRange']['start'] is None) or
                (self._run_stats['timeRange']['start'] > line_time)):
                self._run_stats['timeRange']['start'] = line_time
            if ((self._run_stats['timeRange']['end'] is None) or
                (self._run_stats['timeRange']['end'] < line_time)):
                self._run_stats['timeRange']['end'] = line_time

        parsed = parser.parse(input)

        if parsed is not None:
            if parsed['supported']:
                self._run_stats['linesAnalyzed'] += 1
                namespace_tuple = self._tuplefy_namespace(parsed['ns'])
                # If the query is for a requested namespace ....
                if self._namespace_requested(parsed['ns']):
                    db_name = namespace_tuple[0]
                    collection_name = namespace_tuple[1]
                    query_report = None
                    if parsed['stats']['millis'] >= self._slowms:
                        try:
                            query_report = self.generate_query_report(self._db_uri,
                                                                      parsed,
                                                                      db_name,
                                                                      collection_name)
                        except Exception as e:
                            #print traceback.print_exc()
                            return 1
                    if query_report is not None:
                        if query_report['recommendation'] is not None:
                            self._run_stats['linesWithRecommendations'] += 1
                        self._report.add_query_occurrence(query_report)
            else:
                self._run_stats['unparsableLineInfo']['unparsableLines'] += 1
                self._run_stats['unparsableLineInfo']['unparsableLinesWithTime'] += 1
                self._run_stats['unparsableLineInfo']['unparsedTimeMillis'] += int(parsed['stats']['millis'])
                self._run_stats['unparsableLineInfo']['unparsedAvgTimeMillis'] = self._run_stats['unparsableLineInfo']['unparsedTimeMillis'] / self._run_stats['unparsableLineInfo']['unparsableLinesWithTime']
        else:
            self._run_stats['unparsableLineInfo']['unparsableLines'] += 1
            self._run_stats['unparsableLineInfo']['unparsableLinesWithoutTime'] += 1

    ############################################################################
    def analyze_profile(self):
        """Analyzes queries from a given log file"""
        profile_parser = ProfileParser()
        databases = self._get_requested_databases()
        connection = pymongo.MongoClient(self._db_uri,
                                         document_class=OrderedDict,
                                         read_preference=pymongo.ReadPreference.PRIMARY_PREFERRED)

        if databases == []:
            try:
                databases = connection.database_names()
            except:
                message = "Error: Could not list databases on server. Please " \
                          +         "check the auth components of your URI or provide " \
                          +         "a namespace filter with -n.\n"
                sys.stderr.write(message)
                databases = []

            for ignore_db in IGNORE_DBS:
                if ignore_db in databases:
                    databases.remove(ignore_db)

        for dataase in databases:

            db = connection[database]

            profile_entries = db['system.profile'].find()

            for profile_entry in profile_entries:
                self._process_query(profile_entry,
                                    profile_parser)

        self._output_aggregated_report(sys.stdout)

        return 0

    ############################################################################
    def watch_profile(self):
        """Analyzes queries from a given log file"""
        profile_parser = ProfileParser()
        databases = self._get_requested_databases()
        connection = pymongo.MongoClient(self._db_uri,
                                         document_class=OrderedDict,
                                         read_preference=pymongo.ReadPreference.PRIMARY_PREFERRED)
        enabled_profile = False

        if databases == []:
            try:
                databases = connection.database_names()
            except:
                message = "Error: Could not list databases on server. Please " \
                          +         "check the auth components of your URI.\n"
                sys.stderr.write(message)
                databases = []

            for ignore_db in IGNORE_DBS:
                if ignore_db in databases:
                    databases.remove(ignore_db)

        if len(databases) != 1:
            message = "Error: Please use namespaces (-n) to specify a single " \
                      +         "database for profile watching.\n"
            sys.stderr.write(message)
            return 1

        database = databases[0]
        db = connection[database]

        initial_profile_level = db.profiling_level()

        if initial_profile_level is pymongo.OFF:
            message = "Profile level currently 0. Dex is setting profile " \
                      +         "level 1. To run --watch at profile level 2, " \
                      +         "enable profile level 2 before running Dex.\n"
            sys.stderr.write(message)
            db.set_profiling_level(DEFAULT_PROFILE_LEVEL)

        output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        try:
            for profile_entry in self._tail_profile(db, WATCH_INTERVAL_SECONDS):
                self._process_query(profile_entry,
                                    profile_parser)
                if time.time() >= output_time:
                    self._output_aggregated_report(sys.stderr)
                    output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        except KeyboardInterrupt:
            sys.stderr.write("Interrupt received\n")
        finally:
            self._output_aggregated_report(sys.stdout)
            if initial_profile_level is pymongo.OFF:
                message = "Dex is resetting profile level to initial value " \
                          +         "of 0. You may wish to drop the system.profile " \
                          +         "collection.\n"
                sys.stderr.write(message)
                db.set_profiling_level(initial_profile_level)

        return 0

    ############################################################################
    def analyze_logfile(self, logfile_path):
        self._run_stats['logSource'] = logfile_path
        """Analyzes queries from a given log file"""
        with open(logfile_path) as obj:
            self.analyze_logfile_object(obj)

        self._output_aggregated_report(sys.stdout)

        return 0

    ############################################################################
    def analyze_logfile_object(self, file_object):
        """Analyzes queries from a given log file"""
        log_parser = LogParser()

        if self._start_time is None:
            self._start_time = datetime.now()
            if self._timeout != 0:
                self._end_time = self._start_time + timedelta(minutes=self._timeout)
            else:
                self._end_time = None

        # For each line in the logfile ...
        for line in file_object:
            if self._end_time is not None and datetime.now() > self._end_time:
                self._run_stats['timedOut'] = True
                self._run_stats['timeoutInMinutes'] = self._timeout
                break
            self._process_query(line, log_parser)

        return 0

    ############################################################################
    def watch_logfile(self, logfile_path):
        """Analyzes queries from the tail of a given log file"""
        self._run_stats['logSource'] = logfile_path
        log_parser = LogParser()

        # For each new line in the logfile ...
        output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        try:
            firstLine = True
            for line in self._tail_file(open(logfile_path),
                                        WATCH_INTERVAL_SECONDS):
                if firstLine:
                    self._run_stats['timeRange']['start'] = log_parser.get_line_time(line)
                self._process_query(line, log_parser)
                self._run_stats['timeRange']['end'] = log_parser.get_line_time(line)
                if time.time() >= output_time:
                    self._output_aggregated_report(sys.stderr)
                    output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        except KeyboardInterrupt:
            sys.stderr.write("Interrupt received\n")
        finally:
            self._output_aggregated_report(sys.stdout)

        return 0

    ############################################################################
    def _get_initial_run_stats(self):
        """Singlesource for initializing an output dict"""
        return OrderedDict([('linesWithRecommendations', 0),
                            ('linesAnalyzed', 0),
                            ('linesRead', 0),
                            ('dexTime', datetime.utcnow()),
                            ('logSource', None),
                            ('timeRange', OrderedDict([('start', None),
                                                       ('end', None)])),
                            ('unparsableLineInfo', OrderedDict([('unparsableLines', 0),
                                                                ('unparsableLinesWithoutTime', 0),
                                                                ('unparsableLinesWithTime', 0),
                                                                ('unparsedTimeMillis', 0),
                                                                ('unparsedAvgTimeMillis', 0)]))])

    ############################################################################
    def _make_aggregated_report(self):
        output = OrderedDict([('runStats', self._run_stats),
                              ('results', self._report.get_reports())])
        return output

    ############################################################################
    def _output_aggregated_report(self, out):
        out.write(pretty_json(self._make_aggregated_report()).replace('"', "'").replace("\\'", '"') + "\n")

    ############################################################################
    def _tail_file(self, file, interval):
        """Tails a file"""
        file.seek(0,2)
        while True:
            where = file.tell()
            line = file.readline()
            if not line:
                time.sleep(interval)
                file.seek(where)
            else:
                yield line

    ############################################################################
    def _tail_profile(self, db, interval):
        """Tails the system.profile collection"""
        latest_doc = None
        while latest_doc is None:
            time.sleep(interval)
            latest_doc = db['system.profile'].find_one()

        current_time = latest_doc['ts']

        while True:
            time.sleep(interval)
            cursor = db['system.profile'].find({'ts': {'$gte': current_time}}).sort('ts', pymongo.ASCENDING)
            for doc in cursor:
                current_time = doc['ts']
                yield doc

    ############################################################################
    def _tuplefy_namespace(self, namespace):
        """Converts a mongodb namespace to a db, collection tuple"""
        namespace_split = namespace.split('.', 1)
        if len(namespace_split) is 1:
            # we treat a single element as a collection name.
            # this also properly tuplefies '*'
            namespace_tuple = ('*', namespace_split[0])
        elif len(namespace_split) is 2:
            namespace_tuple = (namespace_split[0],namespace_split[1])
        else:
            return None
        return namespace_tuple

    ############################################################################
    # Need to add rejection of true regex attempts.
    def _validate_namespaces(self, input_namespaces):
        """Converts a list of db namespaces to a list of namespace tuples,
            supporting basic commandline wildcards"""
        output_namespaces = []
        if input_namespaces == []:
            return output_namespaces
        elif '*' in input_namespaces:
            if len(input_namespaces) > 1:
                warning = 'Warning: Multiple namespaces are '
                warning += 'ignored when one namespace is "*"\n'
                sys.stderr.write(warning)
            return output_namespaces
        else:
            for namespace in input_namespaces:
                if not isinstance(namespace, unicode):
                    namespace = unicode(namespace)
                namespace_tuple = self._tuplefy_namespace(namespace)
                if namespace_tuple is None:
                    warning = 'Warning: Invalid namespace ' + namespace
                    warning += ' will be ignored\n'
                    sys.stderr.write(warning)
                else:
                    if namespace_tuple not in output_namespaces:
                        output_namespaces.append(namespace_tuple)
                    else:
                        warning = 'Warning: Duplicate namespace ' + namespace
                        warning += ' will be ignored\n'
                        sys.stderr.write(warning)
        return output_namespaces

    ############################################################################
    def _namespace_requested(self, namespace):
        """Checks whether the requested_namespaces contain the provided
            namespace"""
        if namespace is None:
            return False
        namespace_tuple = self._tuplefy_namespace(namespace)
        if namespace_tuple[0] in IGNORE_DBS:
            return False
        elif namespace_tuple[1] in IGNORE_COLLECTIONS:
            return False
        else:
            return self._tuple_requested(namespace_tuple)

    ############################################################################
    def _tuple_requested(self, namespace_tuple):
        """Helper for _namespace_requested. Supports limited wildcards"""
        if not isinstance(namespace_tuple[0], unicode):
            encoded_db = unicode(namespace_tuple[0])
        else:
            encoded_db = namespace_tuple[0]
        if not isinstance(namespace_tuple[1], unicode):
            encoded_coll = unicode(namespace_tuple[1])
        else:
            encoded_coll = namespace_tuple[1]

        if namespace_tuple is None:
            return False
        elif len(self._requested_namespaces) is 0:
            return True
        for requested_namespace in self._requested_namespaces:
            if  ((((requested_namespace[0]) == u'*') or
                 (encoded_db == requested_namespace[0])) and
                (((requested_namespace[1]) == u'*') or
                 (encoded_coll == requested_namespace[1]))):
                return True
        return False

    ############################################################################
    def _get_requested_databases(self):
        """Returns a list of databases requested, not including ignored dbs"""
        requested_databases = []
        if ((self._requested_namespaces is not None) and
                (self._requested_namespaces != [])):
            for requested_namespace in self._requested_namespaces:
                if requested_namespace[0] is '*':
                    return []
                elif requested_namespace[0] not in IGNORE_DBS:
                    requested_databases.append(requested_namespace[0])
        return requested_databases
<commit_msg>typo fix<commit_after>################################################################################
#
# Copyright (c) 2012 ObjectLabs Corporation
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
################################################################################

import pymongo
import sys
import time
from utils import pretty_json
from analyzer import QueryAnalyzer, ReportAggregation
from parsers import LogParser, ProfileParser
from datetime import datetime
from datetime import timedelta
import traceback
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

################################################################################
# Configuration
################################################################################

IGNORE_DBS = ['local', 'admin', 'config']
IGNORE_COLLECTIONS = [u'system.namespaces',
                      u'system.profile',
                      u'system.users',
                      u'system.indexes']
WATCH_INTERVAL_SECONDS = 3.0
WATCH_DISPLAY_REFRESH_SECONDS = 30.0
DEFAULT_PROFILE_LEVEL = pymongo.SLOW_ONLY

################################################################################
# Dex
#   Uses a QueryAnalyzer (with included LogParser) to analyze a MongoDB
#   query or logfile
################################################################################
class Dex:

    ############################################################################
    def __init__(self, db_uri, verbose, namespaces_list, slowms, check_indexes, timeout):
        self._check_indexes = check_indexes
        self._query_analyzer = QueryAnalyzer(check_indexes)
        self._db_uri = db_uri
        self._slowms = slowms
        self._verbose = verbose
        self._requested_namespaces = self._validate_namespaces(namespaces_list)
        self._recommendation_cache = []
        self._report = ReportAggregation()
        self._start_time = None
        self._timeout_time = None
        self._timeout = timeout
        self._run_stats = self._get_initial_run_stats()
        self._first_line = True

    ############################################################################
    def generate_query_report(self, db_uri, query, db_name, collection_name):
        """Analyzes a single query"""
        return self._query_analyzer.generate_query_report(db_uri,
                                                          query,
                                                          db_name,
                                                          collection_name)

    ############################################################################
    def _process_query(self, input, parser):
        self._run_stats['linesRead'] += 1

        line_time = parser.get_line_time(input)

        if line_time is not None:
            if ((self._run_stats['timeRange']['start'] is None) or
                (self._run_stats['timeRange']['start'] > line_time)):
                self._run_stats['timeRange']['start'] = line_time
            if ((self._run_stats['timeRange']['end'] is None) or
                (self._run_stats['timeRange']['end'] < line_time)):
                self._run_stats['timeRange']['end'] = line_time

        parsed = parser.parse(input)

        if parsed is not None:
            if parsed['supported']:
                self._run_stats['linesAnalyzed'] += 1
                namespace_tuple = self._tuplefy_namespace(parsed['ns'])
                # If the query is for a requested namespace ....
                if self._namespace_requested(parsed['ns']):
                    db_name = namespace_tuple[0]
                    collection_name = namespace_tuple[1]
                    query_report = None
                    if parsed['stats']['millis'] >= self._slowms:
                        try:
                            query_report = self.generate_query_report(self._db_uri,
                                                                      parsed,
                                                                      db_name,
                                                                      collection_name)
                        except Exception as e:
                            #print traceback.print_exc()
                            return 1
                    if query_report is not None:
                        if query_report['recommendation'] is not None:
                            self._run_stats['linesWithRecommendations'] += 1
                        self._report.add_query_occurrence(query_report)
            else:
                self._run_stats['unparsableLineInfo']['unparsableLines'] += 1
                self._run_stats['unparsableLineInfo']['unparsableLinesWithTime'] += 1
                self._run_stats['unparsableLineInfo']['unparsedTimeMillis'] += int(parsed['stats']['millis'])
                self._run_stats['unparsableLineInfo']['unparsedAvgTimeMillis'] = self._run_stats['unparsableLineInfo']['unparsedTimeMillis'] / self._run_stats['unparsableLineInfo']['unparsableLinesWithTime']
        else:
            self._run_stats['unparsableLineInfo']['unparsableLines'] += 1
            self._run_stats['unparsableLineInfo']['unparsableLinesWithoutTime'] += 1

    ############################################################################
    def analyze_profile(self):
        """Analyzes queries from a given log file"""
        profile_parser = ProfileParser()
        databases = self._get_requested_databases()
        connection = pymongo.MongoClient(self._db_uri,
                                         document_class=OrderedDict,
                                         read_preference=pymongo.ReadPreference.PRIMARY_PREFERRED)

        if databases == []:
            try:
                databases = connection.database_names()
            except:
                message = "Error: Could not list databases on server. Please " \
                          +         "check the auth components of your URI or provide " \
                          +         "a namespace filter with -n.\n"
                sys.stderr.write(message)
                databases = []

            for ignore_db in IGNORE_DBS:
                if ignore_db in databases:
                    databases.remove(ignore_db)

        for database in databases:

            db = connection[database]

            profile_entries = db['system.profile'].find()

            for profile_entry in profile_entries:
                self._process_query(profile_entry,
                                    profile_parser)

        self._output_aggregated_report(sys.stdout)

        return 0

    ############################################################################
    def watch_profile(self):
        """Analyzes queries from a given log file"""
        profile_parser = ProfileParser()
        databases = self._get_requested_databases()
        connection = pymongo.MongoClient(self._db_uri,
                                         document_class=OrderedDict,
                                         read_preference=pymongo.ReadPreference.PRIMARY_PREFERRED)
        enabled_profile = False

        if databases == []:
            try:
                databases = connection.database_names()
            except:
                message = "Error: Could not list databases on server. Please " \
                          +         "check the auth components of your URI.\n"
                sys.stderr.write(message)
                databases = []

            for ignore_db in IGNORE_DBS:
                if ignore_db in databases:
                    databases.remove(ignore_db)

        if len(databases) != 1:
            message = "Error: Please use namespaces (-n) to specify a single " \
                      +         "database for profile watching.\n"
            sys.stderr.write(message)
            return 1

        database = databases[0]
        db = connection[database]

        initial_profile_level = db.profiling_level()

        if initial_profile_level is pymongo.OFF:
            message = "Profile level currently 0. Dex is setting profile " \
                      +         "level 1. To run --watch at profile level 2, " \
                      +         "enable profile level 2 before running Dex.\n"
            sys.stderr.write(message)
            db.set_profiling_level(DEFAULT_PROFILE_LEVEL)

        output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        try:
            for profile_entry in self._tail_profile(db, WATCH_INTERVAL_SECONDS):
                self._process_query(profile_entry,
                                    profile_parser)
                if time.time() >= output_time:
                    self._output_aggregated_report(sys.stderr)
                    output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        except KeyboardInterrupt:
            sys.stderr.write("Interrupt received\n")
        finally:
            self._output_aggregated_report(sys.stdout)
            if initial_profile_level is pymongo.OFF:
                message = "Dex is resetting profile level to initial value " \
                          +         "of 0. You may wish to drop the system.profile " \
                          +         "collection.\n"
                sys.stderr.write(message)
                db.set_profiling_level(initial_profile_level)

        return 0

    ############################################################################
    def analyze_logfile(self, logfile_path):
        self._run_stats['logSource'] = logfile_path
        """Analyzes queries from a given log file"""
        with open(logfile_path) as obj:
            self.analyze_logfile_object(obj)

        self._output_aggregated_report(sys.stdout)

        return 0

    ############################################################################
    def analyze_logfile_object(self, file_object):
        """Analyzes queries from a given log file"""
        log_parser = LogParser()

        if self._start_time is None:
            self._start_time = datetime.now()
            if self._timeout != 0:
                self._end_time = self._start_time + timedelta(minutes=self._timeout)
            else:
                self._end_time = None

        # For each line in the logfile ...
        for line in file_object:
            if self._end_time is not None and datetime.now() > self._end_time:
                self._run_stats['timedOut'] = True
                self._run_stats['timeoutInMinutes'] = self._timeout
                break
            self._process_query(line, log_parser)

        return 0

    ############################################################################
    def watch_logfile(self, logfile_path):
        """Analyzes queries from the tail of a given log file"""
        self._run_stats['logSource'] = logfile_path
        log_parser = LogParser()

        # For each new line in the logfile ...
        output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        try:
            firstLine = True
            for line in self._tail_file(open(logfile_path),
                                        WATCH_INTERVAL_SECONDS):
                if firstLine:
                    self._run_stats['timeRange']['start'] = log_parser.get_line_time(line)
                self._process_query(line, log_parser)
                self._run_stats['timeRange']['end'] = log_parser.get_line_time(line)
                if time.time() >= output_time:
                    self._output_aggregated_report(sys.stderr)
                    output_time = time.time() + WATCH_DISPLAY_REFRESH_SECONDS
        except KeyboardInterrupt:
            sys.stderr.write("Interrupt received\n")
        finally:
            self._output_aggregated_report(sys.stdout)

        return 0

    ############################################################################
    def _get_initial_run_stats(self):
        """Singlesource for initializing an output dict"""
        return OrderedDict([('linesWithRecommendations', 0),
                            ('linesAnalyzed', 0),
                            ('linesRead', 0),
                            ('dexTime', datetime.utcnow()),
                            ('logSource', None),
                            ('timeRange', OrderedDict([('start', None),
                                                       ('end', None)])),
                            ('unparsableLineInfo', OrderedDict([('unparsableLines', 0),
                                                                ('unparsableLinesWithoutTime', 0),
                                                                ('unparsableLinesWithTime', 0),
                                                                ('unparsedTimeMillis', 0),
                                                                ('unparsedAvgTimeMillis', 0)]))])

    ############################################################################
    def _make_aggregated_report(self):
        output = OrderedDict([('runStats', self._run_stats),
                              ('results', self._report.get_reports())])
        return output

    ############################################################################
    def _output_aggregated_report(self, out):
        out.write(pretty_json(self._make_aggregated_report()).replace('"', "'").replace("\\'", '"') + "\n")

    ############################################################################
    def _tail_file(self, file, interval):
        """Tails a file"""
        file.seek(0,2)
        while True:
            where = file.tell()
            line = file.readline()
            if not line:
                time.sleep(interval)
                file.seek(where)
            else:
                yield line

    ############################################################################
    def _tail_profile(self, db, interval):
        """Tails the system.profile collection"""
        latest_doc = None
        while latest_doc is None:
            time.sleep(interval)
            latest_doc = db['system.profile'].find_one()

        current_time = latest_doc['ts']

        while True:
            time.sleep(interval)
            cursor = db['system.profile'].find({'ts': {'$gte': current_time}}).sort('ts', pymongo.ASCENDING)
            for doc in cursor:
                current_time = doc['ts']
                yield doc

    ############################################################################
    def _tuplefy_namespace(self, namespace):
        """Converts a mongodb namespace to a db, collection tuple"""
        namespace_split = namespace.split('.', 1)
        if len(namespace_split) is 1:
            # we treat a single element as a collection name.
            # this also properly tuplefies '*'
            namespace_tuple = ('*', namespace_split[0])
        elif len(namespace_split) is 2:
            namespace_tuple = (namespace_split[0],namespace_split[1])
        else:
            return None
        return namespace_tuple

    ############################################################################
    # Need to add rejection of true regex attempts.
    def _validate_namespaces(self, input_namespaces):
        """Converts a list of db namespaces to a list of namespace tuples,
            supporting basic commandline wildcards"""
        output_namespaces = []
        if input_namespaces == []:
            return output_namespaces
        elif '*' in input_namespaces:
            if len(input_namespaces) > 1:
                warning = 'Warning: Multiple namespaces are '
                warning += 'ignored when one namespace is "*"\n'
                sys.stderr.write(warning)
            return output_namespaces
        else:
            for namespace in input_namespaces:
                if not isinstance(namespace, unicode):
                    namespace = unicode(namespace)
                namespace_tuple = self._tuplefy_namespace(namespace)
                if namespace_tuple is None:
                    warning = 'Warning: Invalid namespace ' + namespace
                    warning += ' will be ignored\n'
                    sys.stderr.write(warning)
                else:
                    if namespace_tuple not in output_namespaces:
                        output_namespaces.append(namespace_tuple)
                    else:
                        warning = 'Warning: Duplicate namespace ' + namespace
                        warning += ' will be ignored\n'
                        sys.stderr.write(warning)
        return output_namespaces

    ############################################################################
    def _namespace_requested(self, namespace):
        """Checks whether the requested_namespaces contain the provided
            namespace"""
        if namespace is None:
            return False
        namespace_tuple = self._tuplefy_namespace(namespace)
        if namespace_tuple[0] in IGNORE_DBS:
            return False
        elif namespace_tuple[1] in IGNORE_COLLECTIONS:
            return False
        else:
            return self._tuple_requested(namespace_tuple)

    ############################################################################
    def _tuple_requested(self, namespace_tuple):
        """Helper for _namespace_requested. Supports limited wildcards"""
        if not isinstance(namespace_tuple[0], unicode):
            encoded_db = unicode(namespace_tuple[0])
        else:
            encoded_db = namespace_tuple[0]
        if not isinstance(namespace_tuple[1], unicode):
            encoded_coll = unicode(namespace_tuple[1])
        else:
            encoded_coll = namespace_tuple[1]

        if namespace_tuple is None:
            return False
        elif len(self._requested_namespaces) is 0:
            return True
        for requested_namespace in self._requested_namespaces:
            if  ((((requested_namespace[0]) == u'*') or
                 (encoded_db == requested_namespace[0])) and
                (((requested_namespace[1]) == u'*') or
                 (encoded_coll == requested_namespace[1]))):
                return True
        return False

    ############################################################################
    def _get_requested_databases(self):
        """Returns a list of databases requested, not including ignored dbs"""
        requested_databases = []
        if ((self._requested_namespaces is not None) and
                (self._requested_namespaces != [])):
            for requested_namespace in self._requested_namespaces:
                if requested_namespace[0] is '*':
                    return []
                elif requested_namespace[0] not in IGNORE_DBS:
                    requested_databases.append(requested_namespace[0])
        return requested_databases
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# rsync-only - A wrapper around rsync for use with the SSH "command" restriction
#
# It's to be used with the SSH authorized_keys(8) "command" restriction. It
# allows you to restrict given user commands to rsync and its (configurable)
# safe options only.
#
# @TODO: Add support for allowing certain subset of rsync options only, most
#        notably: --delete and --server. A generic configurable solution would
#        be even nicer.
#

import os
import sys
import subprocess
from optparse import OptionParser

class RsyncWrapper:
    # Original rsync command given
    command = None

    # Location of the rsync binary
    rsync_bin = None

    def __init__(self, command, rsync_bin=["rsync", "/usr/bin/rsync"]):
        self.rsync_bin = rsync_bin

        if not self.is_rsync(command):
            raise StandardError("Invalid rsync command given: %s" % command)
        else:
            self.command = command

    def is_rsync(self, command):
        """Checks whether given command is indeed rsync"""
        for bin in self.rsync_bin:
            if command.startswith(bin):
                return True
        return False

    def execute(self):
        """Executes the rsync command and returns its return value"""
        return subprocess.call(self.command.split(' '))

def main():
    cli = OptionParser(usage='%prog [options]',
                       description='Limits execution of SSH commands to rsync only')
    cli.add_option('-g', '--debug', dest='debug', action='store_true',
                   help='Turn on debugging mode')

    (options, args) = cli.parse_args()

    try:
        # Require Python >= 2.4
        if sys.version_info[0] < 2 or sys.version_info[1] < 4:
            raise StandardError("Python 2.4.0 or higher is required")

        # If SSH_ORIGINAL_COMMAND is not specified, then the user is assuming
        # this is a normal login shell.
        if not os.environ.has_key('SSH_ORIGINAL_COMMAND'):
            raise StandardError("This is not a login shell. You can only use rsync.")

        rsync = RsyncWrapper(os.environ['SSH_ORIGINAL_COMMAND'])
        retval = rsync.execute()
        cli.exit(retval)

    except StandardError, e:
        if options.debug:
            raise
        else:
            cli.exit(1, "%s\n" % e)

if __name__ == "__main__":
    main()

<commit_msg>Added a @TODO for support being a login shell<commit_after>#!/usr/bin/env python
#
# rsync-only - A wrapper around rsync for use with the SSH "command" restriction
#
# It's to be used with the SSH authorized_keys(8) "command" restriction. It
# allows you to restrict given user commands to rsync and its (configurable)
# safe options only.
#
# @TODO: Add support for allowing certain subset of rsync options only, most
#        notably: --delete and --server. A generic configurable solution would
#        be even nicer.
#
# @TODO: Support being a login shell, which maybe nicer than the SSH pubkey
#        command restriction which is limited to SSH keys only.
#

import os
import sys
import subprocess
from optparse import OptionParser

class RsyncWrapper:
    # Original rsync command given
    command = None

    # Location of the rsync binary
    rsync_bin = None

    def __init__(self, command, rsync_bin=["rsync", "/usr/bin/rsync"]):
        self.rsync_bin = rsync_bin

        if not self.is_rsync(command):
            raise StandardError("Invalid rsync command given: %s" % command)
        else:
            self.command = command

    def is_rsync(self, command):
        """Checks whether given command is indeed rsync"""
        for bin in self.rsync_bin:
            if command.startswith(bin):
                return True
        return False

    def execute(self):
        """Executes the rsync command and returns its return value"""
        return subprocess.call(self.command.split(' '))

def main():
    cli = OptionParser(usage='%prog [options]',
                       description='Limits execution of SSH commands to rsync only')
    cli.add_option('-g', '--debug', dest='debug', action='store_true',
                   help='Turn on debugging mode')

    (options, args) = cli.parse_args()

    try:
        # Require Python >= 2.4
        if sys.version_info[0] < 2 or sys.version_info[1] < 4:
            raise StandardError("Python 2.4.0 or higher is required")

        # If SSH_ORIGINAL_COMMAND is not specified, then the user is assuming
        # this is a normal login shell.
        if not os.environ.has_key('SSH_ORIGINAL_COMMAND'):
            raise StandardError("This is not a login shell. You can only use rsync.")

        rsync = RsyncWrapper(os.environ['SSH_ORIGINAL_COMMAND'])
        retval = rsync.execute()
        cli.exit(retval)

    except StandardError, e:
        if options.debug:
            raise
        else:
            cli.exit(1, "%s\n" % e)

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from airflow.exceptions import AirflowException
from airflow.utils.decorators import apply_defaults
from airflow.contrib.hooks.aws_athena_hook import AWSAthenaHook
from airflow.sensors.base_sensor_operator import BaseSensorOperator

class AthenaSensor(BaseSensorOperator):
    """
    Asks for the state of the Query until it reaches a failure state or success state.
    If it fails, failing the task.

    :param query_execution_id: query_execution_id to check the state of
    :type query_execution_id: str
    :param max_retires: Number of times to poll for query state before
        returning the current state, defaults to None
    :type max_retires: int
    :param aws_conn_id: aws connection to use, defaults to 'aws_default'
    :type aws_conn_id: str
    :param sleep_time: Time to wait between two consecutive call to
        check query status on athena, defaults to 10
    :type sleep_time: int
    """

    INTERMEDIATE_STATES = ('QUEUED', 'RUNNING',)
    FAILURE_STATES = ('FAILED', 'CANCELLED',)
    SUCCESS_STATES = ('SUCCEEDED',)

    template_fields = ['query_execution_id']
    template_ext = ()
    ui_color = '#66c3ff'

    @apply_defaults
    def __init__(self,
                 query_execution_id,
                 max_retires=None,
                 aws_conn_id='aws_default',
                 sleep_time=10,
                 *args, **kwargs):
        super(BaseSensorOperator, self).__init__(*args, **kwargs)
        self.aws_conn_id = aws_conn_id
        self.query_execution_id = query_execution_id
        self.hook = None
        self.sleep_time = sleep_time
        self.max_retires = max_retires

    def poke(self, context):
        self.hook = self.get_hook()
        self.hook.get_conn()
        state = self.hook.poll_query_status(self.query_execution_id, self.max_retires)

        if state in self.FAILURE_STATES:
            raise AirflowException('Athena sensor failed')

        if state in self.INTERMEDIATE_STATES:
            return False
        return True

    def get_hook(self):
        return AWSAthenaHook(self.aws_conn_id, self.sleep_time)
<commit_msg>[AIRFLOW-4019] Fix AWS Athena Sensor object has no attribute 'mode' (#4844)<commit_after># -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from airflow.exceptions import AirflowException
from airflow.utils.decorators import apply_defaults
from airflow.contrib.hooks.aws_athena_hook import AWSAthenaHook
from airflow.sensors.base_sensor_operator import BaseSensorOperator

class AthenaSensor(BaseSensorOperator):
    """
    Asks for the state of the Query until it reaches a failure state or success state.
    If it fails, failing the task.

    :param query_execution_id: query_execution_id to check the state of
    :type query_execution_id: str
    :param max_retires: Number of times to poll for query state before
        returning the current state, defaults to None
    :type max_retires: int
    :param aws_conn_id: aws connection to use, defaults to 'aws_default'
    :type aws_conn_id: str
    :param sleep_time: Time to wait between two consecutive call to
        check query status on athena, defaults to 10
    :type sleep_time: int
    """

    INTERMEDIATE_STATES = ('QUEUED', 'RUNNING',)
    FAILURE_STATES = ('FAILED', 'CANCELLED',)
    SUCCESS_STATES = ('SUCCEEDED',)

    template_fields = ['query_execution_id']
    template_ext = ()
    ui_color = '#66c3ff'

    @apply_defaults
    def __init__(self,
                 query_execution_id,
                 max_retires=None,
                 aws_conn_id='aws_default',
                 sleep_time=10,
                 *args, **kwargs):
        super(AthenaSensor, self).__init__(*args, **kwargs)
        self.aws_conn_id = aws_conn_id
        self.query_execution_id = query_execution_id
        self.hook = None
        self.sleep_time = sleep_time
        self.max_retires = max_retires

    def poke(self, context):
        self.hook = self.get_hook()
        self.hook.get_conn()
        state = self.hook.poll_query_status(self.query_execution_id, self.max_retires)

        if state in self.FAILURE_STATES:
            raise AirflowException('Athena sensor failed')

        if state in self.INTERMEDIATE_STATES:
            return False
        return True

    def get_hook(self):
        return AWSAthenaHook(self.aws_conn_id, self.sleep_time)
<|endoftext|>


<commit_before>"""A tool for adding a new cluster / constellation from photographs."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cv2

def setup_parser(parser):

    parser.add_argument(
        'context',
        type=str,
        default=None,
        help="Path to the context image to add.")

    parser.add_argument(
        'detail',
        type=str,
        default=None,
        help="Path to the detail image to add.")

    parser.add_argument(
        'destination',
        type=str,
        default=None,
        help="New path to create and store the constellation to.")

    parser.add_argument(
        'moles',
        type=str,
        default=None,
        nargs='+',
        help="Names of the moles to store.")

def process_args(args):
    context_image = cv2.imread(args.context)
    detail_image = cv2.imread(args.detail)
    print('{}: {}'.format(args.context, context_image.shape))
    print('{}: {}'.format(args.detail, detail_image.shape))
    raise NotImplementedError()
<commit_msg>add-cluster: doc, comment print block<commit_after>"""A tool for adding a new cluster / constellation from photographs."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cv2

def setup_parser(parser):

    parser.add_argument(
        'context',
        type=str,
        default=None,
        help="Path to the context image to add.")

    parser.add_argument(
        'detail',
        type=str,
        default=None,
        help="Path to the detail image to add.")

    parser.add_argument(
        'destination',
        type=str,
        default=None,
        help="New path to create and store the constellation to.")

    parser.add_argument(
        'moles',
        type=str,
        default=None,
        nargs='+',
        help="Names of the moles to store.")

def process_args(args):
    context_image = cv2.imread(args.context)
    detail_image = cv2.imread(args.detail)

    # print out the dimensions of the images
    print('{}: {}'.format(args.context, context_image.shape))
    print('{}: {}'.format(args.detail, detail_image.shape))
    raise NotImplementedError()
<|endoftext|>


<commit_before>import sys
import os
from spyvm import model, interpreter, primitives, shadow
# from spyvm import classtable
# from spyvm.test.test_interpreter import *
from spyvm import squeakimage
from spyvm import constants

def tinyBenchmarks(space, image_name):
    interp = interpreter.Interpreter(space, image_name=image_name)

    w_object = model.W_SmallInteger(0)

    s_class = w_object.shadow_of_my_class(space)
    w_method = s_class.lookup("tinyBenchmarks")

    assert w_method
    w_frame = w_method.create_frame(space, w_object, [])
    interp.store_w_active_context(w_frame)

    counter = 0

    from spyvm.interpreter import BYTECODE_TABLE
    return interp

def run_benchmarks(interp):
    counter = 0
    try:
        interp.loop()
    except interpreter.ReturnFromTopLevel, e:
        w_result = e.object

        assert isinstance(w_result, model.W_BytesObject)
        print w_result.as_string()
        return 0
    return -1

from spyvm import objspace
space = objspace.ObjSpace()

def entry_point(argv):
    if len(argv) > 1:
        filename = argv[1]
    else:
        print "usage:", argv[0], "<image name>"
        return -1
    reader = squeakimage.reader_for_image(space, squeakimage.Stream(DummyFile(filename)))
    reader.initialize()
    image = squeakimage.SqueakImage()
    image.from_reader(space, reader)
    interp = tinyBenchmarks(space, filename)
    run_benchmarks(interp)
    return 0

# _____ Define and setup target ___

def target(*args):
    return entry_point, None

class DummyFile:
    def __init__(self,filename):
        import os
        fd = os.open(filename, os.O_RDONLY, 0777)
        try:
            content = []
            while 1:
                s = os.read(fd, 4096)
                if not s:
                    break
                content.append(s)
            self.content = "".join(content)
        finally:
            os.close(fd)
    def read(self):
        return self.content
    def close(self):
        pass

if __name__ == "__main__":
    entry_point(sys.argv)
<commit_msg>make targetimageloadingsmalltalk work again, and make it possible to run other methods on int<commit_after>import sys
import os
from spyvm import model, interpreter, primitives, shadow
# from spyvm import classtable
# from spyvm.test.test_interpreter import *
from spyvm import squeakimage
from spyvm import constants

def tinyBenchmarks(space, image_name):
    interp = interpreter.Interpreter(space, image_name=image_name)
    return interp

def run_benchmarks(interp, number, benchmark):
    counter = 0
    w_object = model.W_SmallInteger(number)
    try:
        interp.perform(w_object, "tinyBenchmarks")
    except interpreter.ReturnFromTopLevel, e:
        w_result = e.object

        assert isinstance(w_result, model.W_BytesObject)
        print w_result.as_string()
        return 0
    return -1

from spyvm import objspace
space = objspace.ObjSpace()

def entry_point(argv):
    if len(argv) > 1:
        filename = argv[1]
        if len(argv) > 3:
            number = int(argv[2])
            benchmark = argv[3]
        else:
            number = 0
            benchmark = "tinyBenchmarks"
    else:
        print "usage:", argv[0], "<image name>"
        return -1
    reader = squeakimage.reader_for_image(space, squeakimage.Stream(DummyFile(filename)))
    reader.initialize()
    image = squeakimage.SqueakImage()
    image.from_reader(space, reader)
    interp = interpreter.Interpreter(space, image, filename)
    run_benchmarks(interp, number, benchmark)
    return 0

# _____ Define and setup target ___

def target(*args):
    return entry_point, None

def jitpolicy(self):
    from rpython.jit.codewriter.policy import JitPolicy
    return JitPolicy()

class DummyFile:
    def __init__(self,filename):
        import os
        fd = os.open(filename, os.O_RDONLY, 0777)
        try:
            content = []
            while 1:
                s = os.read(fd, 4096)
                if not s:
                    break
                content.append(s)
            self.content = "".join(content)
        finally:
            os.close(fd)
    def read(self):
        return self.content
    def close(self):
        pass

if __name__ == "__main__":
    entry_point(sys.argv)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import functools
import itertools
import numpy as np
import scipy.signal
from multiprocessing import cpu_count
from multiprocessing import Pool

from pyannote.core import Segment, Timeline, Annotation
from pyannote.core.util import pairwise

from pyannote.database.util import get_annotated

def helper_peak_tune(item_prediction, peak=None, metric=None):
    """Apply peak detection on prediction and evaluate the result

    Parameters
    ----------
    item : dict
        Protocol item.
    prediction : SlidingWindowFeature
    peak : Peak, optional
    metric : DiarizationPurityCoverageFMeasure, optional

    Returns
    -------
    fscore : float
    """

    current_file, predictions = item_prediction

    reference = current_file['annotation']

    hypothesis = peak.apply(predictions).to_annotation()
    # remove (reference) non-speech regions
    hypothesis = hypothesis.crop(reference.get_timeline().support(),
                                 mode='intersection')

    uem = get_annotated(current_file)

    # TODO this might take a very long time when hypothesis contains
    # a large number of segments -- find a way to bypass this call
    # and make sure the internal components are accumulated accordingly
    return metric(reference, hypothesis, uem=uem)

class Peak(object):
    """Peak detection

    Peaks are detected on min_duration windows.
    Only peaks greater than p(1) + alpha * (p(99) - p(1)) are kept
    where p(x) is xth percentile.

    Parameters
    ----------
    alpha : float, optional
        Adaptative threshold coefficient. Defaults to 0.5
    percentile : bool, optional
        When False (default), alpha = 0 correspond to min and = 1 to max
        When True, alpha = 0 corresponds to p(1) and = 1 to p(99).
    min_duration : float, optional
        Defaults to 1 second.

    """
    def __init__(self, alpha=0.5, min_duration=1.0, percentile=False):
        super(Peak, self).__init__()
        self.alpha = alpha
        self.percentile = percentile
        self.min_duration = min_duration

    @classmethod
    def tune(cls, items, get_prediction, purity=0.95,
             n_calls=20, n_random_starts=10, n_jobs=-1):
        """Find best set of hyper-parameters using skopt

        Parameters
        ----------
        items : iterable
            Protocol items used as development set. Typically, one would use
            items = protocol.development()
        get_prediction : callable
            Callable that takes an item as input, and returns a prediction as
            a SlidingWindowFeature instances.
        purity : float, optional
            Target purity. Defaults to 0.95.
        n_calls : int, optional
            Number of trials for hyper-parameter optimization. Defaults to 20.
        n_random_starts : int, optional
            Number of trials with random initialization before being smart.
            Defaults to 10.
        n_jobs : int, optional
            Number of parallel job to use. Set to 1 to not use multithreading.
            Defaults to whichever is minimum between number of CPUs and number
            of items.

        Returns
        -------
        best_parameters : dict
            Best set of parameters.
        best_coverage : float
            Best coverage
        """

        import skopt
        import skopt.space

        from pyannote.metrics.diarization import DiarizationPurityCoverageFMeasure

        # make sure items can be iterated over and over again
        items = list(items)

        # defaults to whichever is minimum between
        # number of CPUs and number of items
        if n_jobs < 0:
            n_jobs = min(cpu_count(), len(items))

        if n_jobs > 1:
            pool = Pool(n_jobs)

        # compute predictions once and for all
        # NOTE could multithreading speed things up? this is unclear as
        # get_prediction is probably already multithreaded, or (better) even
        # precomputed
        predictions = [get_prediction(item) for item in items]

        def objective_function(params):
            alpha, min_duration, = params
            peak = cls(alpha=alpha, min_duration=min_duration)

            metric = DiarizationPurityCoverageFMeasure()
            process_one_file = functools.partial(
                helper_peak_tune, peak=peak, metric=metric)

            if n_jobs > 1:
                results = list(pool.map(process_one_file,
                                        zip(items, predictions)))
            else:
                results = [process_one_file(item_prediction)
                           for item_prediction in zip(items, predictions)]

            p, c, f = metric.compute_metrics()

            if p < purity:
                return 1.
            else:
                return 1. - c

        # 0 < alpha < 1 || 0 < min_duration < 5s
        space = [skopt.space.Real(0., 1., prior='uniform'),
                 skopt.space.Real(0., 5., prior='uniform')]

        res = skopt.gp_minimize(
            objective_function, space, x0=[[0.5, 2.]],
            n_calls=n_calls, n_random_starts=n_random_starts,
            random_state=1337, verbose=False)

        if n_jobs > 1:
            pool.terminate()

        best_coverage = 1. - res.fun
        return {'alpha': res.x[0], 'min_duration': res.x[1]}, best_coverage

    def apply(self, predictions, dimension=0):
        """Peak detection

        Parameter
        ---------
        predictions : SlidingWindowFeature
            Predictions returned by segmentation approaches.

        Returns
        -------
        segmentation : Timeline
            Partition.
        """

        if len(predictions.data.shape) == 1:
            y = predictions.data
        elif predictions.data.shape[1] == 1:
            y = predictions.data[:, 0]
        else:
            y = predictions.data[:, dimension]

        sw = predictions.sliding_window

        precision = sw.step
        order = max(1, int(np.rint(self.min_duration / precision)))
        indices = scipy.signal.argrelmax(y, order=order)[0]

        mini = np.nanpercentile(y, 1) if self.percentile else np.nanmin(y)
        maxi = np.nanpercentile(y, 99) if self.percentile else np.nanmax(y)
        threshold = mini + self.alpha * (maxi - mini)

        peak_time = np.array([sw[i].middle for i in indices if y[i] > threshold])

        n_windows = len(y)
        start_time = sw[0].start
        end_time = sw[n_windows].end

        boundaries = np.hstack([[start_time], peak_time, [end_time]])
        segmentation = Timeline()
        for i, (start, end) in enumerate(pairwise(boundaries)):
            segment = Segment(start, end)
            segmentation.add(segment)

        return segmentation

def helper_binarize_tune(item_prediction,
                         binarizer=None, metric=None, **kwargs):
    """Apply binarizer on prediction and evaluate the result

    Parameters
    ----------
    item : dict
        Protocol item.
    prediction : SlidingWindowFeature
    binarizer : Binarize, optional
    metric : BaseMetric, optional
    **kwargs : dict
        Passed to binarizer.apply()

    Returns
    -------
    value : float
        Metric value
    """

    item, prediction = item_prediction

    uem = get_annotated(item)
    reference = item['annotation']
    hypothesis_timeline = binarizer.apply(prediction, **kwargs)

    # convert from timeline to annnotation
    hypothesis = Annotation(uri=hypothesis_timeline.uri)
    for s, segment in enumerate(hypothesis_timeline):
        hypothesis[segment] = s

    result = metric(reference, hypothesis, uem=uem)

    return result

class Binarize(object):
    """Binarize predictions using onset/offset thresholding

    Parameters
    ----------
    onset : float, optional
        Defaults to 0.5.
    offset : float, optional
        Defaults to 0.5.

    Reference
    ---------
    Gregory Gelly and Jean-Luc Gauvain. "Minimum Word Error Training of
    RNN-based Voice Activity Detection", InterSpeech 2015.
    """

    def __init__(self, onset=0.7, offset=0.7):

        super(Binarize, self).__init__()
        self.onset = onset
        self.offset = offset

        self.min_duration = [0., 0.]
        self.pad_onset = 0.
        self.pad_offset = 0.

    @classmethod
    def tune(cls, items, get_prediction, get_metric=None,
             n_calls=20, n_random_starts=10, n_jobs=-1,
             **kwargs):
        """Find best set of hyper-parameters using skopt

        Parameters
        ----------
        items : iterable
            Protocol items used as development set. Typically, one would use
            items = protocol.development()
        get_prediction : callable
            Callable that takes an item as input, and returns a prediction as
            a SlidingWindowFeature instances.
        get_metric : callable, optional
            Callable that takes no input, and returns a fresh evaluation
            metric. Defaults to pyannote.metrics.detection.DetectionErrorRate.
        n_calls : int, optional
            Number of trials for hyper-parameter optimization. Defaults to 20.
        n_random_starts : int, optional
            Number of trials with random initialization before being smart.
            Defaults to 10.
        n_jobs : int, optional
            Number of parallel job to use. Set to 1 to not use multithreading.
            Defaults to whichever is minimum between number of CPUs and number
            of items.
        **kwargs :
            Optional keyword arguments passed to Binarize.apply().

        Returns
        -------
        best_parameters : dict
            Best set of parameters.
        best_metric : float
            Best metric
        """

        import skopt
        import skopt.space

        if get_metric is None:
            from pyannote.metrics.detection import DetectionErrorRate
            get_metric = DetectionErrorRate

        # make sure items can be iterated over and over again
        items = list(items)

        # defaults to whichever is minimum between
        # number of CPUs and number of items
        if n_jobs < 0:
            n_jobs = min(cpu_count(), len(items))

        if n_jobs > 1:
            pool = Pool(n_jobs)

        # compute predictions once and for all
        # NOTE could multithreading speed things up? this is unclear as
        # get_prediction is probably already multithreaded, or (better) even
        # precomputed
        predictions = [get_prediction(item) for item in items]

        def objective_function(params):
            onset, offset, = params
            binarizer = cls(onset=onset, offset=offset)

            metric = get_metric()
            process_one_file = functools.partial(helper_binarize_tune,
                                                 binarizer=binarizer,
                                                 metric=metric,
                                                 **kwargs)

            if n_jobs > 1:
                results = list(pool.map(process_one_file,
                                        zip(items, predictions)))
            else:
                results = [process_one_file(item_prediction)
                           for item_prediction in zip(items, predictions)]

            return abs(metric)

        space = [skopt.space.Real(0., 1., prior='uniform'),
                 skopt.space.Real(0., 1., prior='uniform')]

        res = skopt.gp_minimize(
            objective_function, space,
            x0=[[0.25, 0.25],[0.5, 0.5], [0.75, 0.75]],
            n_calls=n_calls, n_random_starts=n_random_starts,
            random_state=1337, verbose=False)

        if n_jobs > 1:
            pool.terminate()

        return {'onset': res.x[0], 'offset': res.x[1]}, res.fun

    def apply(self, predictions, dimension=0):
        """
        Parameters
        ----------
        predictions : SlidingWindowFeature
            Must be mono-dimensional
        dimension : int, optional
            Which dimension to process
        """

        if len(predictions.data.shape) == 1:
            data = predictions.data
        elif predictions.data.shape[1] == 1:
            data = predictions.data[:, 0]
        else:
            data = predictions.data[:, dimension]

        n_samples = predictions.getNumber()
        window = predictions.sliding_window
        timestamps = [window[i].middle for i in range(n_samples)]

        # initial state
        start = timestamps[0]
        label = data[0] > self.onset

        # timeline meant to store 'active' segments
        active = Timeline()

        for t, y in zip(timestamps[1:], data[1:]):

            # currently active
            if label:
                # switching from active to inactive
                if y < self.offset:
                    segment = Segment(start - self.pad_onset,
                                      t + self.pad_offset)
                    active.add(segment)
                    start = t
                    label = False

            # currently inactive
            else:
                # switching from inactive to active
                if y > self.onset:
                    start = t
                    label = True

        # if active at the end, add final segment
        if label:
            segment = Segment(start - self.pad_onset, t + self.pad_offset)
            active.add(segment)

        # because of padding, some 'active' segments might be overlapping
        # therefore, we merge those overlapping segments
        active = active.support()

        # remove short 'active' segments
        active = Timeline(
            [s for s in active if s.duration > self.min_duration[1]])

        # fill short 'inactive' segments
        inactive = active.gaps()
        for s in inactive:
            if s.duration < self.min_duration[0]:
                active.add(s)
        active = active.support()

        return active
<commit_msg>feat: add log_scale boolean flag to Peak detection<commit_after>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import functools
import itertools
import numpy as np
import scipy.signal
from multiprocessing import cpu_count
from multiprocessing import Pool

from pyannote.core import Segment, Timeline, Annotation
from pyannote.core.util import pairwise

from pyannote.database.util import get_annotated

def helper_peak_tune(item_prediction, peak=None, metric=None):
    """Apply peak detection on prediction and evaluate the result

    Parameters
    ----------
    item : dict
        Protocol item.
    prediction : SlidingWindowFeature
    peak : Peak, optional
    metric : DiarizationPurityCoverageFMeasure, optional

    Returns
    -------
    fscore : float
    """

    current_file, predictions = item_prediction

    reference = current_file['annotation']

    hypothesis = peak.apply(predictions).to_annotation()
    # remove (reference) non-speech regions
    hypothesis = hypothesis.crop(reference.get_timeline().support(),
                                 mode='intersection')

    uem = get_annotated(current_file)

    # TODO this might take a very long time when hypothesis contains
    # a large number of segments -- find a way to bypass this call
    # and make sure the internal components are accumulated accordingly
    return metric(reference, hypothesis, uem=uem)

class Peak(object):
    """Peak detection

    Peaks are detected on min_duration windows.
    Only peaks greater than p(1) + alpha * (p(99) - p(1)) are kept
    where p(x) is xth percentile.

    Parameters
    ----------
    alpha : float, optional
        Adaptative threshold coefficient. Defaults to 0.5
    percentile : bool, optional
        When False (default), alpha = 0 correspond to min and = 1 to max
        When True, alpha = 0 corresponds to p(1) and = 1 to p(99).
    min_duration : float, optional
        Defaults to 1 second.
    log_scale : bool, optional
        Set to True to indicate that binarized scores are log scaled.
        Defaults to False.

    """
    def __init__(self, alpha=0.5, min_duration=1.0, percentile=False,
                 log_scale=False):
        super(Peak, self).__init__()
        self.alpha = alpha
        self.percentile = percentile
        self.min_duration = min_duration
        self.log_scale = log_scale

    @classmethod
    def tune(cls, items, get_prediction, purity=0.95,
             n_calls=20, n_random_starts=10, n_jobs=-1):
        """Find best set of hyper-parameters using skopt

        Parameters
        ----------
        items : iterable
            Protocol items used as development set. Typically, one would use
            items = protocol.development()
        get_prediction : callable
            Callable that takes an item as input, and returns a prediction as
            a SlidingWindowFeature instances.
        purity : float, optional
            Target purity. Defaults to 0.95.
        n_calls : int, optional
            Number of trials for hyper-parameter optimization. Defaults to 20.
        n_random_starts : int, optional
            Number of trials with random initialization before being smart.
            Defaults to 10.
        n_jobs : int, optional
            Number of parallel job to use. Set to 1 to not use multithreading.
            Defaults to whichever is minimum between number of CPUs and number
            of items.

        Returns
        -------
        best_parameters : dict
            Best set of parameters.
        best_coverage : float
            Best coverage
        """

        import skopt
        import skopt.space

        from pyannote.metrics.diarization import DiarizationPurityCoverageFMeasure

        # make sure items can be iterated over and over again
        items = list(items)

        # defaults to whichever is minimum between
        # number of CPUs and number of items
        if n_jobs < 0:
            n_jobs = min(cpu_count(), len(items))

        if n_jobs > 1:
            pool = Pool(n_jobs)

        # compute predictions once and for all
        # NOTE could multithreading speed things up? this is unclear as
        # get_prediction is probably already multithreaded, or (better) even
        # precomputed
        predictions = [get_prediction(item) for item in items]

        def objective_function(params):
            alpha, min_duration, = params
            peak = cls(alpha=alpha, min_duration=min_duration)

            metric = DiarizationPurityCoverageFMeasure()
            process_one_file = functools.partial(
                helper_peak_tune, peak=peak, metric=metric)

            if n_jobs > 1:
                results = list(pool.map(process_one_file,
                                        zip(items, predictions)))
            else:
                results = [process_one_file(item_prediction)
                           for item_prediction in zip(items, predictions)]

            p, c, f = metric.compute_metrics()

            if p < purity:
                return 1.
            else:
                return 1. - c

        # 0 < alpha < 1 || 0 < min_duration < 5s
        space = [skopt.space.Real(0., 1., prior='uniform'),
                 skopt.space.Real(0., 5., prior='uniform')]

        res = skopt.gp_minimize(
            objective_function, space, x0=[[0.5, 2.]],
            n_calls=n_calls, n_random_starts=n_random_starts,
            random_state=1337, verbose=False)

        if n_jobs > 1:
            pool.terminate()

        best_coverage = 1. - res.fun
        return {'alpha': res.x[0], 'min_duration': res.x[1]}, best_coverage

    def apply(self, predictions, dimension=0):
        """Peak detection

        Parameter
        ---------
        predictions : SlidingWindowFeature
            Predictions returned by segmentation approaches.

        Returns
        -------
        segmentation : Timeline
            Partition.
        """

        if len(predictions.data.shape) == 1:
            y = predictions.data
        elif predictions.data.shape[1] == 1:
            y = predictions.data[:, 0]
        else:
            y = predictions.data[:, dimension]

        if self.log_scale:
            y = np.exp(y)

        sw = predictions.sliding_window

        precision = sw.step
        order = max(1, int(np.rint(self.min_duration / precision)))
        indices = scipy.signal.argrelmax(y, order=order)[0]

        mini = np.nanpercentile(y, 1) if self.percentile else np.nanmin(y)
        maxi = np.nanpercentile(y, 99) if self.percentile else np.nanmax(y)
        threshold = mini + self.alpha * (maxi - mini)

        peak_time = np.array([sw[i].middle for i in indices if y[i] > threshold])

        n_windows = len(y)
        start_time = sw[0].start
        end_time = sw[n_windows].end

        boundaries = np.hstack([[start_time], peak_time, [end_time]])
        segmentation = Timeline()
        for i, (start, end) in enumerate(pairwise(boundaries)):
            segment = Segment(start, end)
            segmentation.add(segment)

        return segmentation

def helper_binarize_tune(item_prediction,
                         binarizer=None, metric=None, **kwargs):
    """Apply binarizer on prediction and evaluate the result

    Parameters
    ----------
    item : dict
        Protocol item.
    prediction : SlidingWindowFeature
    binarizer : Binarize, optional
    metric : BaseMetric, optional
    **kwargs : dict
        Passed to binarizer.apply()

    Returns
    -------
    value : float
        Metric value
    """

    item, prediction = item_prediction

    uem = get_annotated(item)
    reference = item['annotation']
    hypothesis_timeline = binarizer.apply(prediction, **kwargs)

    # convert from timeline to annnotation
    hypothesis = Annotation(uri=hypothesis_timeline.uri)
    for s, segment in enumerate(hypothesis_timeline):
        hypothesis[segment] = s

    result = metric(reference, hypothesis, uem=uem)

    return result

class Binarize(object):
    """Binarize predictions using onset/offset thresholding

    Parameters
    ----------
    onset : float, optional
        Defaults to 0.5.
    offset : float, optional
        Defaults to 0.5.

    Reference
    ---------
    Gregory Gelly and Jean-Luc Gauvain. "Minimum Word Error Training of
    RNN-based Voice Activity Detection", InterSpeech 2015.
    """

    def __init__(self, onset=0.7, offset=0.7):

        super(Binarize, self).__init__()
        self.onset = onset
        self.offset = offset

        self.min_duration = [0., 0.]
        self.pad_onset = 0.
        self.pad_offset = 0.

    @classmethod
    def tune(cls, items, get_prediction, get_metric=None,
             n_calls=20, n_random_starts=10, n_jobs=-1,
             **kwargs):
        """Find best set of hyper-parameters using skopt

        Parameters
        ----------
        items : iterable
            Protocol items used as development set. Typically, one would use
            items = protocol.development()
        get_prediction : callable
            Callable that takes an item as input, and returns a prediction as
            a SlidingWindowFeature instances.
        get_metric : callable, optional
            Callable that takes no input, and returns a fresh evaluation
            metric. Defaults to pyannote.metrics.detection.DetectionErrorRate.
        n_calls : int, optional
            Number of trials for hyper-parameter optimization. Defaults to 20.
        n_random_starts : int, optional
            Number of trials with random initialization before being smart.
            Defaults to 10.
        n_jobs : int, optional
            Number of parallel job to use. Set to 1 to not use multithreading.
            Defaults to whichever is minimum between number of CPUs and number
            of items.
        **kwargs :
            Optional keyword arguments passed to Binarize.apply().

        Returns
        -------
        best_parameters : dict
            Best set of parameters.
        best_metric : float
            Best metric
        """

        import skopt
        import skopt.space

        if get_metric is None:
            from pyannote.metrics.detection import DetectionErrorRate
            get_metric = DetectionErrorRate

        # make sure items can be iterated over and over again
        items = list(items)

        # defaults to whichever is minimum between
        # number of CPUs and number of items
        if n_jobs < 0:
            n_jobs = min(cpu_count(), len(items))

        if n_jobs > 1:
            pool = Pool(n_jobs)

        # compute predictions once and for all
        # NOTE could multithreading speed things up? this is unclear as
        # get_prediction is probably already multithreaded, or (better) even
        # precomputed
        predictions = [get_prediction(item) for item in items]

        def objective_function(params):
            onset, offset, = params
            binarizer = cls(onset=onset, offset=offset)

            metric = get_metric()
            process_one_file = functools.partial(helper_binarize_tune,
                                                 binarizer=binarizer,
                                                 metric=metric,
                                                 **kwargs)

            if n_jobs > 1:
                results = list(pool.map(process_one_file,
                                        zip(items, predictions)))
            else:
                results = [process_one_file(item_prediction)
                           for item_prediction in zip(items, predictions)]

            return abs(metric)

        space = [skopt.space.Real(0., 1., prior='uniform'),
                 skopt.space.Real(0., 1., prior='uniform')]

        res = skopt.gp_minimize(
            objective_function, space,
            x0=[[0.25, 0.25],[0.5, 0.5], [0.75, 0.75]],
            n_calls=n_calls, n_random_starts=n_random_starts,
            random_state=1337, verbose=False)

        if n_jobs > 1:
            pool.terminate()

        return {'onset': res.x[0], 'offset': res.x[1]}, res.fun

    def apply(self, predictions, dimension=0):
        """
        Parameters
        ----------
        predictions : SlidingWindowFeature
            Must be mono-dimensional
        dimension : int, optional
            Which dimension to process
        """

        if len(predictions.data.shape) == 1:
            data = predictions.data
        elif predictions.data.shape[1] == 1:
            data = predictions.data[:, 0]
        else:
            data = predictions.data[:, dimension]

        n_samples = predictions.getNumber()
        window = predictions.sliding_window
        timestamps = [window[i].middle for i in range(n_samples)]

        # initial state
        start = timestamps[0]
        label = data[0] > self.onset

        # timeline meant to store 'active' segments
        active = Timeline()

        for t, y in zip(timestamps[1:], data[1:]):

            # currently active
            if label:
                # switching from active to inactive
                if y < self.offset:
                    segment = Segment(start - self.pad_onset,
                                      t + self.pad_offset)
                    active.add(segment)
                    start = t
                    label = False

            # currently inactive
            else:
                # switching from inactive to active
                if y > self.onset:
                    start = t
                    label = True

        # if active at the end, add final segment
        if label:
            segment = Segment(start - self.pad_onset, t + self.pad_offset)
            active.add(segment)

        # because of padding, some 'active' segments might be overlapping
        # therefore, we merge those overlapping segments
        active = active.support()

        # remove short 'active' segments
        active = Timeline(
            [s for s in active if s.duration > self.min_duration[1]])

        # fill short 'inactive' segments
        inactive = active.gaps()
        for s in inactive:
            if s.duration < self.min_duration[0]:
                active.add(s)
        active = active.support()

        return active
<|endoftext|>