<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

from collections import defaultdict
import itertools
import types
import logging
from zope.index.text.textindex import TextIndex
from zope.index.text.cosineindex import CosineIndex
from zope.index.text.lexicon import Lexicon
from zope.index.text.lexicon import Splitter
import time

import dedupe.tfidf as tfidf

logger = logging.getLogger(__name__)

class Blocker:
    '''Takes in a record and returns all blocks that record belongs to'''
    def __init__(self,
                 predicates,
                 stop_words = defaultdict(set)) :

        self.predicates = predicates
        self.stop_words = defaultdict(set)

        self.tfidf_fields = defaultdict(set)

        for full_predicate in self.predicates :
            for predicate in full_predicate :
                if predicate.type == "TfidfPredicate" :
                    self.tfidf_fields[predicate.field].add(predicate)

    def __call__(self, records):

        for record in records :
            record_id = record[0]

            seen_keys = set([])

            for predicate in self.predicates :
                for block_key in predicate(record) :
                    if (block_key, predicate) not in seen_keys :
                        yield (block_key, predicate), record_id
                    else :
                        seen_keys.add(block_key, predicate)

class DedupeBlocker(Blocker) :

    def tfIdfBlock(self, data, field):
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])

        index = TextIndex(Lexicon(splitter, stop_word_remover))

        index.index = CosineIndex(index.lexicon)

        index_to_id = {}
        base_tokens = {}

        for i, (record_id, doc) in enumerate(data, 1) :
            index_to_id[i] = record_id
            base_tokens[i] = splitter.process([doc])
            index.index_doc(i, doc)

        for predicate in self.tfidf_fields[field] :
            canopy = tfidf.makeCanopy(index,
                                      base_tokens,
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())

class RecordLinkBlocker(Blocker) :
    def tfIdfBlock(self, data_1, data_2, field):
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])

        index = TextIndex(Lexicon(splitter, stop_word_remover))

        index.index = CosineIndex(index.lexicon)

        index_to_id = {}
        base_tokens = {}

        i = 1

        for record_id, doc in data_1 :
            index_to_id[i] = record_id
            base_tokens[i] = splitter.process([doc])
            i += 1

        for record_id, doc in data_2  :
            index_to_id[i] = record_id
            index.index_doc(i, doc)
            i += 1

        for predicate in self.tfidf_fields[field] :
            canopy = tfidf.makeCanopy(index,
                                      base_tokens,
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())

def blockTraining(training_pairs,
                  predicate_set,
                  eta=.1,
                  epsilon=.1,
                  matching = "Dedupe"):
    '''
    Takes in a set of training pairs and predicates and tries to find
    a good set of blocking rules.
    '''

    # Setup

    training_dupes = set(training_pairs['match'])
    training_distinct = set(training_pairs['distinct'])

    if matching == "RecordLink" :
        coverage = RecordLinkCoverage(predicate_set,
                                      training_dupes | training_distinct)

    else :
        coverage = DedupeCoverage(predicate_set,
                                  training_dupes | training_distinct)

    # Compound Predicates
    compound_predicates = itertools.combinations(coverage.overlap, 2)

    for compound_predicate in compound_predicates :
        predicate_1, predicate_2 = compound_predicate
        coverage.overlap[CompoundPredicate(compound_predicate)] =\
            coverage.overlap[predicate_1] & coverage.overlap[predicate_2]

    predicate_set = coverage.overlap.keys()

    coverage_threshold = eta * len(training_distinct)
    logger.info("coverage threshold: %s", coverage_threshold)

    # Only consider predicates that cover at least one duplicate pair
    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               training_dupes)
    predicate_set = dupe_coverage.keys()

    # Within blocks, we will compare every combination of
    # records. Therefore, we want to avoid predicates that make large
    # blocks.
    for pred in predicate_set[:] :
        for block in coverage.blocks[pred] :
            if len(block) > 100 :
                predicate_set.remove(pred)
                continue

    # As an efficency, we can throw away the predicates that cover too
    # many distinct pairs
    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    for pred, pairs in distinct_coverage.items() :
        if len(pairs) > coverage_threshold :
            predicate_set.remove(pred)

    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    final_predicate_set = findOptimumBlocking(training_dupes,
                                              predicate_set,
                                              distinct_coverage,
                                              epsilon,
                                              coverage)

    logger.info('Final predicate set:')
    for predicate in final_predicate_set :
        logger.info(predicate)

    if final_predicate_set:
        return final_predicate_set, coverage.stop_words
    else:
        raise ValueError('No predicate found! We could not learn a single good predicate. Maybe give Dedupe more training data')

def findOptimumBlocking(uncovered_dupes,
                        predicate_set,
                        distinct_coverage,
                        epsilon,
                        coverage):

    # Greedily find the predicates that, at each step, covers the
    # most duplicates and covers the least distinct pairs, due to
    # Chvatal, 1979
    #
    # We would like to optimize the ratio of the probability of of a
    # predicate covering a duplicate pair versus the probability of
    # covering a distinct pair. If we have a uniform prior belief
    # about those probabilities, we can estimate these probabilities as
    #
    # (predicate_covered_dupe_pairs + 1) / (all_dupe_pairs + 2)
    #
    # (predicate_covered_distinct_pairs + 1) / (all_distinct_pairs + 2)
    #
    # When we are trying to find the best predicate among a set of
    # predicates, the denominators factor out and our coverage
    # estimator becomes
    #
    # (predicate_covered_dupe_pairs + 1)/ (predicate_covered_distinct_pairs + 1)

    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               uncovered_dupes)

    uncovered_dupes = set(uncovered_dupes)

    final_predicate_set = []
    while len(uncovered_dupes) > epsilon:

        best_cover = 0
        best_predicate = None
        for predicate in dupe_coverage :
            dupes = len(dupe_coverage[predicate])
            distinct = len(distinct_coverage[predicate])
            cover = (dupes + 1.0)/(distinct + 1.0)
            if cover > best_cover:
                best_cover = cover
                best_predicate = predicate
                best_distinct = distinct
                best_dupes = dupes

        if not best_predicate:
            logger.warning('Ran out of predicates')
            break

        final_predicate_set.append(best_predicate)
        predicate_set.remove(best_predicate)

        uncovered_dupes = uncovered_dupes - dupe_coverage[best_predicate]
        dupe_coverage = coverage.predicateCoverage(predicate_set,
                                                   uncovered_dupes)

        logger.debug(best_predicate)
        logger.debug('cover: %(cover)f, found_dupes: %(found_dupes)d, '
                      'found_distinct: %(found_distinct)d, '
                      'uncovered dupes: %(uncovered)d',
                      {'cover' : best_cover,
                       'found_dupes' : best_dupes,
                       'found_distinct' : best_distinct,
                       'uncovered' : len(uncovered_dupes)
                       })

    return final_predicate_set

class Coverage(object) :

    def coveredBy(self, id_records, record_ids, blocker, pairs) :
        self.overlap = defaultdict(set)
        self.blocks = defaultdict(lambda : defaultdict(set))

        covered_by = defaultdict(set)

        for block_key, record_id in blocker(id_records.items()) :
            covered_by[record_id].add(block_key)

        for record_1, record_2 in pairs :
            id_1 = record_ids[record_1]
            id_2 = record_ids[record_2]

            blocks = covered_by[id_1] & covered_by[id_2]
            for block_key, predicate in blocks :
                self.overlap[predicate].add((record_1, record_2))
                self.blocks[predicate][block_key].add((record_1, record_2))

    def predicateCoverage(self,
                          predicate_set,
                          pairs) :

        coverage = defaultdict(set)
        pairs = set(pairs)

        for predicate in predicate_set :
            covered_pairs = pairs.intersection(self.overlap[predicate])
            if covered_pairs :
                coverage[predicate] = covered_pairs

        return coverage

class DedupeCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :
        self.stop_words = {}

        records = set(itertools.chain(*pairs))

        id_records = dict(itertools.izip(itertools.count(), records))
        record_ids = dict(itertools.izip(records, itertools.count()))

        blocker = DedupeBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            data = ((record_id, record[field])
                    for record_id, record
                    in id_records.items())
            blocker.tfIdfBlock(data, field)

        self.coveredBy(id_records, record_ids, blocker, pairs)

class RecordLinkCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :
        self.stop_words = {}

        data_1 = set([])
        data_2 = set([])

        for record_1, record_2 in pairs :
            data_1.add(record_1)
            data_2.add(record_2)

        i = 0
        id_records_1 = {}
        id_records_2 = {}

        for record in data_1 :
            id_records_1[i] = record
            i += 1

        for record in data_2 :
            id_records_2[i] = record
            i += 1

        id_records = id_records_1.copy()
        id_records.update(id_records_2)

        record_ids = dict((record, record_id)
                          for record_id, record
                          in id_records.items())

        blocker = RecordLinkBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            fields_1 = ((record_id, record[field])
                        for record_id, record
                        in id_records_1.items())
            fields_2 = ((record_id, record[field])
                        for record_id, record
                        in id_records_2.items())
            blocker.tfIdfBlock(fields_1, fields_2, field)

        self.coveredBy(id_records, record_ids, blocker, pairs)

def stopWords(data) :
    index = TextIndex(Lexicon(Splitter()))

    for i, (_, doc) in enumerate(data, 1) :
        index.index_doc(i, doc)

    doc_freq = [(len(index.index._wordinfo[wid]), word)
                for word, wid in index.lexicon.items()]

    doc_freq.sort(reverse=True)

    N = float(index.index.documentCount())
    threshold = int(max(1000, N * 0.05))

    stop_words = set([])

    for frequency, word in doc_freq :
        if frequency > threshold :
            stop_words.add(word)
        else :
            break

    return stop_words

class Predicate(object) :
    def __iter__(self) :
        yield self

    def __repr__(self) :
        return "%s: %s" % (self.type, self.__name__)

class SimplePredicate(Predicate) :
    type = "SimplePredicate"

    def __init__(self, func, field) :
        self.func = func
        self.__name__ = "(%s, %s)" % (func.__name__, field)
        self.field = field

    def __call__(self, instance) :
        record = instance[1]
        for block_key in  self.func(record[self.field]) :
            yield block_key

class TfidfPredicate(Predicate):
    type = "TfidfPredicate"

    def __init__(self, threshold, field):
        self.__name__ = '(%s, %s)' % (threshold, field)
        self.field = field
        self.canopy = {}
        self.threshold = threshold

    def __call__(self, record) :
        record_id = record[0]
        center = self.canopy.get(record_id)
        if center is not None :
            return (unicode(center),)
        else :
            return ()

    def __getstate__(self):
        result = self.__dict__.copy()
        result['canopy'] = {}
        return result

class CompoundPredicate(Predicate) :
    type = "CompoundPredicate"

    def __init__(self, predicates) :
        self.predicates = predicates
        self.__name__ = '(%s)' % ', '.join([str(pred)
                                            for pred in
                                            predicates])

    def __iter__(self) :
        for pred in self.predicates :
            yield pred

    def __call__(self, record) :
        block_keys = []
        for predicate in self.predicates :
            block_keys.append(list(predicate(record)))

        for block_key in itertools.product(*block_keys) :
            yield ':'.join(block_key)

class CustomStopWordRemover(object):
    def __init__(self, stop_words) :
        self.stop_words = stop_words.copy()

    def process(self, lst):
        return [w for w in lst if not w in self.stop_words]

<commit_msg>logging messages for canopy creation<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

from collections import defaultdict
import itertools
import types
import logging
from zope.index.text.textindex import TextIndex
from zope.index.text.cosineindex import CosineIndex
from zope.index.text.lexicon import Lexicon
from zope.index.text.lexicon import Splitter
import time

import dedupe.tfidf as tfidf

logger = logging.getLogger(__name__)

class Blocker:
    '''Takes in a record and returns all blocks that record belongs to'''
    def __init__(self,
                 predicates,
                 stop_words = defaultdict(set)) :

        self.predicates = predicates
        self.stop_words = defaultdict(set)

        self.tfidf_fields = defaultdict(set)

        for full_predicate in self.predicates :
            for predicate in full_predicate :
                if predicate.type == "TfidfPredicate" :
                    self.tfidf_fields[predicate.field].add(predicate)

    def __call__(self, records):

        for record in records :
            record_id = record[0]

            seen_keys = set([])

            for predicate in self.predicates :
                for block_key in predicate(record) :
                    if (block_key, predicate) not in seen_keys :
                        yield (block_key, predicate), record_id
                    else :
                        seen_keys.add(block_key, predicate)

class DedupeBlocker(Blocker) :

    def tfIdfBlock(self, data, field):
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])

        index = TextIndex(Lexicon(splitter, stop_word_remover))

        index.index = CosineIndex(index.lexicon)

        index_to_id = {}
        base_tokens = {}

        for i, (record_id, doc) in enumerate(data, 1) :
            index_to_id[i] = record_id
            base_tokens[i] = splitter.process([doc])
            index.index_doc(i, doc)

        for predicate in self.tfidf_fields[field] :
            logger.info("Canopy: %s", str(predicate))
            canopy = tfidf.makeCanopy(index,
                                      base_tokens,
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())

class RecordLinkBlocker(Blocker) :
    def tfIdfBlock(self, data_1, data_2, field):
        '''Creates TF/IDF canopy of a given set of data'''

        splitter = Splitter()

        stop_word_remover = CustomStopWordRemover(self.stop_words[field])

        index = TextIndex(Lexicon(splitter, stop_word_remover))

        index.index = CosineIndex(index.lexicon)

        index_to_id = {}
        base_tokens = {}

        i = 1

        for record_id, doc in data_1 :
            index_to_id[i] = record_id
            base_tokens[i] = splitter.process([doc])
            i += 1

        for record_id, doc in data_2  :
            index_to_id[i] = record_id
            index.index_doc(i, doc)
            i += 1

        for predicate in self.tfidf_fields[field] :
            logger.info("Canopy: %s", str(predicate))
            canopy = tfidf.makeCanopy(index,
                                      base_tokens,
                                      predicate.threshold)
            predicate.canopy = dict((index_to_id[k], index_to_id[v])
                                    for k, v
                                    in canopy.iteritems())

def blockTraining(training_pairs,
                  predicate_set,
                  eta=.1,
                  epsilon=.1,
                  matching = "Dedupe"):
    '''
    Takes in a set of training pairs and predicates and tries to find
    a good set of blocking rules.
    '''

    # Setup

    training_dupes = set(training_pairs['match'])
    training_distinct = set(training_pairs['distinct'])

    if matching == "RecordLink" :
        coverage = RecordLinkCoverage(predicate_set,
                                      training_dupes | training_distinct)

    else :
        coverage = DedupeCoverage(predicate_set,
                                  training_dupes | training_distinct)

    # Compound Predicates
    compound_predicates = itertools.combinations(coverage.overlap, 2)

    for compound_predicate in compound_predicates :
        predicate_1, predicate_2 = compound_predicate
        coverage.overlap[CompoundPredicate(compound_predicate)] =\
            coverage.overlap[predicate_1] & coverage.overlap[predicate_2]

    predicate_set = coverage.overlap.keys()

    coverage_threshold = eta * len(training_distinct)
    logger.info("coverage threshold: %s", coverage_threshold)

    # Only consider predicates that cover at least one duplicate pair
    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               training_dupes)
    predicate_set = dupe_coverage.keys()

    # Within blocks, we will compare every combination of
    # records. Therefore, we want to avoid predicates that make large
    # blocks.
    for pred in predicate_set[:] :
        for block in coverage.blocks[pred] :
            if len(block) > 100 :
                predicate_set.remove(pred)
                continue

    # As an efficency, we can throw away the predicates that cover too
    # many distinct pairs
    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    for pred, pairs in distinct_coverage.items() :
        if len(pairs) > coverage_threshold :
            predicate_set.remove(pred)

    distinct_coverage = coverage.predicateCoverage(predicate_set,
                                                   training_distinct)

    final_predicate_set = findOptimumBlocking(training_dupes,
                                              predicate_set,
                                              distinct_coverage,
                                              epsilon,
                                              coverage)

    logger.info('Final predicate set:')
    for predicate in final_predicate_set :
        logger.info(predicate)

    if final_predicate_set:
        return final_predicate_set, coverage.stop_words
    else:
        raise ValueError('No predicate found! We could not learn a single good predicate. Maybe give Dedupe more training data')

def findOptimumBlocking(uncovered_dupes,
                        predicate_set,
                        distinct_coverage,
                        epsilon,
                        coverage):

    # Greedily find the predicates that, at each step, covers the
    # most duplicates and covers the least distinct pairs, due to
    # Chvatal, 1979
    #
    # We would like to optimize the ratio of the probability of of a
    # predicate covering a duplicate pair versus the probability of
    # covering a distinct pair. If we have a uniform prior belief
    # about those probabilities, we can estimate these probabilities as
    #
    # (predicate_covered_dupe_pairs + 1) / (all_dupe_pairs + 2)
    #
    # (predicate_covered_distinct_pairs + 1) / (all_distinct_pairs + 2)
    #
    # When we are trying to find the best predicate among a set of
    # predicates, the denominators factor out and our coverage
    # estimator becomes
    #
    # (predicate_covered_dupe_pairs + 1)/ (predicate_covered_distinct_pairs + 1)

    dupe_coverage = coverage.predicateCoverage(predicate_set,
                                               uncovered_dupes)

    uncovered_dupes = set(uncovered_dupes)

    final_predicate_set = []
    while len(uncovered_dupes) > epsilon:

        best_cover = 0
        best_predicate = None
        for predicate in dupe_coverage :
            dupes = len(dupe_coverage[predicate])
            distinct = len(distinct_coverage[predicate])
            cover = (dupes + 1.0)/(distinct + 1.0)
            if cover > best_cover:
                best_cover = cover
                best_predicate = predicate
                best_distinct = distinct
                best_dupes = dupes

        if not best_predicate:
            logger.warning('Ran out of predicates')
            break

        final_predicate_set.append(best_predicate)
        predicate_set.remove(best_predicate)

        uncovered_dupes = uncovered_dupes - dupe_coverage[best_predicate]
        dupe_coverage = coverage.predicateCoverage(predicate_set,
                                                   uncovered_dupes)

        logger.debug(best_predicate)
        logger.debug('cover: %(cover)f, found_dupes: %(found_dupes)d, '
                      'found_distinct: %(found_distinct)d, '
                      'uncovered dupes: %(uncovered)d',
                      {'cover' : best_cover,
                       'found_dupes' : best_dupes,
                       'found_distinct' : best_distinct,
                       'uncovered' : len(uncovered_dupes)
                       })

    return final_predicate_set

class Coverage(object) :

    def coveredBy(self, id_records, record_ids, blocker, pairs) :
        self.overlap = defaultdict(set)
        self.blocks = defaultdict(lambda : defaultdict(set))

        covered_by = defaultdict(set)

        for block_key, record_id in blocker(id_records.items()) :
            covered_by[record_id].add(block_key)

        for record_1, record_2 in pairs :
            id_1 = record_ids[record_1]
            id_2 = record_ids[record_2]

            blocks = covered_by[id_1] & covered_by[id_2]
            for block_key, predicate in blocks :
                self.overlap[predicate].add((record_1, record_2))
                self.blocks[predicate][block_key].add((record_1, record_2))

    def predicateCoverage(self,
                          predicate_set,
                          pairs) :

        coverage = defaultdict(set)
        pairs = set(pairs)

        for predicate in predicate_set :
            covered_pairs = pairs.intersection(self.overlap[predicate])
            if covered_pairs :
                coverage[predicate] = covered_pairs

        return coverage

class DedupeCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :
        self.stop_words = {}

        records = set(itertools.chain(*pairs))

        id_records = dict(itertools.izip(itertools.count(), records))
        record_ids = dict(itertools.izip(records, itertools.count()))

        blocker = DedupeBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            data = ((record_id, record[field])
                    for record_id, record
                    in id_records.items())
            blocker.tfIdfBlock(data, field)

        self.coveredBy(id_records, record_ids, blocker, pairs)

class RecordLinkCoverage(Coverage) :
    def __init__(self, predicate_set, pairs) :
        self.stop_words = {}

        data_1 = set([])
        data_2 = set([])

        for record_1, record_2 in pairs :
            data_1.add(record_1)
            data_2.add(record_2)

        i = 0
        id_records_1 = {}
        id_records_2 = {}

        for record in data_1 :
            id_records_1[i] = record
            i += 1

        for record in data_2 :
            id_records_2[i] = record
            i += 1

        id_records = id_records_1.copy()
        id_records.update(id_records_2)

        record_ids = dict((record, record_id)
                          for record_id, record
                          in id_records.items())

        blocker = RecordLinkBlocker(predicate_set)

        for field in blocker.tfidf_fields :
            fields_1 = ((record_id, record[field])
                        for record_id, record
                        in id_records_1.items())
            fields_2 = ((record_id, record[field])
                        for record_id, record
                        in id_records_2.items())
            blocker.tfIdfBlock(fields_1, fields_2, field)

        self.coveredBy(id_records, record_ids, blocker, pairs)

def stopWords(data) :
    index = TextIndex(Lexicon(Splitter()))

    for i, (_, doc) in enumerate(data, 1) :
        index.index_doc(i, doc)

    doc_freq = [(len(index.index._wordinfo[wid]), word)
                for word, wid in index.lexicon.items()]

    doc_freq.sort(reverse=True)

    N = float(index.index.documentCount())
    threshold = int(max(1000, N * 0.05))

    stop_words = set([])

    for frequency, word in doc_freq :
        if frequency > threshold :
            stop_words.add(word)
        else :
            break

    return stop_words

class Predicate(object) :
    def __iter__(self) :
        yield self

    def __repr__(self) :
        return "%s: %s" % (self.type, self.__name__)

class SimplePredicate(Predicate) :
    type = "SimplePredicate"

    def __init__(self, func, field) :
        self.func = func
        self.__name__ = "(%s, %s)" % (func.__name__, field)
        self.field = field

    def __call__(self, instance) :
        record = instance[1]
        for block_key in  self.func(record[self.field]) :
            yield block_key

class TfidfPredicate(Predicate):
    type = "TfidfPredicate"

    def __init__(self, threshold, field):
        self.__name__ = '(%s, %s)' % (threshold, field)
        self.field = field
        self.canopy = {}
        self.threshold = threshold

    def __call__(self, record) :
        record_id = record[0]
        center = self.canopy.get(record_id)
        if center is not None :
            return (unicode(center),)
        else :
            return ()

    def __getstate__(self):
        result = self.__dict__.copy()
        result['canopy'] = {}
        return result

class CompoundPredicate(Predicate) :
    type = "CompoundPredicate"

    def __init__(self, predicates) :
        self.predicates = predicates
        self.__name__ = '(%s)' % ', '.join([str(pred)
                                            for pred in
                                            predicates])

    def __iter__(self) :
        for pred in self.predicates :
            yield pred

    def __call__(self, record) :
        block_keys = []
        for predicate in self.predicates :
            block_keys.append(list(predicate(record)))

        for block_key in itertools.product(*block_keys) :
            yield ':'.join(block_key)

class CustomStopWordRemover(object):
    def __init__(self, stop_words) :
        self.stop_words = stop_words.copy()

    def process(self, lst):
        return [w for w in lst if not w in self.stop_words]

<|endoftext|>


<commit_before>import numpy
from dedupe.distance.affinegap import normalizedAffineGapDistance as comparator

def getCentroid( attribute_variants, comparator ):
    """
    Takes in a list of attribute values for a field,
    evaluates the centroid using the comparator,
    & returns the centroid (i.e. the 'best' value for the field)
    """

    n = len(attribute_variants)

    distance_matrix = numpy.zeros([n,n])

    # populate distance matrix by looping through elements of matrix triangle
    for i in range (0,n):
        for j in range (0, i):
            distance = comparator(attribute_variants[i], attribute_variants[j])
            distance_matrix[i,j] = distance_matrix[j,i] = distance

    average_distance = distance_matrix.mean(0)

    # there can be ties for minimum, average distance string
    min_dist_indices = numpy.where(average_distance==average_distance.min())[0]

    if len(min_dist_indices) > 1:
        centroid = breakCentroidTie(attribute_variants, min_dist_indices)
    else :
        centroid_index = min_dist_indices[0]
        centroid = attribute_variants[centroid_index]

    return centroid

def breakCentroidTie( attribute_variants, min_dist_indices ):
    """
    finds centroid when there are multiple values w/ min avg distance
    (e.g. any dupe cluster of 2) right now this selects the first among a set of
    ties, but can be modified to break ties in strings by selecting the longest string
    """
    return attribute_variants[min_dist_indices[0]]

def getCanonicalRep( record_cluster ):
    canonical_rep = {}

    for key in record_cluster[0].keys():
        key_values = []
        for record in record_cluster :
            # assume non-empty values always better than empty value for canonical record
            if record[key] != '':
                key_values.append(record[key])
        if key_values:
            canonical_rep[key] = getCentroid(key_values, comparator)
        else:
            canonical_rep[key] = ''

    return canonical_rep
<commit_msg>clean up conditional<commit_after>import numpy
from dedupe.distance.affinegap import normalizedAffineGapDistance as comparator

def getCentroid( attribute_variants, comparator ):
    """
    Takes in a list of attribute values for a field,
    evaluates the centroid using the comparator,
    & returns the centroid (i.e. the 'best' value for the field)
    """

    n = len(attribute_variants)

    distance_matrix = numpy.zeros([n,n])

    # populate distance matrix by looping through elements of matrix triangle
    for i in range (0,n):
        for j in range (0, i):
            distance = comparator(attribute_variants[i], attribute_variants[j])
            distance_matrix[i,j] = distance_matrix[j,i] = distance

    average_distance = distance_matrix.mean(0)

    # there can be ties for minimum, average distance string
    min_dist_indices = numpy.where(average_distance==average_distance.min())[0]

    if len(min_dist_indices) > 1:
        centroid = breakCentroidTie(attribute_variants, min_dist_indices)
    else :
        centroid_index = min_dist_indices[0]
        centroid = attribute_variants[centroid_index]

    return centroid

def breakCentroidTie( attribute_variants, min_dist_indices ):
    """
    finds centroid when there are multiple values w/ min avg distance
    (e.g. any dupe cluster of 2) right now this selects the first among a set of
    ties, but can be modified to break ties in strings by selecting the longest string
    """
    return attribute_variants[min_dist_indices[0]]

def getCanonicalRep( record_cluster ):
    canonical_rep = {}

    for key in record_cluster[0].keys():
        key_values = []
        for record in record_cluster :
            # assume non-empty values always better than empty value for canonical record
            if record[key]:
                key_values.append(record[key])
        if key_values:
            canonical_rep[key] = getCentroid(key_values, comparator)
        else:
            canonical_rep[key] = ''

    return canonical_rep
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

# provides functions for selecting a sample of training data
from __future__ import division
from future.utils import viewitems, viewvalues

import itertools
from . import blocking, predicates, core
import numpy
import logging

logger = logging.getLogger(__name__)

class BlockLearner(object) :
    def learn(self, matches, recall) :
        '''
        Takes in a set of training pairs and predicates and tries to find
        a good set of blocking rules.
        '''
        compound_length = 2

        self.blocker.indexAll({i : record
                               for i, record
                               in enumerate(self.unroll(matches))})

        dupe_cover = cover(self.blocker, matches, compound_length)

        comparison_count = self.comparisons(self.total_cover, compound_length)

        coverable_dupes = set.union(*viewvalues(dupe_cover))
        uncoverable_dupes = [pair for i, pair in enumerate(matches)
                             if i not in coverable_dupes]

        epsilon = int((1.0 - recall) * len(matches))

        if len(uncoverable_dupes) > epsilon :
            logger.warning(OUT_OF_PREDICATES_WARNING)
            logger.debug(uncoverable_dupes)
            epsilon = 0
        else :
            epsilon -= len(uncoverable_dupes)

        searcher = BranchBound(dupe_cover, comparison_count, epsilon, 5000)
        final_predicates = searcher.search(dupe_cover)

        logger.info('Final predicate set:')
        for predicate in final_predicates:
            logger.info(predicate)

        return final_predicates

    def comparisons(self, cover, compound_length) :
        CP = predicates.CompoundPredicate

        block_index = {}
        for predicate, blocks in viewitems(cover):
            block_index[predicate] = {}
            for block_id, blocks in viewitems(blocks) :
                for record_id in self._blocks(blocks) :
                    block_index[predicate].setdefault(record_id,
                                                      set()).add(block_id)

        compounder = self.Compounder(cover, block_index)
        comparison_count = {}
        simple_predicates = sorted(cover, key=str)

        for i in range(2, compound_length+1) :
            for combo in itertools.combinations(simple_predicates, i) :
                comparison_count[CP(combo)] = self.estimate(compounder(combo))
        for pred in simple_predicates :
            comparison_count[pred] = self.estimate(viewvalues(cover[pred]))

        return comparison_count

class Compounder(object) :
    def __init__(self, cover, block_index) :
        self.cover = cover
        self.block_index = block_index
        self._cached_predicate = None
        self._cached_blocks = None

    def __call__(self, compound_predicate) :
        a, b = compound_predicate[:-1], compound_predicate[-1]

        if len(a) > 1 :
            if a == self._cached_predicate :
                a_blocks = self._cached_blocks
            else :
                a_blocks = self._cached_blocks = list(self(a))
                self._cached_predicate = a
        else :
            a_blocks = viewvalues(self.cover[a[0]])

        return self.overlap(a_blocks, b)

class DedupeCompounder(Compounder) :
    def overlap(self, a_blocks, b) :
        b_index = self.block_index[b]
        b_index_get = b_index.get
        cover_b = self.cover[b]
        null_set = set()
        for x_ids in a_blocks :
            seen_y = set()
            for record_id in x_ids :
                b_blocks = b_index_get(record_id, null_set)
                for y in b_blocks :
                    if y not in seen_y :
                        yield x_ids & cover_b[y]
                seen_y |= b_blocks

class DedupeBlockLearner(BlockLearner) :
    Compounder = DedupeCompounder

    def __init__(self, predicates, sampled_records) :
        blocker = blocking.Blocker(predicates)
        blocker.indexAll(sampled_records)

        self.total_cover = self.coveredRecords(blocker, sampled_records)
        self.multiplier = sampled_records.original_length/len(sampled_records)

        self.blocker = blocking.Blocker(predicates)

    @staticmethod
    def unroll(matches) : # pragma: no cover
        return unique((record for pair in matches for record in pair))

    @staticmethod
    def _blocks(blocks) : # pragma: no cover
        return blocks

    @staticmethod
    def coveredRecords(blocker, records) :
        cover = {}

        for predicate in blocker.predicates :
            cover[predicate] = {}
            for id, record in viewitems(records) :
                blocks = predicate(record)
                for block in blocks :
                    cover[predicate].setdefault(block, set()).add(id)

        return cover

    def estimate(self, blocks):
        lengths = numpy.fromiter((len(ids) for ids in blocks), int)
        return numpy.sum((lengths * self.multiplier) *
                         (lengths * self.multiplier - 1) /
                         2)

class RecordLinkCompounder(Compounder) :
    def overlap(self, a_blocks, b) :
        b_index = self.block_index[b]
        b_index_get = b_index.get
        cover_b = self.cover[b]
        null_set = set()
        for first, second in a_blocks:
            seen_y = set()
            for record_id in first:
                b_blocks = b_index_get(record_id, null_set)
                for y in b_blocks:
                    if y not in seen_y :
                        yield first & cover_b[y][0], second & cover_b[y][1]
                    seen_y |= b_blocks

class RecordLinkBlockLearner(BlockLearner) :
    Compounder = RecordLinkCompounder

    def __init__(self, predicates, sampled_records_1, sampled_records_2) :
        blocker = blocking.Blocker(predicates)
        blocker.indexAll(sampled_records_2)

        self.total_cover = self.coveredRecords(blocker,
                                               sampled_records_1,
                                               sampled_records_2)

        self.multiplier_1 = sampled_records_1.original_length/len(sampled_records_1)
        self.multiplier_2 = sampled_records_2.original_length/len(sampled_records_2)

        self.blocker = blocking.Blocker(predicates)

    @staticmethod
    def unroll(matches) : # pragma: no cover
        return unique((record_2 for _, record_2 in matches))

    @staticmethod
    def _blocks(blocks) : # pragma: no cover
        return blocks[0]

    @staticmethod
    def coveredRecords(blocker, records_1, records_2) :
        cover = {}

        for predicate in blocker.predicates :
            cover[predicate] = {}
            for id, record in viewitems(records_2) :
                blocks = predicate(record)
                for block in blocks :
                    cover[predicate].setdefault(block, (set(), set()))[1].add(id)

            current_blocks = set(cover[predicate])
            for id, record in viewitems(records_1) :
                blocks = set(predicate(record))
                for block in blocks & current_blocks :
                    cover[predicate][block][0].add(id)

        return cover

    def estimate(self, blocks):
        A, B = core.iunzip(blocks, 2)
        lengths_A = numpy.fromiter((len(ids) for ids in A), float)
        lengths_B = numpy.fromiter((len(ids) for ids in B), float)
        return numpy.sum((lengths_A * self.multiplier_1) *
                         (lengths_B * self.multiplier_2))

class BranchBound(object) :
    def __init__(self, original_cover, comparison_count, epsilon, max_calls) :
        self.dupes_to_cover = set.union(*original_cover.values())
        self.original_cover = original_cover.copy()
        self.calls = max_calls
        self.comparisons = comparison_count
        self.epsilon = epsilon

        self.cheapest = tuple(original_cover.keys())
        self.cheapest_score = float('inf')

    def search(self, dupe_cover, partial=()) :
        if self.calls <= 0 :
            return self.cheapest

        self.calls -= 1

        uncovered_dupes = self.dupes_to_cover.difference(*(self.original_cover[pred]
                                                           for pred in partial))

        if len(uncovered_dupes) <= self.epsilon :
            partial_score = self.score(partial)
            if partial_score < self.cheapest_score :
                print(partial_score, self.cheapest_score)
                self.cheapest = partial
                self.cheapest_score = partial_score

        elif dupe_cover and (self.lower_bound(partial, dupe_cover) <= self.cheapest_score):
            cost = lambda p : self.comparisons[p]/len(dupe_cover[p])
            best_predicate = min(dupe_cover, key=cost)

            remaining = remaining_cover(dupe_cover, dupe_cover[best_predicate])

            self.search(remaining, partial + (best_predicate,))

            reduced = self.dominates(dupe_cover, best_predicate)

            uncoverable_dupes = uncovered_dupes.difference(*reduced.values())
            if len(uncoverable_dupes) <= self.epsilon:
                self.search(reduced, partial)

        return self.cheapest

    def score(self, partial):
        return sum(self.comparisons[p] for p in partial)

    def lower_bound(self, partial, dupe_cover) :
        return self.score(partial) + min(self.comparisons[p] for p in dupe_cover)

    def dominates(self, coverage, dominator):
        remaining = coverage.copy()

        dominant_cost = self.comparisons[dominator]
        dominant_cover = coverage[dominator]

        for pred, cover in viewitems(coverage):
             if (dominant_cost <= self.comparisons[pred] and
                 dominant_cover >= cover):
                 del remaining[pred]

        return remaining

def cover(blocker, pairs, compound_length) : # pragma: no cover
    cover = coveredPairs(blocker.predicates, pairs)
    cover = compound(cover, compound_length)
    cover = remaining_cover(cover)
    return cover

def coveredPairs(predicates, pairs) :
    cover = {}

    for predicate in predicates :
        cover[predicate] = {i for i, (record_1, record_2)
                            in enumerate(pairs)
                            if (set(predicate(record_1)) &
                                set(predicate(record_2)))}

    return cover

def compound(cover, compound_length) :
    simple_predicates = sorted(cover, key=str)
    CP = predicates.CompoundPredicate

    for i in range(2, compound_length+1) :
        compound_predicates = itertools.combinations(simple_predicates, i)

        for compound_predicate in compound_predicates :
            a, b = compound_predicate[:-1], compound_predicate[-1]
            if len(a) == 1 :
                a = a[0]

            if a in cover:
                compound_cover = cover[a] & cover[b]
                if compound_cover:
                    cover[CP(compound_predicate)] = compound_cover

    return cover

def remaining_cover(coverage, covered=set()):
    remaining = {}
    for predicate, uncovered in viewitems(coverage):
        still_uncovered = uncovered - covered
        if still_uncovered:
            remaining[predicate] = still_uncovered

    return remaining

def unique(seq):
    """Return the unique elements of a collection even if those elements are
       unhashable and unsortable, like dicts and sets"""
    cleaned = []
    for each in seq:
        if each not in cleaned:
            cleaned.append(each)
    return cleaned

OUT_OF_PREDICATES_WARNING = "Ran out of predicates: Dedupe tries to find blocking rules that will work well with your data. Sometimes it can't find great ones, and you'll get this warning. It means that there are some pairs of true records that dedupe may never compare. If you are getting bad results, try increasing the `max_comparison` argument to the train method"

NO_PREDICATES_ERROR = "No predicate found! We could not learn a single good predicate. Maybe give Dedupe more training data or increasing the `max_comparisons` argument to the train method"
<commit_msg>use python 3.5 no python 3.6<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# provides functions for selecting a sample of training data
from __future__ import division
from future.utils import viewitems, viewvalues

import itertools
from . import blocking, predicates, core
import numpy
import logging

logger = logging.getLogger(__name__)

class BlockLearner(object) :
    def learn(self, matches, recall) :
        '''
        Takes in a set of training pairs and predicates and tries to find
        a good set of blocking rules.
        '''
        compound_length = 2

        self.blocker.indexAll({i : record
                               for i, record
                               in enumerate(self.unroll(matches))})

        dupe_cover = cover(self.blocker, matches, compound_length)

        comparison_count = self.comparisons(self.total_cover, compound_length)

        coverable_dupes = set.union(*viewvalues(dupe_cover))
        uncoverable_dupes = [pair for i, pair in enumerate(matches)
                             if i not in coverable_dupes]

        epsilon = int((1.0 - recall) * len(matches))

        if len(uncoverable_dupes) > epsilon :
            logger.warning(OUT_OF_PREDICATES_WARNING)
            logger.debug(uncoverable_dupes)
            epsilon = 0
        else :
            epsilon -= len(uncoverable_dupes)

        searcher = BranchBound(dupe_cover, comparison_count, epsilon, 5000)
        final_predicates = searcher.search(dupe_cover)

        logger.info('Final predicate set:')
        for predicate in final_predicates:
            logger.info(predicate)

        return final_predicates

    def comparisons(self, cover, compound_length) :
        CP = predicates.CompoundPredicate

        block_index = {}
        for predicate, blocks in viewitems(cover):
            block_index[predicate] = {}
            for block_id, blocks in viewitems(blocks) :
                for record_id in self._blocks(blocks) :
                    block_index[predicate].setdefault(record_id,
                                                      set()).add(block_id)

        compounder = self.Compounder(cover, block_index)
        comparison_count = {}
        simple_predicates = sorted(cover, key=str)

        for i in range(2, compound_length+1) :
            for combo in itertools.combinations(simple_predicates, i) :
                comparison_count[CP(combo)] = self.estimate(compounder(combo))
        for pred in simple_predicates :
            comparison_count[pred] = self.estimate(viewvalues(cover[pred]))

        return comparison_count

class Compounder(object) :
    def __init__(self, cover, block_index) :
        self.cover = cover
        self.block_index = block_index
        self._cached_predicate = None
        self._cached_blocks = None

    def __call__(self, compound_predicate) :
        a, b = compound_predicate[:-1], compound_predicate[-1]

        if len(a) > 1 :
            if a == self._cached_predicate :
                a_blocks = self._cached_blocks
            else :
                a_blocks = self._cached_blocks = list(self(a))
                self._cached_predicate = a
        else :
            a_blocks = viewvalues(self.cover[a[0]])

        return self.overlap(a_blocks, b)

class DedupeCompounder(Compounder) :
    def overlap(self, a_blocks, b) :
        b_index = self.block_index[b]
        b_index_get = b_index.get
        cover_b = self.cover[b]
        null_set = set()
        for x_ids in a_blocks :
            seen_y = set()
            for record_id in x_ids :
                b_blocks = b_index_get(record_id, null_set)
                for y in b_blocks :
                    if y not in seen_y :
                        yield x_ids & cover_b[y]
                seen_y |= b_blocks

class DedupeBlockLearner(BlockLearner) :
    Compounder = DedupeCompounder

    def __init__(self, predicates, sampled_records) :
        blocker = blocking.Blocker(predicates)
        blocker.indexAll(sampled_records)

        self.total_cover = self.coveredRecords(blocker, sampled_records)
        self.multiplier = sampled_records.original_length/len(sampled_records)

        self.blocker = blocking.Blocker(predicates)

    @staticmethod
    def unroll(matches) : # pragma: no cover
        return unique((record for pair in matches for record in pair))

    @staticmethod
    def _blocks(blocks) : # pragma: no cover
        return blocks

    @staticmethod
    def coveredRecords(blocker, records) :
        cover = {}

        for predicate in blocker.predicates :
            cover[predicate] = {}
            for id, record in viewitems(records) :
                blocks = predicate(record)
                for block in blocks :
                    cover[predicate].setdefault(block, set()).add(id)

        return cover

    def estimate(self, blocks):
        lengths = numpy.fromiter((len(ids) for ids in blocks), int)
        return numpy.sum((lengths * self.multiplier) *
                         (lengths * self.multiplier - 1) /
                         2)

class RecordLinkCompounder(Compounder) :
    def overlap(self, a_blocks, b) :
        b_index = self.block_index[b]
        b_index_get = b_index.get
        cover_b = self.cover[b]
        null_set = set()
        for first, second in a_blocks:
            seen_y = set()
            for record_id in first:
                b_blocks = b_index_get(record_id, null_set)
                for y in b_blocks:
                    if y not in seen_y :
                        yield first & cover_b[y][0], second & cover_b[y][1]
                    seen_y |= b_blocks

class RecordLinkBlockLearner(BlockLearner) :
    Compounder = RecordLinkCompounder

    def __init__(self, predicates, sampled_records_1, sampled_records_2) :
        blocker = blocking.Blocker(predicates)
        blocker.indexAll(sampled_records_2)

        self.total_cover = self.coveredRecords(blocker,
                                               sampled_records_1,
                                               sampled_records_2)

        self.multiplier_1 = sampled_records_1.original_length/len(sampled_records_1)
        self.multiplier_2 = sampled_records_2.original_length/len(sampled_records_2)

        self.blocker = blocking.Blocker(predicates)

    @staticmethod
    def unroll(matches) : # pragma: no cover
        return unique((record_2 for _, record_2 in matches))

    @staticmethod
    def _blocks(blocks) : # pragma: no cover
        return blocks[0]

    @staticmethod
    def coveredRecords(blocker, records_1, records_2) :
        cover = {}

        for predicate in blocker.predicates :
            cover[predicate] = {}
            for id, record in viewitems(records_2) :
                blocks = predicate(record)
                for block in blocks :
                    cover[predicate].setdefault(block, (set(), set()))[1].add(id)

            current_blocks = set(cover[predicate])
            for id, record in viewitems(records_1) :
                blocks = set(predicate(record))
                for block in blocks & current_blocks :
                    cover[predicate][block][0].add(id)

        return cover

    def estimate(self, blocks):
        A, B = core.iunzip(blocks, 2)
        lengths_A = numpy.fromiter((len(ids) for ids in A), float)
        lengths_B = numpy.fromiter((len(ids) for ids in B), float)
        return numpy.sum((lengths_A * self.multiplier_1) *
                         (lengths_B * self.multiplier_2))

class BranchBound(object) :
    def __init__(self, original_cover, comparison_count, epsilon, max_calls) :
        self.dupes_to_cover = set.union(*original_cover.values())
        self.original_cover = original_cover.copy()
        self.calls = max_calls
        self.comparisons = comparison_count
        self.epsilon = epsilon

        self.cheapest = tuple(original_cover.keys())
        self.cheapest_score = float('inf')

    def search(self, dupe_cover, partial=()) :
        if self.calls <= 0 :
            return self.cheapest

        self.calls -= 1

        uncovered_dupes = self.dupes_to_cover.difference(*(self.original_cover[pred]
                                                           for pred in partial))

        if len(uncovered_dupes) <= self.epsilon :
            partial_score = self.score(partial)
            if partial_score < self.cheapest_score :
                self.cheapest = partial
                self.cheapest_score = partial_score

        elif dupe_cover and (self.lower_bound(partial, dupe_cover) <= self.cheapest_score):
            cost = lambda p : self.comparisons[p]/len(dupe_cover[p])
            best_predicate = min(dupe_cover, key=cost)

            remaining = remaining_cover(dupe_cover, dupe_cover[best_predicate])

            self.search(remaining, partial + (best_predicate,))

            reduced = self.dominates(dupe_cover, best_predicate)

            uncoverable_dupes = uncovered_dupes.difference(*reduced.values())
            if len(uncoverable_dupes) <= self.epsilon:
                self.search(reduced, partial)

        return self.cheapest

    def score(self, partial):
        return sum(self.comparisons[p] for p in partial)

    def lower_bound(self, partial, dupe_cover) :
        return self.score(partial) + min(self.comparisons[p] for p in dupe_cover)

    def dominates(self, coverage, dominator):
        remaining = coverage.copy()

        dominant_cost = self.comparisons[dominator]
        dominant_cover = coverage[dominator]

        for pred, cover in viewitems(coverage):
             if (dominant_cost <= self.comparisons[pred] and
                 dominant_cover >= cover):
                 del remaining[pred]

        return remaining

def cover(blocker, pairs, compound_length) : # pragma: no cover
    cover = coveredPairs(blocker.predicates, pairs)
    cover = compound(cover, compound_length)
    cover = remaining_cover(cover)
    return cover

def coveredPairs(predicates, pairs) :
    cover = {}

    for predicate in predicates :
        cover[predicate] = {i for i, (record_1, record_2)
                            in enumerate(pairs)
                            if (set(predicate(record_1)) &
                                set(predicate(record_2)))}

    return cover

def compound(cover, compound_length) :
    simple_predicates = sorted(cover, key=str)
    CP = predicates.CompoundPredicate

    for i in range(2, compound_length+1) :
        compound_predicates = itertools.combinations(simple_predicates, i)

        for compound_predicate in compound_predicates :
            a, b = compound_predicate[:-1], compound_predicate[-1]
            if len(a) == 1 :
                a = a[0]

            if a in cover:
                compound_cover = cover[a] & cover[b]
                if compound_cover:
                    cover[CP(compound_predicate)] = compound_cover

    return cover

def remaining_cover(coverage, covered=set()):
    remaining = {}
    for predicate, uncovered in viewitems(coverage):
        still_uncovered = uncovered - covered
        if still_uncovered:
            remaining[predicate] = still_uncovered

    return remaining

def unique(seq):
    """Return the unique elements of a collection even if those elements are
       unhashable and unsortable, like dicts and sets"""
    cleaned = []
    for each in seq:
        if each not in cleaned:
            cleaned.append(each)
    return cleaned

OUT_OF_PREDICATES_WARNING = "Ran out of predicates: Dedupe tries to find blocking rules that will work well with your data. Sometimes it can't find great ones, and you'll get this warning. It means that there are some pairs of true records that dedupe may never compare. If you are getting bad results, try increasing the `max_comparison` argument to the train method"

NO_PREDICATES_ERROR = "No predicate found! We could not learn a single good predicate. Maybe give Dedupe more training data or increasing the `max_comparisons` argument to the train method"
<|endoftext|>


<commit_before>#!/usr/bin/env python
import calendar
import os
import signal
import subprocess
import sys
import time

import deimos.cleanup
import deimos.config
import deimos.containerizer
import deimos.containerizer.docker
from deimos.err import Err
import deimos.flock
from deimos.logger import log
import deimos.sig
import deimos.usage

def cli(argv=None):
    deimos.sig.install(lambda _: None)
    if argv is None:
        argv = sys.argv
    sub = argv[1] if len(argv) > 1 else None

    if sub in ["-h", "--help", "help"]:
        print format_help()
        return 0

    conf = deimos.config.load_configuration()

    if sub == "config":
        log.info("Final configuration:")
        for _, conf in conf.items():
            print "%r" % conf
        return 0

    if sub == "locks":
        deimos.flock.lock_browser(os.path.join(conf.state.root, "mesos"))
        return 0

    if sub == "state":
        cleanup = deimos.cleanup.Cleanup(conf.state.root)
        t, rm = time.time(), False
        for arg in argv[2:]:
            if arg == "--rm":
                rm = True
                continue
            t = calendar.timegm(time.strptime(arg, "%Y-%m-%dT%H:%M:%SZ"))
        if rm:
            return cleanup.remove(t)
        else:
            for d in cleanup.dirs(t):
                sys.stdout.write(d + "\n")
            return 0

    if sub not in deimos.containerizer.methods():
        print >>sys.stderr, format_help()
        print >>sys.stderr, "** Please specify a subcommand **".center(79)
        log.error("Bad ARGV: %r" % argv[1:])
        return 1

    deimos.docker.options = conf.docker.argv()
    containerizer = deimos.containerizer.docker.Docker(
        container_settings=conf.containers,
        optimistic_unpack=conf.uris.unpack,
        state_root=conf.state.root
    )

    deimos.usage.report()
    try:
        result = deimos.containerizer.stdio(containerizer, *argv[1:])
        deimos.usage.report()
        if result is not None:
            if isinstance(result, bool):
                return 0 if result else 1
            if isinstance(result, int):
                return result
            if isinstance(result, str):
                sys.stdout.write(result)
            else:
                for item in result:
                    sys.stdout.write(str(item) + "\n")
    except Err as e:
        log.error("%s.%s: %s", type(e).__module__, type(e).__name__, str(e))
        return 4
    except subprocess.CalledProcessError as e:
        log.error(str(e))
        return 4
    except Exception:
        log.exception("Unhandled failure in %s", sub)
        return 8
    return 0

def format_help():
    return """
 USAGE: deimos launch  (--no-fork)?
        deimos usage
        deimos destroy
        deimos wait    (--docker <docker-id>)?
        deimos locks
        deimos state

  Deimos provides Mesos integration for Docker, allowing Docker to be used as
  an external containerizer.

 deimos launch (--no-fork)?

  Launches a container and runs the executor or command specified in the
  TaskInfo, passed in on standard in.

  The launch subcommand always watches the launched container and logs changes
  in its lifecycle. By default, it forks off a child to do the watching, as
  part of the contract external containerizers have with Mesos. With
  --no-fork, launch will watch the container and log in the foreground. This
  can be helpful during debugging.

 deimos usage

  Generates a protobuf description of the resources used by the container.

 deimos destroy

  Shuts down the specified container.

 deimos wait (--docker <docker-id>)?

  Without --docker, reads STDIN to find the container to watch.

  With --docker, waits for a particular Docker CID to exit and records the
  result in the state directory, returning an appropriate exit code.

 deimos locks

  List file locks taken by Deimos, associating each file with a PID, an inode,
  and a lock level. The same file may appear multiple times.

 deimos state (--rm)?

  List stale state directories (those with an exit file). With --rm, removes
  stale states.

 deimos config

  Load and display the configuration.

""".strip("\n")

if __name__ == "__main__":
    sys.exit(cli(sys.argv))
<commit_msg>Fix some errors in the docs.<commit_after>#!/usr/bin/env python
import calendar
import os
import signal
import subprocess
import sys
import time

import deimos.cleanup
import deimos.config
import deimos.containerizer
import deimos.containerizer.docker
from deimos.err import Err
import deimos.flock
from deimos.logger import log
import deimos.sig
import deimos.usage

def cli(argv=None):
    deimos.sig.install(lambda _: None)
    if argv is None:
        argv = sys.argv
    sub = argv[1] if len(argv) > 1 else None

    if sub in ["-h", "--help", "help"]:
        print format_help()
        return 0

    conf = deimos.config.load_configuration()

    if sub == "config":
        log.info("Final configuration:")
        for _, conf in conf.items():
            print "%r" % conf
        return 0

    if sub == "locks":
        deimos.flock.lock_browser(os.path.join(conf.state.root, "mesos"))
        return 0

    if sub == "state":
        cleanup = deimos.cleanup.Cleanup(conf.state.root)
        t, rm = time.time(), False
        for arg in argv[2:]:
            if arg == "--rm":
                rm = True
                continue
            t = calendar.timegm(time.strptime(arg, "%Y-%m-%dT%H:%M:%SZ"))
        if rm:
            return cleanup.remove(t)
        else:
            for d in cleanup.dirs(t):
                sys.stdout.write(d + "\n")
            return 0

    if sub not in deimos.containerizer.methods():
        print >>sys.stderr, format_help()
        print >>sys.stderr, "** Please specify a subcommand **".center(79)
        log.error("Bad ARGV: %r" % argv[1:])
        return 1

    deimos.docker.options = conf.docker.argv()
    containerizer = deimos.containerizer.docker.Docker(
        container_settings=conf.containers,
        optimistic_unpack=conf.uris.unpack,
        state_root=conf.state.root
    )

    deimos.usage.report()
    try:
        result = deimos.containerizer.stdio(containerizer, *argv[1:])
        deimos.usage.report()
        if result is not None:
            if isinstance(result, bool):
                return 0 if result else 1
            if isinstance(result, int):
                return result
            if isinstance(result, str):
                sys.stdout.write(result)
            else:
                for item in result:
                    sys.stdout.write(str(item) + "\n")
    except Err as e:
        log.error("%s.%s: %s", type(e).__module__, type(e).__name__, str(e))
        return 4
    except subprocess.CalledProcessError as e:
        log.error(str(e))
        return 4
    except Exception:
        log.exception("Unhandled failure in %s", sub)
        return 8
    return 0

def format_help():
    return """
 USAGE: deimos launch (--no-fork)?
        deimos usage
        deimos destroy
        deimos wait
        deimos observe <mesos-container-id>
        deimos locks
        deimos state

  Deimos provides Mesos integration for Docker, allowing Docker to be used as
  an external containerizer.

 deimos launch (--no-fork)?

  Launches a container and runs the executor or command specified in the
  TaskInfo, passed in on standard in.

  The launch subcommand always watches the launched container and logs changes
  in its lifecycle. By default, it forks off a child to do the watching, as
  part of the contract external containerizers have with Mesos. With
  --no-fork, launch will watch the container and log in the foreground. This
  can be helpful during debugging.

 deimos usage

  Generates a protobuf description of the resources used by the container.

 deimos destroy

  Shuts down the specified container.

 deimos wait

  Reads STDIN to find the container to watch.

 deimos observe <mesos-container-id>

  Observes the Mesos container ID, in a way that blocks all calls to `wait`.
  It is for internal use...probably don't want to play with this one.

 deimos locks

  List file locks taken by Deimos, associating each file with a PID, an inode,
  and a lock level. The same file may appear multiple times.

 deimos state (--rm)?

  List stale state directories (those with an exit file). With --rm, removes
  stale states.

 deimos config

  Load and display the configuration.

""".strip("\n")

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


<commit_before><commit_msg>mr=10<commit_after><|endoftext|>


<commit_before># (C) Copyright 2015 Hewlett Packard Enterprise Development Company LP

"""Classes for monitoring the monitoring server stack.

    Covering mon-persister, mon-api and mon-thresh.
    Kafka, mysql, vertica and influxdb are covered by other detection plugins. Mon-notification uses statsd.
"""

import logging
import yaml

import monasca_setup.agent_config
import monasca_setup.detection
from monasca_setup.detection import find_process_cmdline
from monasca_setup.detection import find_process_name
from monasca_setup.detection import watch_process
from monasca_setup.detection import watch_process_by_username

log = logging.getLogger(__name__)

class MonAgent(monasca_setup.detection.Plugin):
    """Detect the Monsaca agent engine and setup some simple checks."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        self.available = True
        agent_process_list = ['monasca-collector', 'monasca-forwarder', 'monasca-statsd']
        for process in agent_process_list:
            if find_process_cmdline(process) is None:
                self.available = False
                return

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tEnabling the Monasca Agent process check")
        return watch_process_by_username('mon-agent', 'monasca-agent', 'monitoring',
                                         'monasca-agent')

    def dependencies_installed(self):
        return True

class MonAPI(monasca_setup.detection.Plugin):
    """Detect mon_api and setup monitoring."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        monasca_api = monasca_setup.detection.find_process_cmdline('monasca-api')
        if monasca_api is not None:
            # monasca-api can show up in urls and be an arg to this setup program, check port also
            # Find the right port from the config, this is specific to the Java version
            try:
                with open('/etc/monasca/api-config.yml', 'r') as config:
                    self.api_config = yaml.load(config.read())
                api_port = self.api_config['server']['applicationConnectors'][0]['port']
            except Exception:
                api_port = 8070
                log.warn('Failed parsing /etc/monasca/api-config.yml, defaulting api port to {0}'.format(api_port))
            for conn in monasca_api.connections('inet'):
                if conn.laddr[1] == api_port:
                    self.available = True
                    return

    def build_config(self):
        """Build the config as a Plugins object and return."""
        config = monasca_setup.agent_config.Plugins()

        log.info("\tEnabling the Monasca api process check")
        config.merge(watch_process(['monasca-api'], 'monitoring', 'monasca-api', exact_match=False))

        log.info("\tEnabling the Monasca api healthcheck")
        config.merge(dropwizard_health_check('monitoring', 'monasca-api', 'http://localhost:8081/healthcheck'))

        log.info("\tEnabling the Monasca api metrics")
        whitelist = [
            {
                "name": "jvm.memory.total.max",
                "path": "gauges/jvm.memory.total.max/value",
                "type": "gauge"
            },
            {
                "name": "jvm.memory.total.used",
                "path": "gauges/jvm.memory.total.used/value",
                "type": "gauge"
            },
            {
                "name": "metrics.published",
                "path": "meters/monasca.api.app.MetricService.metrics.published/count",
                "type": "rate"
            }
        ]

        if not self._is_hibernate_on():
            # if hibernate is not used, it is mysql with DBI
            # for that case having below entries makes sense
            log.debug('MonApi has not enabled Hibernate, adding DBI metrics')
            whitelist.extend([
                {
                    "name": "raw-sql.time.avg",
                    "path": "timers/org.skife.jdbi.v2.DBI.raw-sql/mean",
                    "type": "gauge"
                },
                {
                    "name": "raw-sql.time.max",
                    "path": "timers/org.skife.jdbi.v2.DBI.raw-sql/max",
                    "type": "gauge"
                }
            ])

        config.merge(dropwizard_metrics('monitoring',
                                        'monasca-api',
                                        'http://localhost:8081/metrics',
                                        whitelist))

        return config

    def dependencies_installed(self):
        return True

    def _is_hibernate_on(self):
        # check if api_config has been declared in __init__
        # if not it means that configuration file was not found
        # or monasca-api Python implementation is running

        api_config = getattr(self, 'api_config', None)
        if api_config is None:
            return False

        hibernate_cfg = self.api_config.get('hibernate', None)
        if hibernate_cfg is None:
            return False

        return hibernate_cfg.get('supportEnabled', False)

class MonNotification(monasca_setup.detection.Plugin):
    """Detect the Monsaca notification engine and setup some simple checks."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        if find_process_cmdline('monasca-notification') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tEnabling the Monasca Notification healthcheck")
        return watch_process_by_username('mon-notification', 'monasca-notification', 'monitoring',
                                         'monasca-notification')

    def dependencies_installed(self):
        return True

class MonPersister(monasca_setup.detection.Plugin):
    """Detect mon_persister and setup monitoring."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        if find_process_cmdline('monasca-persister') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        """Read persister-config.yml file to find the exact numThreads."""
        try:
            with open('/etc/monasca/persister-config.yml', 'r') as config:
                self.persister_config = yaml.load(config.read())
        except Exception:
            log.exception('Failed parsing /etc/monasca/persister-config.yml')
            self.available = False
            return

        alarm_num_threads = self.persister_config['alarmHistoryConfiguration']['numThreads']
        metric_num_threads = self.persister_config['metricConfiguration']['numThreads']

        database_type = self.persister_config['databaseConfiguration']['databaseType']

        config = monasca_setup.agent_config.Plugins()

        log.info("\tEnabling the Monasca persister process check")
        config.merge(watch_process(['monasca-persister'], 'monitoring', 'monasca-persister', exact_match=False))

        log.info("\tEnabling the Monasca persister healthcheck")
        config.merge(dropwizard_health_check('monitoring', 'monasca-persister', 'http://localhost:8091/healthcheck'))

        log.info("\tEnabling the Monasca persister metrics")
        whitelist = [
            {
                "name": "jvm.memory.total.max",
                "path": "gauges/jvm.memory.total.max/value",
                "type": "gauge"},
            {
                "name": "jvm.memory.total.used",
                "path": "gauges/jvm.memory.total.used/value",
                "type": "gauge"}
        ]

        # Generate initial whitelist based on the database type
        if database_type == 'influxdb':
            pass
        elif database_type == 'vertica':
            whitelist.extend([
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-hit-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-hit-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-miss-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-miss-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-hit-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-hit-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-miss-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-miss-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-hit-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-hit-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-miss-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-miss-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.measurement-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.measurement-meter/count",
                    "type": "rate"}
            ])
        else:
            log.warn('Failed finding database type in /etc/monasca/persister-config.yml')

        # Dynamic Whitelist
        for idx in range(alarm_num_threads):
            new_thread = {"name": "alarm-state-transitions-added-to-batch-counter[{0}]".format(idx),
                          "path": "counters/monasca.persister.pipeline.event.AlarmStateTransitionHandler[alarm-state-transition-{0}].alarm-state-transitions-added-to-batch-counter/count".format(idx),
                          "type": "rate"
                          }
            whitelist.append(new_thread)

        for idx in range(metric_num_threads):
            new_thread = {"name": "metrics-added-to-batch-counter[{0}]".format(idx),
                          "path": "counters/monasca.persister.pipeline.event.MetricHandler[metric-{0}].metrics-added-to-batch-counter/count".format(idx),
                          "type": "rate"
                          }
            whitelist.append(new_thread)

        config.merge(dropwizard_metrics('monitoring', 'monasca-persister', 'http://localhost:8091/metrics', whitelist))
        return config

    def dependencies_installed(self):
        return True

class MonThresh(monasca_setup.detection.Plugin):
    """Detect the running mon-thresh and monitor."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        if find_process_cmdline('backtype.storm.daemon') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tWatching the mon-thresh process.")
        config = monasca_setup.agent_config.Plugins()
        for process in ['backtype.storm.daemon.nimbus', 'backtype.storm.daemon.supervisor', 'backtype.storm.daemon.worker']:
            if find_process_cmdline(process) is not None:
                config.merge(watch_process([process], 'monitoring', 'apache-storm', exact_match=False, detailed=False))
        config.merge(watch_process_by_username('storm', 'monasca-thresh', 'monitoring', 'apache-storm'))
        return config

    def dependencies_installed(self):
        return True

def dropwizard_health_check(service, component, url):
    """Setup a dropwizard heathcheck to be watched by the http_check plugin."""
    config = monasca_setup.agent_config.Plugins()
    config['http_check'] = {'init_config': None,
                            'instances': [{'name': "{0}-{1} healthcheck".format(service, component),
                                           'url': url,
                                           'timeout': 1,
                                           'include_content': False,
                                           'dimensions': {'service': service, 'component': component}}]}
    return config

def dropwizard_metrics(service, component, url, whitelist):
    """Setup a dropwizard metrics check"""
    config = monasca_setup.agent_config.Plugins()
    config['http_metrics'] = {'init_config': None,
                              'instances': [{'name': "{0}-{1} metrics".format(service, component),
                                             'url': url,
                                             'timeout': 1,
                                             'dimensions': {'service': service, 'component': component},
                                             'whitelist': whitelist}]}
    return config

class MonInfluxDB(monasca_setup.detection.Plugin):
    """Detect InfluxDB and setup some simple checks."""

    def _detect(self):
        """Run detection, set self.available True if the service is detected.
        """

        if find_process_name('influxd') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tEnabling the Monasca InfluxDB check")
        return watch_process(['influxd'], 'monitoring', 'influxd',
                             exact_match=False)

    def dependencies_installed(self):
        return True
<commit_msg>Increasing Monasca healthcheck timeout<commit_after># (C) Copyright 2015-2016 Hewlett Packard Enterprise Development Company LP

"""Classes for monitoring the monitoring server stack.

    Covering mon-persister, mon-api and mon-thresh.
    Kafka, mysql, vertica and influxdb are covered by other detection plugins. Mon-notification uses statsd.
"""

import logging
import yaml

import monasca_setup.agent_config
import monasca_setup.detection
from monasca_setup.detection import find_process_cmdline
from monasca_setup.detection import find_process_name
from monasca_setup.detection import watch_process
from monasca_setup.detection import watch_process_by_username

log = logging.getLogger(__name__)

class MonAgent(monasca_setup.detection.Plugin):
    """Detect the Monsaca agent engine and setup some simple checks."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        self.available = True
        agent_process_list = ['monasca-collector', 'monasca-forwarder', 'monasca-statsd']
        for process in agent_process_list:
            if find_process_cmdline(process) is None:
                self.available = False
                return

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tEnabling the Monasca Agent process check")
        return watch_process_by_username('mon-agent', 'monasca-agent', 'monitoring',
                                         'monasca-agent')

    def dependencies_installed(self):
        return True

class MonAPI(monasca_setup.detection.Plugin):
    """Detect mon_api and setup monitoring."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        monasca_api = monasca_setup.detection.find_process_cmdline('monasca-api')
        if monasca_api is not None:
            # monasca-api can show up in urls and be an arg to this setup program, check port also
            # Find the right port from the config, this is specific to the Java version
            try:
                with open('/etc/monasca/api-config.yml', 'r') as config:
                    self.api_config = yaml.load(config.read())
                api_port = self.api_config['server']['applicationConnectors'][0]['port']
            except Exception:
                api_port = 8070
                log.warn('Failed parsing /etc/monasca/api-config.yml, defaulting api port to {0}'.format(api_port))
            for conn in monasca_api.connections('inet'):
                if conn.laddr[1] == api_port:
                    self.available = True
                    return

    def build_config(self):
        """Build the config as a Plugins object and return."""
        config = monasca_setup.agent_config.Plugins()

        log.info("\tEnabling the Monasca api process check")
        config.merge(watch_process(['monasca-api'], 'monitoring', 'monasca-api', exact_match=False))

        log.info("\tEnabling the Monasca api healthcheck")
        config.merge(dropwizard_health_check('monitoring', 'monasca-api', 'http://localhost:8081/healthcheck'))

        log.info("\tEnabling the Monasca api metrics")
        whitelist = [
            {
                "name": "jvm.memory.total.max",
                "path": "gauges/jvm.memory.total.max/value",
                "type": "gauge"
            },
            {
                "name": "jvm.memory.total.used",
                "path": "gauges/jvm.memory.total.used/value",
                "type": "gauge"
            },
            {
                "name": "metrics.published",
                "path": "meters/monasca.api.app.MetricService.metrics.published/count",
                "type": "rate"
            }
        ]

        if not self._is_hibernate_on():
            # if hibernate is not used, it is mysql with DBI
            # for that case having below entries makes sense
            log.debug('MonApi has not enabled Hibernate, adding DBI metrics')
            whitelist.extend([
                {
                    "name": "raw-sql.time.avg",
                    "path": "timers/org.skife.jdbi.v2.DBI.raw-sql/mean",
                    "type": "gauge"
                },
                {
                    "name": "raw-sql.time.max",
                    "path": "timers/org.skife.jdbi.v2.DBI.raw-sql/max",
                    "type": "gauge"
                }
            ])

        config.merge(dropwizard_metrics('monitoring',
                                        'monasca-api',
                                        'http://localhost:8081/metrics',
                                        whitelist))

        return config

    def dependencies_installed(self):
        return True

    def _is_hibernate_on(self):
        # check if api_config has been declared in __init__
        # if not it means that configuration file was not found
        # or monasca-api Python implementation is running

        api_config = getattr(self, 'api_config', None)
        if api_config is None:
            return False

        hibernate_cfg = self.api_config.get('hibernate', None)
        if hibernate_cfg is None:
            return False

        return hibernate_cfg.get('supportEnabled', False)

class MonNotification(monasca_setup.detection.Plugin):
    """Detect the Monsaca notification engine and setup some simple checks."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        if find_process_cmdline('monasca-notification') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tEnabling the Monasca Notification healthcheck")
        return watch_process_by_username('mon-notification', 'monasca-notification', 'monitoring',
                                         'monasca-notification')

    def dependencies_installed(self):
        return True

class MonPersister(monasca_setup.detection.Plugin):
    """Detect mon_persister and setup monitoring."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        if find_process_cmdline('monasca-persister') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        """Read persister-config.yml file to find the exact numThreads."""
        try:
            with open('/etc/monasca/persister-config.yml', 'r') as config:
                self.persister_config = yaml.load(config.read())
        except Exception:
            log.exception('Failed parsing /etc/monasca/persister-config.yml')
            self.available = False
            return

        alarm_num_threads = self.persister_config['alarmHistoryConfiguration']['numThreads']
        metric_num_threads = self.persister_config['metricConfiguration']['numThreads']

        database_type = self.persister_config['databaseConfiguration']['databaseType']

        config = monasca_setup.agent_config.Plugins()

        log.info("\tEnabling the Monasca persister process check")
        config.merge(watch_process(['monasca-persister'], 'monitoring', 'monasca-persister', exact_match=False))

        log.info("\tEnabling the Monasca persister healthcheck")
        config.merge(dropwizard_health_check('monitoring', 'monasca-persister', 'http://localhost:8091/healthcheck'))

        log.info("\tEnabling the Monasca persister metrics")
        whitelist = [
            {
                "name": "jvm.memory.total.max",
                "path": "gauges/jvm.memory.total.max/value",
                "type": "gauge"},
            {
                "name": "jvm.memory.total.used",
                "path": "gauges/jvm.memory.total.used/value",
                "type": "gauge"}
        ]

        # Generate initial whitelist based on the database type
        if database_type == 'influxdb':
            pass
        elif database_type == 'vertica':
            whitelist.extend([
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-hit-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-hit-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-miss-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-cache-miss-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-hit-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-hit-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-miss-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.definition-dimension-cache-miss-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-hit-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-hit-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-miss-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.dimension-cache-miss-meter/count",
                    "type": "rate"},
                {
                    "name": "monasca.persister.repository.vertica.VerticaMetricRepo.measurement-meter",
                    "path": "meters/monasca.persister.repository.vertica.VerticaMetricRepo.measurement-meter/count",
                    "type": "rate"}
            ])
        else:
            log.warn('Failed finding database type in /etc/monasca/persister-config.yml')

        # Dynamic Whitelist
        for idx in range(alarm_num_threads):
            new_thread = {"name": "alarm-state-transitions-added-to-batch-counter[{0}]".format(idx),
                          "path": "counters/monasca.persister.pipeline.event.AlarmStateTransitionHandler[alarm-state-transition-{0}].alarm-state-transitions-added-to-batch-counter/count".format(idx),
                          "type": "rate"
                          }
            whitelist.append(new_thread)

        for idx in range(metric_num_threads):
            new_thread = {"name": "metrics-added-to-batch-counter[{0}]".format(idx),
                          "path": "counters/monasca.persister.pipeline.event.MetricHandler[metric-{0}].metrics-added-to-batch-counter/count".format(idx),
                          "type": "rate"
                          }
            whitelist.append(new_thread)

        config.merge(dropwizard_metrics('monitoring', 'monasca-persister', 'http://localhost:8091/metrics', whitelist))
        return config

    def dependencies_installed(self):
        return True

class MonThresh(monasca_setup.detection.Plugin):
    """Detect the running mon-thresh and monitor."""
    def _detect(self):
        """Run detection, set self.available True if the service is detected."""
        if find_process_cmdline('backtype.storm.daemon') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tWatching the mon-thresh process.")
        config = monasca_setup.agent_config.Plugins()
        for process in ['backtype.storm.daemon.nimbus', 'backtype.storm.daemon.supervisor', 'backtype.storm.daemon.worker']:
            if find_process_cmdline(process) is not None:
                config.merge(watch_process([process], 'monitoring', 'apache-storm', exact_match=False, detailed=False))
        config.merge(watch_process_by_username('storm', 'monasca-thresh', 'monitoring', 'apache-storm'))
        return config

    def dependencies_installed(self):
        return True

def dropwizard_health_check(service, component, url):
    """Setup a dropwizard heathcheck to be watched by the http_check plugin."""
    config = monasca_setup.agent_config.Plugins()
    config['http_check'] = {'init_config': None,
                            'instances': [{'name': "{0}-{1} healthcheck".format(service, component),
                                           'url': url,
                                           'timeout': 5,
                                           'include_content': False,
                                           'dimensions': {'service': service, 'component': component}}]}
    return config

def dropwizard_metrics(service, component, url, whitelist):
    """Setup a dropwizard metrics check"""
    config = monasca_setup.agent_config.Plugins()
    config['http_metrics'] = {'init_config': None,
                              'instances': [{'name': "{0}-{1} metrics".format(service, component),
                                             'url': url,
                                             'timeout': 5,
                                             'dimensions': {'service': service, 'component': component},
                                             'whitelist': whitelist}]}
    return config

class MonInfluxDB(monasca_setup.detection.Plugin):
    """Detect InfluxDB and setup some simple checks."""

    def _detect(self):
        """Run detection, set self.available True if the service is detected.
        """

        if find_process_name('influxd') is not None:
            self.available = True

    def build_config(self):
        """Build the config as a Plugins object and return."""
        log.info("\tEnabling the Monasca InfluxDB check")
        return watch_process(['influxd'], 'monitoring', 'influxd',
                             exact_match=False)

    def dependencies_installed(self):
        return True
<|endoftext|>


<commit_before>"""
A script to show off certain Raspberry Pi features automatically,
without needing to interact with the Pi to start commands.

I'm intending this to be run at startup, perhaps if no keyboard
is connected?
"""

import os.path
import subprocess

hat_product_file = "/proc/device-tree/hat/product"

def get_hat_name():
    if not os.path.exists(hat_product_file):
        return None

    with file(hat_product_file) as f:
        line = f.read()
        return line

hat_name = get_hat_name()
if not hat_name:
    print "No HAT connected"
    exit
print hat_name

script_name = None
if hat_name.startswith("Unicorn HAT"):
    script_name = "/home/pi/Pimoroni/unicornhat/rainbow.py"

print script_name
if os.path.exists(script_name):
    subprocess.call(['sudo', 'python', script_name])
<commit_msg>Record how to run the script via cron - not yet tested.<commit_after>"""
A script to show off certain Raspberry Pi features automatically,
without needing to interact with the Pi to start commands.

I'm intending this to be run at startup, perhaps if no keyboard
is connected?

To run this automatically, see
http://www.raspberrypi-spy.co.uk/2013/07/running-a-python-script-at-boot-using-cron/

Add this line to the root contab file
@reboot python /home/pi/develop/raspi_code/demo/demo_my_pi.py &
"""

import os.path
import subprocess

hat_product_file = "/proc/device-tree/hat/product"

def get_hat_name():
    if not os.path.exists(hat_product_file):
        return None

    with file(hat_product_file) as f:
        line = f.read()
        return line

hat_name = get_hat_name()
if not hat_name:
    print "No HAT connected"
    exit
print hat_name

script_name = None
if hat_name.startswith("Unicorn HAT"):
    script_name = "/home/pi/Pimoroni/unicornhat/rainbow.py"

print script_name
if os.path.exists(script_name):
    subprocess.call(['sudo', 'python', script_name])
<|endoftext|>