<commit_before>#!/usr/bin/env python
import re, glob, sys,os
import argparse
import itertools as it
import numpy as np
import scipy.sparse
import cPickle as pickle
from sklearn.naive_bayes import MultinomialNB
from content_sources import arxiv, bibtex, rss

def find_keywords(text):

    keywords=re.sub('[{}:?!@#$%^&*\(\)_.\\/,\'\"]','',text).upper()
    prepositions = open('data/prepositions.dat').read().upper().split()

    for p in prepositions:
        keywords = re.sub(r"\b{!s}\b".format(p),' ',keywords)

    return keywords.split()

def test(argv):
    #train the algorithm
    #read in titles we like
    good_titles_kw = find_keywords(open('data/allTitles.dat').read())
    #read in titles we don't
    bad_titles_kw = find_keywords(open('data/arxivTitles.500.dat').read())

    #remore duplicate words
    good_kw_set = set(good_titles_kw)
    bad_kw_set = set(bad_titles_kw)

    #make a set of all the words
    all_kw_set = good_kw_set.union(bad_kw_set)

    #count the occurences of each of the the keywords
    good_freq = np.zeros(len(all_kw_set))
    bad_freq = np.zeros(len(all_kw_set))
    all_freq = np.zeros(len(all_kw_set))

    ##good_txt = ' '.join(good_titles_kw)
    ##bad_txt = ' '.join(bad_titles_kw)

    for i,kw in enumerate(all_kw_set):
        good_freq[i] += good_titles_kw.count(kw)# len(re.findall(r"\b{!s}\b".format(kw), good_txt))
        bad_freq[i] += bad_titles_kw.count(kw) # len(re.findall(r"\b{!s}\b".format(kw), bad_txt))

    all_freq = good_freq+bad_freq

    #compute the base rates of words in the good or bad set
    base_good = np.sum(good_freq)/np.sum(all_freq)
    base_bad = np.sum(bad_freq)/np.sum(all_freq)

    #compute the "evidence probabilites"
    evidence_prob = all_freq/np.sum(all_freq)

    #compute likelihood of word being good and bad
    good_likelihood  = good_freq/all_freq
    bad_likelihood  = bad_freq/all_freq

    #make a dictionary to easily access likelihoods of our keywords
    gl_dict = {kw:lh for kw,lh in zip(all_kw_set,good_likelihood)}
    bl_dict = {kw:lh for kw,lh in zip(all_kw_set,bad_likelihood)}

    #read in new titles
    new_titles = open('data/arxivTitles.feb20.300.dat').readlines()
    nt_good_scores = np.empty(len(new_titles))
    nt_bad_scores = np.empty(len(new_titles))

    #compute score for each new title
    for i,nt in enumerate(new_titles):
        nt_kw_set = set(find_keywords(nt))
        nt_good_scores[i] = base_good*np.prod([gl_dict[kw] for kw in nt_kw_set if kw in gl_dict])
        nt_bad_scores[i] = base_bad*np.prod([bl_dict[kw] for kw in nt_kw_set if kw in bl_dict])

    for b,g,t in zip(nt_bad_scores,nt_good_scores,new_titles):
        print(b,g,'GOOD!' if g>b and g>0 else 'BAD!',t)

#read existing knowledge from a database
def parse_knowledge(knowledge_db):
    pass

# \param nb = MultinomialNB classifier
def filter_content(content,method,naive_bayes,keywords):

    new_samples = [find_keywords(entry[method])  for entry in content]

    #compute vector for each new entry
    X = scipy.sparse.lil_matrix((len(new_samples),len(keywords)))

    for j,kw in enumerate(keywords):
        for i,gs in enumerate(new_samples):
            if kw in gs:
                X[i,j]+=1

    categories = naive_bayes.predict(X)
    return [e for c,e in zip(categories,content) if c =='good']

def get_content(sources):
    all_content = list()
    for src in sources:
        try:
            src.fetch()
        except:
            print("Fetch of content from {!r} has failed".format(src))

        content = None
        try:
            content = src.parse()
        except:
            print("parsing of content from {!r} has failed".format(src))

        if content:
            all_content += content

    return all_content

def load_knowledge(knowledge):
    #existing naive_bayes object and keyword list
    nb=None
    kw=list()
    if knowledge is not None:
        if not os.path.isdir(knowledge):
            print("Knowledge bust be a directory")
            exit()

        kfiles = glob.glob(knowledge+'/*')
        if not 'nb.p' in kfiles:
            print("Knowledge does not contain nb.p (pickled naive bayes object)")
            exit()
        if not 'kw.p' in kfiles:
            print("Knowledge does not contain kw.p (pickled keyword list)")
            exit()

    else:
        knowledge =os.path.expanduser('~/.shakespeare')

    if os.path.exists(knowledge):
        nb=pickle.load(open(knowledge+'/nb.p'))
        kw=pickle.load(open(knowledge+'/kw.p'))

    return(nb,kw)

#training suite
def train(good_sources, bad_sources,method,naive_bayes=None,keywords=list()):
    #train the algorithm
    good_samples = [find_keywords(entry[method]) for entry in good_sources]
    bad_samples = [find_keywords(entry[method]) for entry in bad_sources]

    new_kws = set([kw for kw in it.chain(*good_samples)]+ [kw for kw in it.chain(*bad_samples)])

    all_kw = keywords + [nkw for nkw in new_kws if nkw not in keywords]

    X = scipy.sparse.lil_matrix((len(good_samples) + len(bad_samples),len(all_kw)))

    offset = len(good_samples)
    for j,kw in enumerate(all_kw):
        for i,gs in enumerate(good_samples):
            X[i,j]=gs.count(kw)
        for i,bs in enumerate(bad_samples):
            X[i+offset,j]=bs.count(kw)
    y = ['good']*len(good_samples) + ['bad']*len(bad_samples)

    if naive_bayes is None:
        nb = MultinomialNB()
        nb.fit(X,y)
    else:
        nb = naive_bayes
        nb.partial_fit(X,y,classes=['good','bad'])

    return nb, all_kw

def to_markdown(content,output_file):
    try:
        with open(output_file,'w') as outf:
            outf.write('# Relevant articles\n')
            for article in content:
                outf.write("## {}\n".format(article['title']))
                outf.write("* authors: {}\n".format(article['author']))
                outf.write("* abstract: {}\n".format(article['abstract']))
                outf.write("* [link]({})\n\n".format(article['url']))
    except:
        print("Failed to write markdown file")

def review_content(good_content,content,method,review_all=False):
    to_review=[]
    if review_all:
        to_review = content
    else:
        to_review = good_content

    human_class=[]
    for entry in to_review:
        print("Is \"{}\" a good entry?".format(entry[method]))
        decision = raw_input('Y/n?').lower()
        human_class.append(decision=='y')
    return human_class

def main(argv):

    #add command line options for sources, output prefs, database of "good" keywords
    parser = argparse.ArgumentParser()
    parser.add_argument('-o','--output', help='output file name. only supports markdown right now.',dest ='output',default=None)
    parser.add_argument('-b','--bibtex', help='bibtex files to fetch',dest='bibfiles', nargs='*',default=list())
    parser.add_argument('-j','--journals', help='journals to fetch. Currently supports {}.'.format(' '.join(rss.rss_feeds.keys())),
                        nargs='*',dest='journals',default=list())
    parser.add_argument('-a','--arXiv', help='arXiv categories to fetch',
                      nargs='*',dest='arXiv',default=list())
    parser.add_argument('--all_sources', help='flag to search from all sources.',action ='store_true')
    parser.add_argument('--all_good_sources', help='flag to search from good sources. Specfied in your config file.',action ='store_true')
    parser.add_argument('--train', help='flag to train. All sources beside "--train-input-good" are treated as bad/irrelevant papers',action ='store_true')
    parser.add_argument('-g','--train_input_good', help='bibtex file containing relevant articles.',dest ='good_source',default=None)
    parser.add_argument('-m','--method', help='Methods to try to find relevent papers. Right now, only all, title, author, and abstract are valid fields',
                        dest='method',default='title')
    parser.add_argument('-k', '--knowledge',
                        help='path to database containing information about good and bad keywords. \
                              If you are training, you must specifiy this, as it will be where your output is written ',
                        dest='knowledge',default=None)
    parser.add_argument('--overwrite-knowledge', help='flag to overwrite knowledge,if training',action ='store_true',default=True, dest='overwrite_knowledge')
    parser.add_argument('--feedback', help='flag to give feedback after sorting content',action ='store_true',default=False, dest='feedback')
    parser.add_argument('--review_all', help='review all the new selections. Otherwise, you will only review the good selections',action ='store_true',default=False, dest='review_all')
    args = parser.parse_args(argv)

    if  not args.method in ['title','abstract','author, all']:
            print("Invalid method. Options are title, abstract, author, and all")
            exit()

    method = args.method

    #Set up training if that's what we're doing
    if args.train:

        #check to make sure we have a good training input
        if args.good_source is None:
            print("When training, you must specify one good source")
            exit()
        if not os.path.exists(args.good_source):
            print("Specified training input does not exist")
            exit()
        if not os.path.isfile(args.good_source):
            print("Specified training input is not a file")
            exit()
        if not os.path.splitext(args.good_source)[1] == '.bib' :
            print("Training input must be in bibtex format")
            exit()

        #load the existing knowledge
        nb,kw = load_knowledge(args.knowledge)

        #make sure you have a good destination to write your training output

        good_content = get_content([bibtex.BibTex(args.good_source)])

        bad_content = get_content([arxiv.ArXiv(cat) for cat in args.arXiv] +
                                  [bibtex.BibTex(bibfile) for bibfile in args.bibfiles] +
                                  [rss.JournalFeed(journal) for journal in args.journals])

        #train, and write out knowledge (naive_bayes class and keywords)
        nb, kw = train(good_content,bad_content,method,naive_bayes=nb, keywords=kw)
        pickle.dump(nb,open(knowledge+'/nb.p','w'))
        pickle.dump(kw,open(knowledge+'/kw.p','w'))

    #we are filtering new content through our existing knowledge
    else:

        #load the old knowledge
        nb,kw = load_knowledge(args.knowledge)
        sources = [arxiv.ArXiv(cat) for cat in args.arXiv] + \
                  [ bibtex.BibTex (bibfile) for bibfile in args.bibfiles] + \
                  [rss.JournalFeed(journal) for journal in args.journals]

        new_content = get_content(sources)

        good_content = filter_content(new_content,method,nb,kw)

        if (args.output):
            to_markdown(good_content,args.output)
        else:
            print(good_content)

        if(args.feedback):
            review_content(good_content,new_content,method,args.review_all)

if __name__=="__main__":
    main(sys.argv[1:])
<commit_msg>completed human feedback loop<commit_after>#!/usr/bin/env python
import re, glob, sys,os
import argparse
import itertools as it
import numpy as np
import scipy.sparse
import cPickle as pickle
from sklearn.naive_bayes import MultinomialNB
from content_sources import arxiv, bibtex, rss

def find_keywords(text):

    keywords=re.sub('[{}:?!@#$%^&*\(\)_.\\/,\'\"]','',text).upper()
    prepositions = open('data/prepositions.dat').read().upper().split()

    for p in prepositions:
        keywords = re.sub(r"\b{!s}\b".format(p),' ',keywords)

    return keywords.split()

def test(argv):
    #train the algorithm
    #read in titles we like
    good_titles_kw = find_keywords(open('data/allTitles.dat').read())
    #read in titles we don't
    bad_titles_kw = find_keywords(open('data/arxivTitles.500.dat').read())

    #remore duplicate words
    good_kw_set = set(good_titles_kw)
    bad_kw_set = set(bad_titles_kw)

    #make a set of all the words
    all_kw_set = good_kw_set.union(bad_kw_set)

    #count the occurences of each of the the keywords
    good_freq = np.zeros(len(all_kw_set))
    bad_freq = np.zeros(len(all_kw_set))
    all_freq = np.zeros(len(all_kw_set))

    ##good_txt = ' '.join(good_titles_kw)
    ##bad_txt = ' '.join(bad_titles_kw)

    for i,kw in enumerate(all_kw_set):
        good_freq[i] += good_titles_kw.count(kw)# len(re.findall(r"\b{!s}\b".format(kw), good_txt))
        bad_freq[i] += bad_titles_kw.count(kw) # len(re.findall(r"\b{!s}\b".format(kw), bad_txt))

    all_freq = good_freq+bad_freq

    #compute the base rates of words in the good or bad set
    base_good = np.sum(good_freq)/np.sum(all_freq)
    base_bad = np.sum(bad_freq)/np.sum(all_freq)

    #compute the "evidence probabilites"
    evidence_prob = all_freq/np.sum(all_freq)

    #compute likelihood of word being good and bad
    good_likelihood  = good_freq/all_freq
    bad_likelihood  = bad_freq/all_freq

    #make a dictionary to easily access likelihoods of our keywords
    gl_dict = {kw:lh for kw,lh in zip(all_kw_set,good_likelihood)}
    bl_dict = {kw:lh for kw,lh in zip(all_kw_set,bad_likelihood)}

    #read in new titles
    new_titles = open('data/arxivTitles.feb20.300.dat').readlines()
    nt_good_scores = np.empty(len(new_titles))
    nt_bad_scores = np.empty(len(new_titles))

    #compute score for each new title
    for i,nt in enumerate(new_titles):
        nt_kw_set = set(find_keywords(nt))
        nt_good_scores[i] = base_good*np.prod([gl_dict[kw] for kw in nt_kw_set if kw in gl_dict])
        nt_bad_scores[i] = base_bad*np.prod([bl_dict[kw] for kw in nt_kw_set if kw in bl_dict])

    for b,g,t in zip(nt_bad_scores,nt_good_scores,new_titles):
        print(b,g,'GOOD!' if g>b and g>0 else 'BAD!',t)

#read existing knowledge from a database
def parse_knowledge(knowledge_db):
    pass

# \param nb = MultinomialNB classifier
def filter_content(content,
                   method,
                   naive_bayes,
                   keywords):

    new_samples = [find_keywords(entry[method]) for entry in content]

    #compute vector for each new entry
    X = scipy.sparse.lil_matrix((len(new_samples),len(keywords)))

    for j,kw in enumerate(keywords):
        for i,ns in enumerate(new_samples):
            X[i,j]=ns.count(kw)

    categories = naive_bayes.predict(X)
    return [e for c,e in zip(categories,content) if c =='good']

def get_content(sources):
    all_content = list()
    for src in sources:
        try:
            src.fetch()
        except:
            print("Fetch of content from {!r} has failed".format(src))

        content = None
        try:
            content = src.parse()
        except:
            print("parsing of content from {!r} has failed".format(src))

        if content:
            all_content += content

    return all_content

def load_knowledge(knowledge):
    #existing naive_bayes object and keyword list
    nb=None
    kw=list()
    if knowledge is not None:
        if not os.path.isdir(knowledge):
            print("Knowledge bust be a directory")
            exit()

        kfiles = glob.glob(knowledge+'/*')
        if not 'nb.p' in kfiles:
            print("Knowledge does not contain nb.p (pickled naive bayes object)")
            exit()
        if not 'kw.p' in kfiles:
            print("Knowledge does not contain kw.p (pickled keyword list)")
            exit()

    else:
        knowledge =os.path.expanduser('~/.shakespeare')

    if os.path.exists(knowledge):
        nb=pickle.load(open(knowledge+'/nb.p'))
        kw=pickle.load(open(knowledge+'/kw.p'))

    return(nb,kw, knowledge)

#training suite
def train(good_sources, bad_sources,method,naive_bayes=None,keywords=list()):
    #train the algorithm
    good_samples = find_keywords(' '.join([entry[method] for entry in good_sources]))
    bad_samples = find_keywords(' '.join([entry[method] for entry in bad_sources]))

    #if we have an exists knowledge base to append this new information to, do so
    if naive_bayes:
        new_kws = set(good_samples+bad_samples).difference(keywords)

        #make some call to naive_bayes.partial_fit in here
        X = np.concatenate((naive_bayes.feature_count_, np.zeros((naive_bayes.feature_count_.shape[0],naive_bayes.feature_count_.shape[1]+len(new_kws)))),1)
        keywords += list(new_kws)

        all_kw = keywords + [nkw for nkw in new_kws if nkw not in keywords]

    else:
        all_kw = list(set(good_samples+bad_samples))
        X = np.zeros((2,len(all_kw)))

    for j,kw in enumerate(all_kw):
        X[0,j] += good_samples.count(kw)
        X[1,j] += bad_samples.count(kw)

    y = ['good','bad']

    naive_bayes = MultinomialNB()
    naive_bayes.fit(X,y)

    return naive_bayes, all_kw

def to_markdown(content,output_file):
    try:
        with open(output_file,'w') as outf:
            outf.write('# Relevant articles\n')
            for article in content:
                outf.write("## {}\n".format(article['title']))
                outf.write("* authors: {}\n".format(article['author']))
                outf.write("* abstract: {}\n".format(article['abstract']))
                outf.write("* [link]({})\n\n".format(article['url']))
    except:
        print("Failed to write markdown file")

def review_content(good_content,content,method,review_all=False):
    to_review=[]
    if review_all:
        to_review = content
    else:
        to_review = good_content

    human_class=[]
    for entry in to_review:
        print("Is \"{}\" a good entry?".format(entry[method]))
        decision = raw_input('Y/n?').lower()
        human_class.append('good' if decision=='y' else 'bad')
    return human_class, to_review

def main(argv):

    #add command line options for sources, output prefs, database of "good" keywords
    parser = argparse.ArgumentParser()
    parser.add_argument('-o','--output', help='output file name. only supports markdown right now.',dest ='output',default=None)
    parser.add_argument('-b','--bibtex', help='bibtex files to fetch',dest='bibfiles', nargs='*',default=list())
    parser.add_argument('-j','--journals', help='journals to fetch. Currently supports {}.'.format(' '.join(rss.rss_feeds.keys())),
                        nargs='*',dest='journals',default=list())
    parser.add_argument('-a','--arXiv', help='arXiv categories to fetch',
                      nargs='*',dest='arXiv',default=list())
    parser.add_argument('--all_sources', help='flag to search from all sources.',action ='store_true')
    parser.add_argument('--all_good_sources', help='flag to search from good sources. Specfied in your config file.',action ='store_true')
    parser.add_argument('--train', help='flag to train. All sources beside "--train-input-good" are treated as bad/irrelevant papers',action ='store_true')
    parser.add_argument('-g','--train_input_good', help='bibtex file containing relevant articles.',dest ='good_source',default=None)
    parser.add_argument('-m','--method', help='Methods to try to find relevent papers. Right now, only all, title, author, and abstract are valid fields',
                        dest='method',default='title')
    parser.add_argument('-k', '--knowledge',
                        help='path to database containing information about good and bad keywords. \
                              If you are training, you must specifiy this, as it will be where your output is written ',
                        dest='knowledge',default=None)
    parser.add_argument('--overwrite-knowledge', help='flag to overwrite knowledge,if training',action ='store_true',default=True, dest='overwrite_knowledge')
    parser.add_argument('--feedback', help='flag to give feedback after sorting content',action ='store_true',default=False, dest='feedback')
    parser.add_argument('--review_all', help='review all the new selections. Otherwise, you will only review the good selections',action ='store_true',default=False, dest='review_all')
    args = parser.parse_args(argv)

    if  not args.method in ['title','abstract','author, all']:
            print("Invalid method. Options are title, abstract, author, and all")
            exit()

    method = args.method

    #Set up training if that's what we're doing
    if args.train:

        #check to make sure we have a good training input
        if args.good_source is None:
            print("When training, you must specify one good source")
            exit()
        if not os.path.exists(args.good_source):
            print("Specified training input does not exist")
            exit()
        if not os.path.isfile(args.good_source):
            print("Specified training input is not a file")
            exit()
        if not os.path.splitext(args.good_source)[1] == '.bib' :
            print("Training input must be in bibtex format")
            exit()

        #load the existing knowledge
        nb,kw,knowledge = load_knowledge(args.knowledge)
        if args.overwrite_knowledge:
            nb=None
            kw=list()

        good_content = get_content([bibtex.BibTex(args.good_source)])

        bad_content = get_content([arxiv.ArXiv(cat) for cat in args.arXiv] +
                                  [bibtex.BibTex(bibfile) for bibfile in args.bibfiles] +
                                   [rss.JournalFeed(journal) for journal in args.journals])

        #train, and write out knowledge (naive_bayes class and keywords)
        nb, kw = train(good_content,bad_content,method,naive_bayes=nb, keywords=kw)
        pickle.dump(nb,open(knowledge+'/nb.p','w'))
        pickle.dump(kw,open(knowledge+'/kw.p','w'))

    #we are filtering new content through our existing knowledge
    else:

        #load the old knowledge
        nb,kw,knowledge = load_knowledge(args.knowledge)
        sources = [arxiv.ArXiv(cat) for cat in args.arXiv] + \
                  [ bibtex.BibTex (bibfile) for bibfile in args.bibfiles] + \
                  [rss.JournalFeed(journal) for journal in args.journals]

        new_content = get_content(sources)

        good_content = filter_content(new_content,method,nb,kw)

        if (args.output):
            to_markdown(good_content,args.output)
        else:
            print(good_content)

        if(args.feedback):
            human_class, reviewed_content = review_content(good_content,new_content,method,args.review_all)
            good_content = [entry for cat,entry in zip(human_class,reviewed_content) if cat=='good']
            bad_content = [entry for cat,entry in zip(human_class,reviewed_content) if cat=='bad']
            nb, kw = train(good_content,bad_content,method,naive_bayes=nb, keywords=kw)
            pickle.dump(nb,open(knowledge+'/nb.p','w'))
            pickle.dump(kw,open(knowledge+'/kw.p','w'))

if __name__=="__main__":
    main(sys.argv[1:])
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
##
# parallel.py: Tools for distributing computation.
##
# © 2014 Chris Ferrie (csferrie@gmail.com) and
#        Christopher E. Granade (cgranade@gmail.com)
#
# This file is a part of the Qinfer project.
# Licensed under the AGPL version 3.
##
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

## FEATURES ##################################################################

from __future__ import absolute_import
from __future__ import division # Ensures that a/b is always a float.

## EXPORTS ###################################################################

__all__ = ['DirectViewParallelizedModel']

## IMPORTS ###################################################################

import numpy as np
from qinfer.abstract_model import Model
from qinfer.derived_models import DerivedModel

import warnings

try:
    import ipyparallel as ipp
    interactive = ipp.interactive
except ImportError:
    try:
        import IPython.parallel as ipp
        interactive = ipp.interactive
    except (ImportError, AttributeError):
        import warnings
        warnings.warn(
            "Could not import IPython parallel. "
            "Parallelization support will be disabled."
        )
        ipp = None
        interactive = lambda fn: fn

## LOGGING ###################################################################

import logging
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

## CLASSES ###################################################################

class DirectViewParallelizedModel(DerivedModel):
    r"""
    Given an instance of a :class:`Model`, parallelizes execution of that model's
    likelihood by breaking the ``modelparams`` array into segments and
    executing a segment on each member of a :class:`~ipyparallel.DirectView`.

    This :class:`Model` assumes that it has ownership over the DirectView, such
    that no other processes will send tasks during the lifetime of the Model.

    :param qinfer.Model serial_model: Model to be parallelized. This
        model will be distributed to the engines in the direct view, such that
        the model must support pickling.
    :param ipyparallel.DirectView direct_view: Direct view onto the engines
        that will be used to parallelize evaluation of the model's likelihood
        function.
    :param bool purge_client: If ``True``, then this model will purge results
        and metadata from the IPython client whenever the model cache is cleared.
        This is useful for solving memory leaks caused by very large numbers of
        calls to ``likelihood``. By default, this is disabled, since enabling
        this option can cause data loss if the client is being sent other tasks
        during the operation of this model.
    :param int serial_threshold: Sets the number of model vectors below which
        the serial model is to be preferred. By default, this is set to ``10 *
        n_engines``, where ``n_engines`` is the number of engines exposed by
        ``direct_view``.
    """

    ## INITIALIZER ##

    def __init__(self, serial_model, direct_view, purge_client=False, serial_theshold=None):
        if ipp is None:
            raise RuntimeError(
                "This model requires IPython parallelization support, "
                "but an error was raised importing IPython.parallel."
            )

        self._dv = direct_view
        self._purge_client = purge_client
        self._serial_threshold = (
            10 * self.n_engines
            if serial_theshold is None else int(serial_theshold)
        )

        super(DirectViewParallelizedModel, self).__init__(serial_model)

    ## SPECIAL METHODS ##

    def __getstate__(self):
        # Since instances of this class will be pickled as they are passed to
        # remote engines, we need to be careful not to include _dv
        return {
            '_underlying_model': self._underlying_model,
            '_dv': None,
            '_call_count': self._call_count,
            '_sim_count': self._sim_count,
            '_serial_threshold': self._serial_threshold
        }

    ## PROPERTIES ##

    # Provide _serial_model as a back-compat.
    @property
    def _serial_model(self):
        warnings.warn("_serial_model is deprecated in favor of _underlying_model.",
            DeprecationWarning
        )
        return self._underlying_model
    @_serial_model.setter
    def _serial_model(self, value):
        warnings.warn("_serial_model is deprecated in favor of _underlying_model.",
            DeprecationWarning
        )
        self._underlying_model = value

    @property
    def n_engines(self):
        """
        The number of engines seen by the direct view owned by this parallelized
        model.

        :rtype: int
        """
        return len(self._dv) if self._dv is not None else 0

    ## METHODS ##

    def clear_cache(self):
        """
        Clears any cache associated with the serial model and the engines
        seen by the direct view.
        """
        self.underlying_model.clear_cache()
        try:
            logger.info('DirectView results has {} items. Clearing.'.format(
                len(self._dv.results)
            ))
            self._dv.purge_results('all')
            if self._purge_client:
                self._dv.client.purge_everything()
        except:
            pass

    def likelihood(self, outcomes, modelparams, expparams):
        # By calling the superclass implementation, we can consolidate
        # call counting there.
        """
        Returns the likelihood for the underlying (serial) model, distributing
        the model parameter array across the engines controlled by this
        parallelized model.
        """
        super(DirectViewParallelizedModel, self).likelihood(outcomes, modelparams, expparams)

        if modelparams.shape[1] <= self._serial_threshold:
            return self._serial_model.likelihood(outcomes, modelparams, expparams)

        if self._dv is None:
            raise RuntimeError(
                "No direct view provided; this may be because the instance was "
                "loaded from a pickle or NumPy saved array without providing a "
                "new direct view."
            )

        # If there's less models than some threshold, just use the serial model.
        # By default, we'll set that threshold to be the number of engines * 10.

        # Need to decorate with interactive to overcome namespace issues with
        # remote engines.
        @interactive
        def serial_likelihood(mps, sm, os, eps):
            return sm.likelihood(os, mps, eps)

        # TODO: check whether there's a better way to pass the extra parameters
        # that doesn't use so much memory.
        # The trick is that serial_likelihood will be pickled, so we need to be
        # careful about closures.
        L = self._dv.map_sync(
            serial_likelihood,
            np.array_split(modelparams, self.n_engines, axis=0),
            [self.underlying_model] * self.n_engines,
            [outcomes] * self.n_engines,
            [expparams] * self.n_engines,
        )

        return np.concatenate(L, axis=1)

<commit_msg>Fixed an index bug and fixed a typo while at it.<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
##
# parallel.py: Tools for distributing computation.
##
# © 2014 Chris Ferrie (csferrie@gmail.com) and
#        Christopher E. Granade (cgranade@gmail.com)
#
# This file is a part of the Qinfer project.
# Licensed under the AGPL version 3.
##
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

## FEATURES ##################################################################

from __future__ import absolute_import
from __future__ import division # Ensures that a/b is always a float.

## EXPORTS ###################################################################

__all__ = ['DirectViewParallelizedModel']

## IMPORTS ###################################################################

import numpy as np
from qinfer.abstract_model import Model
from qinfer.derived_models import DerivedModel

import warnings

try:
    import ipyparallel as ipp
    interactive = ipp.interactive
except ImportError:
    try:
        import IPython.parallel as ipp
        interactive = ipp.interactive
    except (ImportError, AttributeError):
        import warnings
        warnings.warn(
            "Could not import IPython parallel. "
            "Parallelization support will be disabled."
        )
        ipp = None
        interactive = lambda fn: fn

## LOGGING ###################################################################

import logging
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

## CLASSES ###################################################################

class DirectViewParallelizedModel(DerivedModel):
    r"""
    Given an instance of a :class:`Model`, parallelizes execution of that model's
    likelihood by breaking the ``modelparams`` array into segments and
    executing a segment on each member of a :class:`~ipyparallel.DirectView`.

    This :class:`Model` assumes that it has ownership over the DirectView, such
    that no other processes will send tasks during the lifetime of the Model.

    If you are having trouble pickling your model, consider switching to
    ``dill`` by calling ``dview.use_dill()``. This mode gives more support
    for closures.

    :param qinfer.Model serial_model: Model to be parallelized. This
        model will be distributed to the engines in the direct view, such that
        the model must support pickling.
    :param ipyparallel.DirectView direct_view: Direct view onto the engines
        that will be used to parallelize evaluation of the model's likelihood
        function.
    :param bool purge_client: If ``True``, then this model will purge results
        and metadata from the IPython client whenever the model cache is cleared.
        This is useful for solving memory leaks caused by very large numbers of
        calls to ``likelihood``. By default, this is disabled, since enabling
        this option can cause data loss if the client is being sent other tasks
        during the operation of this model.
    :param int serial_threshold: Sets the number of model vectors below which
        the serial model is to be preferred. By default, this is set to ``10 *
        n_engines``, where ``n_engines`` is the number of engines exposed by
        ``direct_view``.
    """

    ## INITIALIZER ##

    def __init__(self, serial_model, direct_view, purge_client=False, serial_threshold=None):
        if ipp is None:
            raise RuntimeError(
                "This model requires IPython parallelization support, "
                "but an error was raised importing IPython.parallel."
            )

        self._dv = direct_view
        self._purge_client = purge_client
        self._serial_threshold = (
            10 * self.n_engines
            if serial_threshold is None else int(serial_threshold)
        )

        super(DirectViewParallelizedModel, self).__init__(serial_model)

    ## SPECIAL METHODS ##

    def __getstate__(self):
        # Since instances of this class will be pickled as they are passed to
        # remote engines, we need to be careful not to include _dv
        return {
            '_underlying_model': self._underlying_model,
            '_dv': None,
            '_call_count': self._call_count,
            '_sim_count': self._sim_count,
            '_serial_threshold': self._serial_threshold
        }

    ## PROPERTIES ##

    # Provide _serial_model as a back-compat.
    @property
    def _serial_model(self):
        warnings.warn("_serial_model is deprecated in favor of _underlying_model.",
            DeprecationWarning
        )
        return self._underlying_model
    @_serial_model.setter
    def _serial_model(self, value):
        warnings.warn("_serial_model is deprecated in favor of _underlying_model.",
            DeprecationWarning
        )
        self._underlying_model = value

    @property
    def n_engines(self):
        """
        The number of engines seen by the direct view owned by this parallelized
        model.

        :rtype: int
        """
        return len(self._dv) if self._dv is not None else 0

    ## METHODS ##

    def clear_cache(self):
        """
        Clears any cache associated with the serial model and the engines
        seen by the direct view.
        """
        self.underlying_model.clear_cache()
        try:
            logger.info('DirectView results has {} items. Clearing.'.format(
                len(self._dv.results)
            ))
            self._dv.purge_results('all')
            if self._purge_client:
                self._dv.client.purge_everything()
        except:
            pass

    def likelihood(self, outcomes, modelparams, expparams):
        # By calling the superclass implementation, we can consolidate
        # call counting there.
        """
        Returns the likelihood for the underlying (serial) model, distributing
        the model parameter array across the engines controlled by this
        parallelized model.
        """
        super(DirectViewParallelizedModel, self).likelihood(outcomes, modelparams, expparams)

        if modelparams.shape[0] <= self._serial_threshold:
            return self._serial_model.likelihood(outcomes, modelparams, expparams)

        if self._dv is None:
            raise RuntimeError(
                "No direct view provided; this may be because the instance was "
                "loaded from a pickle or NumPy saved array without providing a "
                "new direct view."
            )

        # If there's less models than some threshold, just use the serial model.
        # By default, we'll set that threshold to be the number of engines * 10.

        # Need to decorate with interactive to overcome namespace issues with
        # remote engines.
        @interactive
        def serial_likelihood(mps, sm, os, eps):
            return sm.likelihood(os, mps, eps)

        # TODO: check whether there's a better way to pass the extra parameters
        # that doesn't use so much memory.
        # The trick is that serial_likelihood will be pickled, so we need to be
        # careful about closures.
        L = self._dv.map_sync(
            serial_likelihood,
            np.array_split(modelparams, self.n_engines, axis=0),
            [self.underlying_model] * self.n_engines,
            [outcomes] * self.n_engines,
            [expparams] * self.n_engines,
        )

        return np.concatenate(L, axis=1)

<|endoftext|>


<commit_before># -*- Mode: Python -*-

import charset
import regular
import automata
from pprint import pprint as pp
import sys
W = sys.stdout.write

# two approaches for generating a lexer:
#  1) generate tables and dispatch through them at run-time.
#  2) generate range-test code from the dfa.

# #2 should be faster, especially if we try a little harder.
# idea: feed a bunch of sample data through the lexer, collecting
#   probabilities at each transition, use this to reorder them.
#   [a branch predictor will do the same thing, is it worth it?]

# I think tables are better once the lexer hits a certain size.

class lexer:

    def __init__ (self, lexicon):
        # combine all the regexps into one, with each sub-expression wrapped with TOKEN
        tokens = regular.n_or ([ (regular.TOKEN, regular.parse (x), y) for (x, y) in lexicon ])
        #pp (tokens)
        c0 = automata.regexp_to_nfa_converter()
        nfa, start, end = c0.go (tokens)
        #pp (nfa)
        c1 = automata.nfa_dfa_converter (nfa, start, end)
        dfa, finals, actions = c1.go (c0.tokens)
        flat = c1.coalesce()
        #print 'dfa='
        #pp (flat)
        #print 'finals=', finals
        self.dfa = flat
        #print 'actions=', actions
        # XXX finals + actions == redundant
        self.finals = finals
        self.actions = {}
        for state, action in actions:
            if not self.actions.has_key (state):
                # give preference to tokens defined earlier in the lexicon
                self.actions[state] = action
        #print 'actions=', self.actions

    def sort_transitions (self, trans):
        # sort transitions in a way that generates the most efficient set of
        #   range tests.
        tests = [ (sym.as_ranges()[1], sym, ts) for (sym, ts) in trans ]
        tests.sort (lambda a,b: cmp (len(a[0]), len(b[0])))
        return [ (sym, ts) for (ranges, sym, ts) in tests ]

    def read (self, file):
        # this is useful for debugging your lexicon.  feed it data you'd
        #   like to lex.  when it works the way you like, then call
        #   gen_irken/c/etc() on it.
        s = file.read()
        state = 0
        current = []
        good = None
        slen = len(s)
        i = 0
        for ch in s:
            while 1:
                entry = self.dfa[state]
                ostate = state
                for sym, ts in entry:
                    if sym.has (ch):
                        state = ts
                        break
                print '%d %r => %d' % (ostate, ch, state)
                #W ('[%d|%d]' % (ord(ch),state,))
                action = self.actions.get (state, None)
                if not good and action:
                    # transition in
                    good = action
                    break
                elif good and not action:
                    # transition out
                    print good, repr(''.join (current))
                    current = []
                    state = 0
                    good = None
                else:
                    # update final state if one...
                    good = action
                    break
            current.append (ch)
            i += 1

    def generate (self):
        print     "    def step (self, ch):"
        for i in range (len (self.dfa)):
            if i == 0:
                print "        if self.state == %d:" % i
            else:
                print "        elif self.state == %d:" % i
            state_len = len (self.dfa[i])
            trans = self.sort_transitions (self.dfa[i])
            for j in range (state_len):
                sym, ts = trans[j]
                if state_len == 1:
                    print "            self.state = %d" % (ts,)
                elif j == state_len - 1:
                    # last range is always a catch-all
                    print "            else:"
                    print "                self.state = %d" % (ts,)
                else:
                    invert, ranges = sym.as_ranges()
                    test = []
                    for start, end in ranges:
                        if start == end:
                            test.append ("ch == %d" % (start,))
                        else:
                            test.append ("%d <= ch <= %d" % (start, end))
                    test = ' or '.join (test)
                    if invert:
                        test = 'not (%s)' % (test,)
                    if j == 0:
                        print "            if %s:" % (test,)
                    else:
                        print "            elif %s:" % (test,)
                    print "                self.state = %d" % (ts,)
            if self.actions.has_key (i):
                print "                # TOKEN %s" % (self.actions[i],)

    def gen_c (self):
        print "int step (int state, char ch)"
        print "{"
        for i in range (len (self.dfa)):
            if i == 0:
                print "  if (state == %d) {" % i
            else:
                print "  } else if (state == %d) {" % i
            state_len = len (self.dfa[i])
            trans = self.sort_transitions (self.dfa[i])
            for j in range (state_len):
                sym, ts = trans[j]
                if state_len == 1:
                    print "    state = %d;" % (ts,)
                elif j == state_len - 1:
                    # last range is always a catch-all
                    print "    } else {"
                    print "      state = %d;" % (ts,)
                    print "    }"
                else:
                    invert, ranges = sym.as_ranges()
                    test = []
                    for start, end in ranges:
                        if start == end:
                            test.append ("(ch == %d)" % (start,))
                        else:
                            test.append ("(ch >= %d && ch <= %d)" % (start, end))
                    test = ' || '.join (test)
                    if invert:
                        test = '!(%s)' % (test,)
                    if j == 0:
                        print "    if (%s) {" % (test,)
                    else:
                        print "    } else if (%s) {" % (test,)
                    print "      state = %d;" % (ts,)
        print "  }"
        print "  return state;"
        print "}"

    def find_sink (self):
        # is the sink is always the second state?
        for i in range (len (self.dfa)):
            entry = self.dfa[i]
            if len(entry) == 1:
                [(sym, ts)] = entry
                if sym == charset.DOT and ts == i:
                    return i
        else:
            raise ValueError

    def gen_irken (self, file):
        # generate a table-based dfa, using strings
        assert (len (self.dfa) < 256)
        W = file.write
        W (";; generated by lexer.py\n")
        tables = []
        for i in range (len (self.dfa)):
            table = [1] * 256
            for sym, ts in self.dfa[i]:
                for i in range (256):
                    if sym[i]:
                        table[i] = ts
            table = ''.join ([chr(x) for x in table])
            tables.append (table)
        unique = {}
        for t in tables:
            if not unique.has_key (t):
                unique[t] = len (unique)
        indirect = []
        for t in tables:
            indirect.append (unique[t])
        W ("(define dfa \n")
        W ("  (let* (\n")
        items = unique.items()
        items.sort (lambda a,b: cmp (a[1], b[1]))
        for table, index in items:
            W ("         (t%d \"%s\")\n" % (index, ''.join (["\\x%02x" % (ord (x),) for x in table])))
        W ("        )\n")
        W ("   #(%s)))\n" % (" ".join (["t%d" % (x,) for x in indirect])))

        # find the sink state
        sink = self.find_sink()

        W ("(define finals\n")
        W ("  #(\n")
        for i in range (len (self.dfa)):
            f = self.actions.get (i, None)
            if f:
                W ("  '%s\n" % f)
            else:
                W ("  'not-final\n")
        W ("  ))\n")
        W ("(define (step ch state)\n")
        W ("  (char->ascii (string-ref dfa[state] (char->ascii ch))))\n")
        return tables

keywords = '|'.join ('and is in not if then else elif yield while for try def class'.split())
augassign = ('(' + '|'.join (r'\+ - \* / % & \| \^ << >> \*\* //'.split()) + ')=')

python_lexicon = [
    (r'[0-9]+',                 'number'),
    (r'[A-Za-z_][A-Za-z0-9_]*', 'ident'),
    (r'"([^\\\n"]|(\\.))*"',    'string1'),
    (r"'([^\\\n']|(\\.))*'",    'string2'),
    (r'[ \t]*#[^\n]*\n',        'comment'),
    (r'[-+]',                   'addop'),
    (r'([*/%]|//)',             'mulop'),
    (r'<|>|<=|>=|==',           'compare'),
    (r'<<|>>',                  'shift'),
    (r'\*\*',                   'power'),
    (r'[ \t]+',                 'whitespace'),
    (r'[\n]',                   'newline'),
    (r'\.',                     'getattr'),
    (r'=',                      'assign'),
    (r'\(',                     'lparen'),
    (r'\)',                     'rparen'),
    (r',',                      'comma'),
    (r':',                      'colon'),
    (r';',                      'semicolon'),
    (r'{',                      'lbracket'),
    (r'}',                      'rbracket'),
    (r'\[',                     'lbrace'),
    (r'\]',                     'rbrace'),
    (r'&',                      'bitand'),
    (r'\|',                     'bitor'),
    (r'^',                      'bitxor'),
    (r'~',                      'bitnot'),
    (keywords,                  'keyword'),
    (augassign,                 'augassign'),
    ]

if __name__ == '__main__':
    m = lexer (python_lexicon)
    #m.gen_scheme_code()
    #m.gen_scheme_table()
    #m.gen_irken (open ("lexstep.scm", 'wb'))
    m.read (open ("../nodes.py"))
<commit_msg>gen lexstep.scm<commit_after># -*- Mode: Python -*-

import charset
import regular
import automata
from pprint import pprint as pp
import sys
W = sys.stdout.write

# two approaches for generating a lexer:
#  1) generate tables and dispatch through them at run-time.
#  2) generate range-test code from the dfa.

# #2 should be faster, especially if we try a little harder.
# idea: feed a bunch of sample data through the lexer, collecting
#   probabilities at each transition, use this to reorder them.
#   [a branch predictor will do the same thing, is it worth it?]

# I think tables are better once the lexer hits a certain size.

class lexer:

    def __init__ (self, lexicon):
        # combine all the regexps into one, with each sub-expression wrapped with TOKEN
        tokens = regular.n_or ([ (regular.TOKEN, regular.parse (x), y) for (x, y) in lexicon ])
        #pp (tokens)
        c0 = automata.regexp_to_nfa_converter()
        nfa, start, end = c0.go (tokens)
        #pp (nfa)
        c1 = automata.nfa_dfa_converter (nfa, start, end)
        dfa, finals, actions = c1.go (c0.tokens)
        flat = c1.coalesce()
        #print 'dfa='
        #pp (flat)
        #print 'finals=', finals
        self.dfa = flat
        #print 'actions=', actions
        # XXX finals + actions == redundant
        self.finals = finals
        self.actions = {}
        for state, action in actions:
            if not self.actions.has_key (state):
                # give preference to tokens defined earlier in the lexicon
                self.actions[state] = action
        #print 'actions=', self.actions

    def sort_transitions (self, trans):
        # sort transitions in a way that generates the most efficient set of
        #   range tests.
        tests = [ (sym.as_ranges()[1], sym, ts) for (sym, ts) in trans ]
        tests.sort (lambda a,b: cmp (len(a[0]), len(b[0])))
        return [ (sym, ts) for (ranges, sym, ts) in tests ]

    def read (self, file):
        # this is useful for debugging your lexicon.  feed it data you'd
        #   like to lex.  when it works the way you like, then call
        #   gen_irken/c/etc() on it.
        s = file.read()
        state = 0
        current = []
        good = None
        slen = len(s)
        i = 0
        for ch in s:
            while 1:
                entry = self.dfa[state]
                ostate = state
                for sym, ts in entry:
                    if sym.has (ch):
                        state = ts
                        break
                print '%d %r => %d' % (ostate, ch, state)
                #W ('[%d|%d]' % (ord(ch),state,))
                action = self.actions.get (state, None)
                if not good and action:
                    # transition in
                    good = action
                    break
                elif good and not action:
                    # transition out
                    print good, repr(''.join (current))
                    current = []
                    state = 0
                    good = None
                else:
                    # update final state if one...
                    good = action
                    break
            current.append (ch)
            i += 1

    def generate (self):
        print     "    def step (self, ch):"
        for i in range (len (self.dfa)):
            if i == 0:
                print "        if self.state == %d:" % i
            else:
                print "        elif self.state == %d:" % i
            state_len = len (self.dfa[i])
            trans = self.sort_transitions (self.dfa[i])
            for j in range (state_len):
                sym, ts = trans[j]
                if state_len == 1:
                    print "            self.state = %d" % (ts,)
                elif j == state_len - 1:
                    # last range is always a catch-all
                    print "            else:"
                    print "                self.state = %d" % (ts,)
                else:
                    invert, ranges = sym.as_ranges()
                    test = []
                    for start, end in ranges:
                        if start == end:
                            test.append ("ch == %d" % (start,))
                        else:
                            test.append ("%d <= ch <= %d" % (start, end))
                    test = ' or '.join (test)
                    if invert:
                        test = 'not (%s)' % (test,)
                    if j == 0:
                        print "            if %s:" % (test,)
                    else:
                        print "            elif %s:" % (test,)
                    print "                self.state = %d" % (ts,)
            if self.actions.has_key (i):
                print "                # TOKEN %s" % (self.actions[i],)

    def gen_c (self):
        print "int step (int state, char ch)"
        print "{"
        for i in range (len (self.dfa)):
            if i == 0:
                print "  if (state == %d) {" % i
            else:
                print "  } else if (state == %d) {" % i
            state_len = len (self.dfa[i])
            trans = self.sort_transitions (self.dfa[i])
            for j in range (state_len):
                sym, ts = trans[j]
                if state_len == 1:
                    print "    state = %d;" % (ts,)
                elif j == state_len - 1:
                    # last range is always a catch-all
                    print "    } else {"
                    print "      state = %d;" % (ts,)
                    print "    }"
                else:
                    invert, ranges = sym.as_ranges()
                    test = []
                    for start, end in ranges:
                        if start == end:
                            test.append ("(ch == %d)" % (start,))
                        else:
                            test.append ("(ch >= %d && ch <= %d)" % (start, end))
                    test = ' || '.join (test)
                    if invert:
                        test = '!(%s)' % (test,)
                    if j == 0:
                        print "    if (%s) {" % (test,)
                    else:
                        print "    } else if (%s) {" % (test,)
                    print "      state = %d;" % (ts,)
        print "  }"
        print "  return state;"
        print "}"

    def find_sink (self):
        # is the sink is always the second state?
        for i in range (len (self.dfa)):
            entry = self.dfa[i]
            if len(entry) == 1:
                [(sym, ts)] = entry
                if sym == charset.DOT and ts == i:
                    return i
        else:
            raise ValueError

    def gen_irken (self, file):
        # generate a table-based dfa, using strings
        assert (len (self.dfa) < 256)
        W = file.write
        W (";; generated by lexer.py\n")
        tables = []
        for i in range (len (self.dfa)):
            table = [1] * 256
            for sym, ts in self.dfa[i]:
                for i in range (256):
                    if sym[i]:
                        table[i] = ts
            table = ''.join ([chr(x) for x in table])
            tables.append (table)
        unique = {}
        for t in tables:
            if not unique.has_key (t):
                unique[t] = len (unique)
        indirect = []
        for t in tables:
            indirect.append (unique[t])
        W ("(define dfa \n")
        W ("  (let* (\n")
        items = unique.items()
        items.sort (lambda a,b: cmp (a[1], b[1]))
        for table, index in items:
            W ("         (t%d \"%s\")\n" % (index, ''.join (["\\x%02x" % (ord (x),) for x in table])))
        W ("        )\n")
        W ("   #(%s)))\n" % (" ".join (["t%d" % (x,) for x in indirect])))

        # find the sink state
        sink = self.find_sink()

        W ("(define finals\n")
        W ("  #(\n")
        for i in range (len (self.dfa)):
            f = self.actions.get (i, None)
            if f:
                W ("  '%s\n" % f)
            else:
                W ("  'not-final\n")
        W ("  ))\n")
        W ("(define (step ch state)\n")
        W ("  (char->ascii (string-ref dfa[state] (char->ascii ch))))\n")
        return tables

keywords = '|'.join ('and is in not if then else elif yield while for try def class'.split())
augassign = ('(' + '|'.join (r'\+ - \* / % & \| \^ << >> \*\* //'.split()) + ')=')

python_lexicon = [
    (r'[0-9]+',                 'number'),
    (r'[A-Za-z_][A-Za-z0-9_]*', 'ident'),
    (r'"([^\\\n"]|(\\.))*"',    'string1'),
    (r"'([^\\\n']|(\\.))*'",    'string2'),
    (r'[ \t]*#[^\n]*\n',        'comment'),
    (r'[-+]',                   'addop'),
    (r'([*/%]|//)',             'mulop'),
    (r'<|>|<=|>=|==',           'compare'),
    (r'<<|>>',                  'shift'),
    (r'\*\*',                   'power'),
    (r'[ \t]+',                 'whitespace'),
    (r'[\n]',                   'newline'),
    (r'\.',                     'getattr'),
    (r'=',                      'assign'),
    (r'\(',                     'lparen'),
    (r'\)',                     'rparen'),
    (r',',                      'comma'),
    (r':',                      'colon'),
    (r';',                      'semicolon'),
    (r'{',                      'lbracket'),
    (r'}',                      'rbracket'),
    (r'\[',                     'lbrace'),
    (r'\]',                     'rbrace'),
    (r'&',                      'bitand'),
    (r'\|',                     'bitor'),
    (r'^',                      'bitxor'),
    (r'~',                      'bitnot'),
    (keywords,                  'keyword'),
    (augassign,                 'augassign'),
    ]

if __name__ == '__main__':
    m = lexer (python_lexicon)
    #m.gen_scheme_code()
    #m.gen_scheme_table()
    m.gen_irken (open ("lexstep.scm", 'wb'))
    #m.read (open ("../nodes.py"))
<|endoftext|>


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

__author__ = "@tsmall888"
__license__ = "MIT"

from gevent import monkey
monkey.patch_all(thread=False)

import re
import json
import censys
from censys import *
import shodan
import argparse
from netaddr import *
from gevent import spawn, joinall
from gevent.queue import JoinableQueue
from collections import OrderedDict
from termcolor import colored

import config

class PassiveScan(object):
	def __init__(self, args):
		self.__ip = args.ip
		self.__threads = args.threads
		self.__json = args.json
		self.__errors = args.errors
		self.__ports = args.ports
		self.__apis = []
		if args.shodan: self.__apis.append(shodan.Shodan(config.SHODAN_API_KEY))
		if args.censys: self.__apis.append(censys.ipv4.CensysIPv4(config.CENSYS_API_ID, config.CENSYS_SECRET))

	def start(self):
		if not self.__threads:
			self.__threads = len(IPNetwork(self.__ip)) if len(IPNetwork(self.__ip)) <= 10 else 10
		if len(IPNetwork(self.__ip)) < int(self.__threads):
			print "Please decrease number of threads to number of hosts <= %s" % len(IPNetwork(self.__ip))
			exit()

		queue = JoinableQueue()
		[queue.put(str(ip)) for ip in IPNetwork(self.__ip)]

		workers = [spawn(self.get_ip_info, queue, self.__apis) for t in range(int(self.__threads))]

		queue.join()

	def get_ip_info(self, queue, apis):
		while not queue.empty():
			data = OrderedDict()
			data["ip"] = queue.get()

			for api in apis:
				if api.__class__.__name__ == "Shodan":
					data["shodan"] = self.search_shodan(data["ip"], api)
					if self.__errors == False:
						if "error" in data["shodan"]:
							data.pop("shodan", None)
				elif api.__class__.__name__ == "CensysIPv4":
					data["censys"] = self.search_censys(data["ip"], api)
					if self.__errors == False:
						if "error" in data["censys"]:
							data.pop("censys", None)

			if self.__errors == False:
				if not "shodan" in data and not "censys" in data:
					data = None

			if data:
				self.output(data)

			queue.task_done()

	def search_shodan(self, ip, api):
		data = OrderedDict()

		try:
			result = api.host(ip)
		except shodan.APIError, e:
			data["error"] = str(e)
		else:
			if isinstance(result, dict):
				if "ip_str" in result:
					if "last_update" in result:
						data["last_update"] = result["last_update"]
					if "asn" in result:
						data["asn"] = int(re.findall("\d+", result["asn"])[0])
					if "isp" in result:
						data["isp"] = result["isp"]
					if "org" in result:
						data["org"] = result["org"]
					if "country_code" in result:
						data["country_code"] = result["country_code"]
					if "hostnames" in result and result["hostnames"]:
						data["hostnames"] = result["hostnames"]
					if "ports" in result:
						data["ports"] = result["ports"]
		finally:
			return data

	def search_censys(self, ip, api):
		data = OrderedDict()

		try:
			result = api.view(ip)
		except censys.base.CensysException, e:
			data["error"] = str(e)
		else:
			if isinstance(result, dict):
				if "ip" in result:
					if "updated_at" in result:
						data["last_update"] = result["updated_at"]
					if "autonomous_system" in result:
						if "asn" in result["autonomous_system"]:
							data["asn"] = result["autonomous_system"]["asn"]
						if "description" in result["autonomous_system"]:
							data["as_name"] = result["autonomous_system"]["name"]
					if "location" in result:
						if "country_code" in result["location"]:
							data["country_code"] = result["location"]["country_code"]
					if "protocols" in result:
						data["ports"] = [int(port.split("/")[0]) for port in result["protocols"]]
		finally:
			return data

	def output(self, data):
		if self.__json:
			print json.dumps(data, indent=4, sort_keys=False)

		else:
			for key, value in data.iteritems():
				if key is "ip":
					print colored(key, 'yellow'), colored(value, 'green')

				if key in ['shodan', 'censys']:
					print "  %s:" % (colored(key, 'yellow'))

				if isinstance(value, dict):
					for key, value in value.iteritems():
						if key == "ports":
							if self.__ports:
								if key == "ports":
									print "    %s: %s" % (colored(key, 'yellow'), " ".join([colored(port, 'red') if port in self.__ports else colored(port, 'green') for port in value]))
								else:
									print "    %s: %s" % (colored(key, 'yellow'), colored(value, 'green'))
						else:
							print "    %s: %s" % (colored(key, 'yellow'), colored(value, 'green'))

def main():
	help = "Passive scan"
	parser = argparse.ArgumentParser(description=help, epilog="Author: %s" % __author__)
	parser.add_argument("-ip", help="target ip or net", required=False)
	parser.add_argument("--shodan", "-s", help="use shodan", required=False, action='store_true')
	parser.add_argument("--censys", "-c", help="use censys", required=False, action='store_true')
	parser.add_argument("--threads", "-t", help="number of threads", required=False, type=int)
	parser.add_argument("--json", help='json output', required=False, action='store_true')
	parser.add_argument("--errors", "-e", help='show APIs errors', required=False, action='store_true')
	parser.add_argument("--ports", "-p", type=int, nargs='*', help='specify intersting ports', required=False)
	args = parser.parse_args()

	PassiveScan(args).start()

if __name__ == '__main__':
    main()
<commit_msg>moved to python3<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__author__ = "@t33m"
__license__ = "MIT"

from gevent import monkey
monkey.patch_all(thread=False)

import re
import json
import censys
from censys import *
import shodan
import argparse
from netaddr import *
from gevent import spawn, joinall
from gevent.queue import JoinableQueue
from collections import OrderedDict
from termcolor import colored

import config

class PassiveScan(object):
	def __init__(self, args):
		self.__ip = args.ip
		self.__threads = args.threads
		self.__json = args.json
		self.__errors = args.errors
		self.__ports = args.ports
		self.__apis = []
		if args.shodan: self.__apis.append(shodan.Shodan(config.SHODAN_API_KEY))
		if args.censys: self.__apis.append(censys.ipv4.CensysIPv4(config.CENSYS_API_ID, config.CENSYS_SECRET))

	def start(self):
		if not self.__threads:
			self.__threads = len(IPNetwork(self.__ip)) if len(IPNetwork(self.__ip)) <= 10 else 10
		if len(IPNetwork(self.__ip)) < int(self.__threads):
			print("Please decrease number of threads to number of hosts <= %s" % len(IPNetwork(self.__ip)))
			exit()

		queue = JoinableQueue()
		[queue.put(str(ip)) for ip in IPNetwork(self.__ip)]

		workers = [spawn(self.get_ip_info, queue, self.__apis) for t in range(int(self.__threads))]

		queue.join()

	def get_ip_info(self, queue, apis):
		while not queue.empty():
			data = OrderedDict()
			data["ip"] = queue.get()

			for api in apis:
				if api.__class__.__name__ == "Shodan":
					data["shodan"] = self.search_shodan(data["ip"], api)
					if self.__errors == False:
						if "error" in data["shodan"]:
							data.pop("shodan", None)
				elif api.__class__.__name__ == "CensysIPv4":
					data["censys"] = self.search_censys(data["ip"], api)
					if self.__errors == False:
						if "error" in data["censys"]:
							data.pop("censys", None)

			if self.__errors == False:
				if not "shodan" in data and not "censys" in data:
					data = None

			if data:
				self.output(data)

			queue.task_done()

	def search_shodan(self, ip, api):
		data = OrderedDict()

		try:
			result = api.host(ip)
		except shodan.APIError as e:
			data["error"] = str(e)
		else:
			if isinstance(result, dict):
				if "ip_str" in result:
					if "last_update" in result:
						data["last_update"] = result["last_update"]
					if "asn" in result:
						data["asn"] = int(re.findall("\d+", result["asn"])[0])
					if "isp" in result:
						data["isp"] = result["isp"]
					if "org" in result:
						data["org"] = result["org"]
					if "country_code" in result:
						data["country_code"] = result["country_code"]
					if "hostnames" in result and result["hostnames"]:
						data["hostnames"] = result["hostnames"]
					if "ports" in result:
						data["ports"] = result["ports"]
		finally:
			return data

	def search_censys(self, ip, api):
		data = OrderedDict()

		try:
			result = api.view(ip)
		except censys.base.CensysException as e:
			data["error"] = str(e)
		else:
			if isinstance(result, dict):
				if "ip" in result:
					if "updated_at" in result:
						data["last_update"] = result["updated_at"]
					if "autonomous_system" in result:
						if "asn" in result["autonomous_system"]:
							data["asn"] = result["autonomous_system"]["asn"]
						if "description" in result["autonomous_system"]:
							data["as_name"] = result["autonomous_system"]["name"]
					if "location" in result:
						if "country_code" in result["location"]:
							data["country_code"] = result["location"]["country_code"]
					if "protocols" in result:
						data["ports"] = [int(port.split("/")[0]) for port in result["protocols"]]
		finally:
			return data

	def output(self, data):
		if self.__json:
			print(json.dumps(data, indent=4, sort_keys=False))

		else:
			for key, value in data.iteritems():
				if key is "ip":
					print(colored(key, 'yellow'), colored(value, 'green'))

				if key in ['shodan', 'censys']:
					print("  %s:" % (colored(key, 'yellow')))

				if isinstance(value, dict):
					for key, value in value.iteritems():
						if key == "ports":
							if self.__ports:
								if key == "ports":
									print("    %s: %s" % (colored(key, 'yellow'), " ".join([colored(port, 'red') if port in self.__ports else colored(port, 'green') for port in value])))
								else:
									print("    %s: %s" % (colored(key, 'yellow'), colored(value, 'green')))
						else:
							print("    %s: %s" % (colored(key, 'yellow'), colored(value, 'green')))

def main():
	help = "Passive scan"
	parser = argparse.ArgumentParser(description=help, epilog="Author: %s" % __author__)
	parser.add_argument("-ip", help="target ip or net", required=False)
	parser.add_argument("--shodan", "-s", help="use shodan", required=False, action='store_true')
	parser.add_argument("--censys", "-c", help="use censys", required=False, action='store_true')
	parser.add_argument("--threads", "-t", help="number of threads", required=False, type=int)
	parser.add_argument("--json", help='json output', required=False, action='store_true')
	parser.add_argument("--errors", "-e", help='show APIs errors', required=False, action='store_true')
	parser.add_argument("--ports", "-p", type=int, nargs='*', help='specify intersting ports', required=False)
	args = parser.parse_args()

	PassiveScan(args).start()

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


<commit_before># coding: utf-8

from dtest import Tester
from tools import since, no_vnodes
from assertions import assert_unavailable
from cassandra import ConsistencyLevel, WriteTimeout
from cassandra.query import SimpleStatement

import time
from threading import Thread
from ccmlib.cluster import Cluster

@since('2.0.6')
class TestPaxos(Tester):

    def prepare(self, ordered=False, create_keyspace=True, use_cache=False, nodes=1, rf=1):
        cluster = self.cluster

        if (ordered):
            cluster.set_partitioner("org.apache.cassandra.dht.ByteOrderedPartitioner")

        if (use_cache):
            cluster.set_configuration_options(values={ 'row_cache_size_in_mb' : 100 })

        cluster.populate(nodes).start()
        node1 = cluster.nodelist()[0]
        time.sleep(0.2)

        cursor = self.patient_cql_connection(node1, version="3.0.0")
        if create_keyspace:
            self.create_ks(cursor, 'ks', rf)
        return cursor

    def replica_availability_test(self):
        #See CASSANDRA-8640
        session = self.prepare(nodes=3, rf=3)
        session.execute("CREATE TABLE test (k int PRIMARY KEY, v int)")
        session.execute("INSERT INTO test (k, v) VALUES (0, 0) IF NOT EXISTS")

        self.cluster.nodelist()[2].stop()
        session.execute("INSERT INTO test (k, v) VALUES (1, 1) IF NOT EXISTS")

        self.cluster.nodelist()[1].stop()
        assert_unavailable(session.execute, "INSERT INTO test (k, v) VALUES (2, 2) IF NOT EXISTS")

        self.cluster.nodelist()[1].start(wait_for_binary_proto=True)
        session.execute("INSERT INTO test (k, v) VALUES (3, 3) IF NOT EXISTS")

        self.cluster.nodelist()[2].start(wait_for_binary_proto=True)
        session.execute("INSERT INTO test (k, v) VALUES (4, 4) IF NOT EXISTS")

    @no_vnodes()
    def cluster_availability_test(self):
        #Warning, a change in partitioner or a change in CCM token allocation
        #may require the partition keys of these inserts to be changed.
        #This must not use vnodes as it relies on assumed token values.

        session = self.prepare(nodes=3)
        session.execute("CREATE TABLE test (k int PRIMARY KEY, v int)")
        session.execute("INSERT INTO test (k, v) VALUES (0, 0) IF NOT EXISTS")

        self.cluster.nodelist()[2].stop()
        session.execute("INSERT INTO test (k, v) VALUES (1, 1) IF NOT EXISTS")

        self.cluster.nodelist()[1].stop()
        session.execute("INSERT INTO test (k, v) VALUES (3, 2) IF NOT EXISTS")

        self.cluster.nodelist()[1].start()
        session.execute("INSERT INTO test (k, v) VALUES (5, 5) IF NOT EXISTS")

        self.cluster.nodelist()[2].start()
        session.execute("INSERT INTO test (k, v) VALUES (6, 6) IF NOT EXISTS")

    def contention_test_multi_iterations(self):
        self._contention_test(8, 100)

    ##Warning, this test will require you to raise the open
    ##file limit on OSX. Use 'ulimit -n 1000'
    def contention_test_many_threds(self):
        self._contention_test(300, 1)

    def _contention_test(self, threads, iterations):
        """ Test threads repeatedly contending on the same row """

        verbose = False

        cursor = self.prepare(nodes=3)
        cursor.execute("CREATE TABLE test (k int, v int static, id int, PRIMARY KEY (k, id))")
        cursor.execute("INSERT INTO test(k, v) VALUES (0, 0)");

        class Worker(Thread):
            def __init__(self, wid, cursor, iterations, query):
                Thread.__init__(self)
                self.wid = wid
                self.iterations = iterations
                self.query = query
                self.cursor = cursor
                self.errors = 0
                self.retries = 0

            def run(self):
                global worker_done
                i = 0
                prev = 0
                while i < self.iterations:
                    done = False
                    while not done:
                        try:
                            res = self.cursor.execute(self.query, (prev+1, prev, self.wid ))
                            if verbose:
                                print "[%3d] CAS %3d -> %3d (res: %s)" % (self.wid, prev, prev+1, str(res))
                            if res[0][0] is True:
                                done = True
                                prev = prev + 1
                            else:
                                self.retries = self.retries + 1
                                # There is 2 conditions, so 2 reasons to fail: if we failed because the row with our
                                # worker ID already exists, it means we timeout earlier but our update did went in,
                                # so do consider this as a success
                                prev = res[0][3]
                                if res[0][2] is not None:
                                    if verbose:
                                        print "[%3d] Update was inserted on previous try (res = %s)" % (self.wid, str(res))
                                    done = True
                        except WriteTimeout as e:
                            if verbose:
                                print "[%3d] TIMEOUT (%s)" % (self.wid, str(e))
                            # This means a timeout: just retry, if it happens that our update was indeed persisted,
                            # we'll figure it out on the next run.
                            self.retries = self.retries + 1
                        except Exception as e:
                            if verbose:
                                print "[%3d] ERROR: %s" % (self.wid, str(e))
                            self.errors = self.errors + 1
                            done = True
                    i = i + 1
                    # Clean up for next iteration
                    while True:
                        try:
                            self.cursor.execute("DELETE FROM test WHERE k = 0 AND id = %d IF EXISTS" % self.wid)
                            break;
                        except WriteTimeout as e:
                            pass

        nodes = self.cluster.nodelist()
        workers = []

        c = self.patient_cql_connection(nodes[0], version="3.0.0", keyspace='ks')
        q = c.prepare("""
                BEGIN BATCH
                   UPDATE test SET v = ? WHERE k = 0 IF v = ?;
                   INSERT INTO test (k, id) VALUES (0, ?) IF NOT EXISTS;
                APPLY BATCH
            """)

        for n in range(0, threads):
            workers.append(Worker(n, c, iterations, q))

        start = time.time()

        for w in workers:
            w.start()

        for w in workers:
            w.join()

        if verbose:
            runtime = time.time() - start
            print "runtime:", runtime

        query = SimpleStatement("SELECT v FROM test WHERE k = 0", consistency_level=ConsistencyLevel.ALL)
        rows = cursor.execute(query)
        value = rows[0][0]

        errors = 0
        retries = 0
        for w in workers:
            errors = errors + w.errors
            retries = retries + w.retries

        assert (value == threads * iterations) and (errors == 0), "value=%d, errors=%d, retries=%d" % (value, errors, retries)

<commit_msg>Wait until cluster thinks node are up before querying<commit_after># coding: utf-8

from dtest import Tester
from tools import since, no_vnodes
from assertions import assert_unavailable
from cassandra import ConsistencyLevel, WriteTimeout
from cassandra.query import SimpleStatement

import time
from threading import Thread
from ccmlib.cluster import Cluster

@since('2.0.6')
class TestPaxos(Tester):

    def prepare(self, ordered=False, create_keyspace=True, use_cache=False, nodes=1, rf=1):
        cluster = self.cluster

        if (ordered):
            cluster.set_partitioner("org.apache.cassandra.dht.ByteOrderedPartitioner")

        if (use_cache):
            cluster.set_configuration_options(values={ 'row_cache_size_in_mb' : 100 })

        cluster.populate(nodes).start()
        node1 = cluster.nodelist()[0]
        time.sleep(0.2)

        cursor = self.patient_cql_connection(node1, version="3.0.0")
        if create_keyspace:
            self.create_ks(cursor, 'ks', rf)
        return cursor

    def replica_availability_test(self):
        #See CASSANDRA-8640
        session = self.prepare(nodes=3, rf=3)
        session.execute("CREATE TABLE test (k int PRIMARY KEY, v int)")
        session.execute("INSERT INTO test (k, v) VALUES (0, 0) IF NOT EXISTS")

        self.cluster.nodelist()[2].stop()
        session.execute("INSERT INTO test (k, v) VALUES (1, 1) IF NOT EXISTS")

        self.cluster.nodelist()[1].stop()
        assert_unavailable(session.execute, "INSERT INTO test (k, v) VALUES (2, 2) IF NOT EXISTS")

        self.cluster.nodelist()[1].start(wait_for_binary_proto=True, wait_other_notice=True)
        session.execute("INSERT INTO test (k, v) VALUES (3, 3) IF NOT EXISTS")

        self.cluster.nodelist()[2].start(wait_for_binary_proto=True)
        session.execute("INSERT INTO test (k, v) VALUES (4, 4) IF NOT EXISTS")

    @no_vnodes()
    def cluster_availability_test(self):
        #Warning, a change in partitioner or a change in CCM token allocation
        #may require the partition keys of these inserts to be changed.
        #This must not use vnodes as it relies on assumed token values.

        session = self.prepare(nodes=3)
        session.execute("CREATE TABLE test (k int PRIMARY KEY, v int)")
        session.execute("INSERT INTO test (k, v) VALUES (0, 0) IF NOT EXISTS")

        self.cluster.nodelist()[2].stop()
        session.execute("INSERT INTO test (k, v) VALUES (1, 1) IF NOT EXISTS")

        self.cluster.nodelist()[1].stop()
        session.execute("INSERT INTO test (k, v) VALUES (3, 2) IF NOT EXISTS")

        self.cluster.nodelist()[1].start()
        session.execute("INSERT INTO test (k, v) VALUES (5, 5) IF NOT EXISTS")

        self.cluster.nodelist()[2].start()
        session.execute("INSERT INTO test (k, v) VALUES (6, 6) IF NOT EXISTS")

    def contention_test_multi_iterations(self):
        self._contention_test(8, 100)

    ##Warning, this test will require you to raise the open
    ##file limit on OSX. Use 'ulimit -n 1000'
    def contention_test_many_threds(self):
        self._contention_test(300, 1)

    def _contention_test(self, threads, iterations):
        """ Test threads repeatedly contending on the same row """

        verbose = False

        cursor = self.prepare(nodes=3)
        cursor.execute("CREATE TABLE test (k int, v int static, id int, PRIMARY KEY (k, id))")
        cursor.execute("INSERT INTO test(k, v) VALUES (0, 0)");

        class Worker(Thread):
            def __init__(self, wid, cursor, iterations, query):
                Thread.__init__(self)
                self.wid = wid
                self.iterations = iterations
                self.query = query
                self.cursor = cursor
                self.errors = 0
                self.retries = 0

            def run(self):
                global worker_done
                i = 0
                prev = 0
                while i < self.iterations:
                    done = False
                    while not done:
                        try:
                            res = self.cursor.execute(self.query, (prev+1, prev, self.wid ))
                            if verbose:
                                print "[%3d] CAS %3d -> %3d (res: %s)" % (self.wid, prev, prev+1, str(res))
                            if res[0][0] is True:
                                done = True
                                prev = prev + 1
                            else:
                                self.retries = self.retries + 1
                                # There is 2 conditions, so 2 reasons to fail: if we failed because the row with our
                                # worker ID already exists, it means we timeout earlier but our update did went in,
                                # so do consider this as a success
                                prev = res[0][3]
                                if res[0][2] is not None:
                                    if verbose:
                                        print "[%3d] Update was inserted on previous try (res = %s)" % (self.wid, str(res))
                                    done = True
                        except WriteTimeout as e:
                            if verbose:
                                print "[%3d] TIMEOUT (%s)" % (self.wid, str(e))
                            # This means a timeout: just retry, if it happens that our update was indeed persisted,
                            # we'll figure it out on the next run.
                            self.retries = self.retries + 1
                        except Exception as e:
                            if verbose:
                                print "[%3d] ERROR: %s" % (self.wid, str(e))
                            self.errors = self.errors + 1
                            done = True
                    i = i + 1
                    # Clean up for next iteration
                    while True:
                        try:
                            self.cursor.execute("DELETE FROM test WHERE k = 0 AND id = %d IF EXISTS" % self.wid)
                            break;
                        except WriteTimeout as e:
                            pass

        nodes = self.cluster.nodelist()
        workers = []

        c = self.patient_cql_connection(nodes[0], version="3.0.0", keyspace='ks')
        q = c.prepare("""
                BEGIN BATCH
                   UPDATE test SET v = ? WHERE k = 0 IF v = ?;
                   INSERT INTO test (k, id) VALUES (0, ?) IF NOT EXISTS;
                APPLY BATCH
            """)

        for n in range(0, threads):
            workers.append(Worker(n, c, iterations, q))

        start = time.time()

        for w in workers:
            w.start()

        for w in workers:
            w.join()

        if verbose:
            runtime = time.time() - start
            print "runtime:", runtime

        query = SimpleStatement("SELECT v FROM test WHERE k = 0", consistency_level=ConsistencyLevel.ALL)
        rows = cursor.execute(query)
        value = rows[0][0]

        errors = 0
        retries = 0
        for w in workers:
            errors = errors + w.errors
            retries = retries + w.retries

        assert (value == threads * iterations) and (errors == 0), "value=%d, errors=%d, retries=%d" % (value, errors, retries)

<|endoftext|>


<commit_before>import datetime
import logging
import math
import traceback

import pytz

from bson import json_util, ObjectId

import rmc.shared.constants as c

def json_loads(json_str):
    return json_util.loads(json_str)

def json_dumps(obj):
    return json_util.dumps(obj).replace('</', '<\\/')

def dict_to_list(dikt):
    # TODO(jlfwong): This function's name is horribly misleading about what it
    # does - rename and document
    update_with_name = lambda key, val: dict(val, **{ 'name': key })
    return [update_with_name(k, v) for k, v in dikt.iteritems()]

def get_term_id_for_date(the_date):
    # From http://ugradcalendar.uwaterloo.ca/page/uWaterloo-Calendar-Events-and-Academic-Deadlines
    # Seems should be usually right; just not sure of Spring term always
    # starting on May...
    # TODO(david): uWaterloo specific
    term_start_months = [9, 5, 1]

    # Find the month this term started
    for month in term_start_months:
        if the_date.month >= month:
            start_month = month
            break

    return "%d_%02d" % (the_date.year, start_month)

def get_current_term_id():
    return get_term_id_for_date(datetime.datetime.now())

# Ported Ruby's Statistics2.pnormaldist(qn) to Python
# http://stackoverflow.com/questions/6116770/whats-the-equivalent-of-rubys-pnormaldist-statistics-function-in-haskell
# inverse of normal distribution ([2])
# Pr( (-\infty, x] ) = qn -> x
def pnormaldist(qn):
    b = [1.570796288, 0.03706987906, -0.8364353589e-3,
        -0.2250947176e-3, 0.6841218299e-5, 0.5824238515e-5,
        -0.104527497e-5, 0.8360937017e-7, -0.3231081277e-8,
        0.3657763036e-10, 0.6936233982e-12]

    if qn < 0.0 or 1.0 < qn:
        logging.error("Error : qn <= 0 or qn >= 1  in pnorm()!")
        return 0.0;

    if qn == 0.5:
        return 0.0

    w1 = qn
    if qn > 0.5:
        w1 = 1.0 - w1
    w3 = -math.log(4.0 * w1 * (1.0 - w1))
    w1 = b[0]
    for i in range(1, 11):
        w1 += b[i] * w3**i;

    if qn > 0.5:
        return math.sqrt(w1 * w3)

    return -math.sqrt(w1 * w3)

def get_sorting_score(phat, n, confidence=c.RATINGS_CONFIDENCE):
    """
    Get the score used for sorting by ratings

    Returns the lower bound on Wilson Score. See
    http://evanmiller.org/how-not-to-sort-by-average-rating.html

    Args:
        phat: The observed proportion of positive ratings (0 <= phat <= 1)
        n: The total number of ratings
        confidence: How much confidences we want for this to be the lower bound?
    """
    if n == 0:
        return 0

    try:
        if confidence == c.RATINGS_CONFIDENCE:
            z = 1.9599639715843482
        else:
            z = pnormaldist(1-(1-confidence)/2)
        # Modified to optimize for our data model
        #phat = 1.0*pos/n
        retVal = (phat + z*z/(2*n) -
                z * math.sqrt((phat*(1-phat)+z*z/(4*n))/n))/(1+z*z/n)
    except:
        # This should never happen, so we should debug this case as soon as we
        # get the error.
        # Returning phat should be the same as calling the function with n=INF.
        # While this is bad, it's better than nothing and we can fix it
        # with the re-aggregator.
        logging.error('get_sorting_score(%s, %s, %s) threw an exception'
                % (phat, n, confidence))
        logging.error(' '.join(traceback.format_stack()))
        retVal = max(0, min(1, phat))
    return retVal

def flatten_dict(dikt):
    """Flatten dict into 1 level by JSON-encoding all non-primitive values."""
    flattened = {}
    for k, v in dikt.iteritems():
        if isinstance(v, dict) or isinstance(v, list):
            flattened[k] = json_util.dumps(v)
        elif isinstance(v, ObjectId):
            flattened[k] = str(v)
        else:
            flattened[k] = v
    return flattened

def eastern_to_utc(date):
    tz = pytz.timezone('US/Eastern')
    return utc_date(date, tz)

def utc_date(date, tz):
    return tz.normalize(tz.localize(date)).astimezone(pytz.utc)

def to_dict(doc, fields):
    # TODO(jlfwong): This looks like it's only used in one place and should
    # be killed off
    """Warning: Using this convenience fn is probably not as efficient as the
    plain old manually building up a dict."""
    def map_field(prop):
        val = getattr(doc, prop)
        return val.to_dict() if hasattr(val, 'to_dict') else val

    return { f: map_field(f) for f in fields }

def freshness_filter(objs, to_date_func):
    # Only return results from within the past 2 non-leap years
    date_limit = datetime.datetime.now() - datetime.timedelta(days=2 * 365)

    return filter(lambda obj: to_date_func(obj) and to_date_func(obj) >= date_limit, objs)
<commit_msg>Add num_days param to freshness_filter<commit_after>import datetime
import logging
import math
import traceback

import pytz

from bson import json_util, ObjectId

import rmc.shared.constants as c

NUM_DAYS_FRESH_DATA = 2 * 365

def json_loads(json_str):
    return json_util.loads(json_str)

def json_dumps(obj):
    return json_util.dumps(obj).replace('</', '<\\/')

def dict_to_list(dikt):
    # TODO(jlfwong): This function's name is horribly misleading about what it
    # does - rename and document
    update_with_name = lambda key, val: dict(val, **{ 'name': key })
    return [update_with_name(k, v) for k, v in dikt.iteritems()]

def get_term_id_for_date(the_date):
    # From http://ugradcalendar.uwaterloo.ca/page/uWaterloo-Calendar-Events-and-Academic-Deadlines
    # Seems should be usually right; just not sure of Spring term always
    # starting on May...
    # TODO(david): uWaterloo specific
    term_start_months = [9, 5, 1]

    # Find the month this term started
    for month in term_start_months:
        if the_date.month >= month:
            start_month = month
            break

    return "%d_%02d" % (the_date.year, start_month)

def get_current_term_id():
    return get_term_id_for_date(datetime.datetime.now())

# Ported Ruby's Statistics2.pnormaldist(qn) to Python
# http://stackoverflow.com/questions/6116770/whats-the-equivalent-of-rubys-pnormaldist-statistics-function-in-haskell
# inverse of normal distribution ([2])
# Pr( (-\infty, x] ) = qn -> x
def pnormaldist(qn):
    b = [1.570796288, 0.03706987906, -0.8364353589e-3,
        -0.2250947176e-3, 0.6841218299e-5, 0.5824238515e-5,
        -0.104527497e-5, 0.8360937017e-7, -0.3231081277e-8,
        0.3657763036e-10, 0.6936233982e-12]

    if qn < 0.0 or 1.0 < qn:
        logging.error("Error : qn <= 0 or qn >= 1  in pnorm()!")
        return 0.0;

    if qn == 0.5:
        return 0.0

    w1 = qn
    if qn > 0.5:
        w1 = 1.0 - w1
    w3 = -math.log(4.0 * w1 * (1.0 - w1))
    w1 = b[0]
    for i in range(1, 11):
        w1 += b[i] * w3**i;

    if qn > 0.5:
        return math.sqrt(w1 * w3)

    return -math.sqrt(w1 * w3)

def get_sorting_score(phat, n, confidence=c.RATINGS_CONFIDENCE):
    """
    Get the score used for sorting by ratings

    Returns the lower bound on Wilson Score. See
    http://evanmiller.org/how-not-to-sort-by-average-rating.html

    Args:
        phat: The observed proportion of positive ratings (0 <= phat <= 1)
        n: The total number of ratings
        confidence: How much confidences we want for this to be the lower bound?
    """
    if n == 0:
        return 0

    try:
        if confidence == c.RATINGS_CONFIDENCE:
            z = 1.9599639715843482
        else:
            z = pnormaldist(1-(1-confidence)/2)
        # Modified to optimize for our data model
        #phat = 1.0*pos/n
        retVal = (phat + z*z/(2*n) -
                z * math.sqrt((phat*(1-phat)+z*z/(4*n))/n))/(1+z*z/n)
    except:
        # This should never happen, so we should debug this case as soon as we
        # get the error.
        # Returning phat should be the same as calling the function with n=INF.
        # While this is bad, it's better than nothing and we can fix it
        # with the re-aggregator.
        logging.error('get_sorting_score(%s, %s, %s) threw an exception'
                % (phat, n, confidence))
        logging.error(' '.join(traceback.format_stack()))
        retVal = max(0, min(1, phat))
    return retVal

def flatten_dict(dikt):
    """Flatten dict into 1 level by JSON-encoding all non-primitive values."""
    flattened = {}
    for k, v in dikt.iteritems():
        if isinstance(v, dict) or isinstance(v, list):
            flattened[k] = json_util.dumps(v)
        elif isinstance(v, ObjectId):
            flattened[k] = str(v)
        else:
            flattened[k] = v
    return flattened

def eastern_to_utc(date):
    tz = pytz.timezone('US/Eastern')
    return utc_date(date, tz)

def utc_date(date, tz):
    return tz.normalize(tz.localize(date)).astimezone(pytz.utc)

def to_dict(doc, fields):
    # TODO(jlfwong): This looks like it's only used in one place and should
    # be killed off
    """Warning: Using this convenience fn is probably not as efficient as the
    plain old manually building up a dict."""
    def map_field(prop):
        val = getattr(doc, prop)
        return val.to_dict() if hasattr(val, 'to_dict') else val

    return { f: map_field(f) for f in fields }

def freshness_filter(objs, to_date_func, num_days=None):
    """Return results from within the past num_days days."""
    if num_days is None:
        num_days = NUM_DAYS_FRESH_DATA

    date_limit = datetime.datetime.now() - datetime.timedelta(days=num_days)

    return filter(lambda obj: to_date_func(obj) and to_date_func(obj) >= date_limit, objs)
<|endoftext|>