<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2010 Radim Rehurek <radimrehurek@seznam.cz>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html

"""
Automated tests for checking transformation algorithms (the models package).
"""

import logging
import unittest
import os
import os.path
import tempfile

import six
import numpy
import scipy.linalg

from gensim.corpora import mmcorpus, Dictionary
from gensim.models import ldamodel, ldamulticore
from gensim import matutils

module_path = os.path.dirname(__file__) # needed because sample data files are located in the same folder
datapath = lambda fname: os.path.join(module_path, 'test_data', fname)

# set up vars used in testing ("Deerwester" from the web tutorial)
texts = [['human', 'interface', 'computer'],
 ['survey', 'user', 'computer', 'system', 'response', 'time'],
 ['eps', 'user', 'interface', 'system'],
 ['system', 'human', 'system', 'eps'],
 ['user', 'response', 'time'],
 ['trees'],
 ['graph', 'trees'],
 ['graph', 'minors', 'trees'],
 ['graph', 'minors', 'survey']]
dictionary = Dictionary(texts)
corpus = [dictionary.doc2bow(text) for text in texts]

def testfile():
    # temporary data will be stored to this file
    return os.path.join(tempfile.gettempdir(), 'gensim_models.tst')

class TestLdaModel(unittest.TestCase):
    def setUp(self):
        self.corpus = mmcorpus.MmCorpus(datapath('testcorpus.mm'))
        self.class_ = ldamodel.LdaModel
        self.model = self.class_(corpus, id2word=dictionary, num_topics=2, passes=100)

    def testTransform(self):
        passed = False
        # sometimes, LDA training gets stuck at a local minimum
        # in that case try re-training the model from scratch, hoping for a
        # better random initialization
        for i in range(5): # restart at most 5 times
            # create the transformation model
            model = self.class_(id2word=dictionary, num_topics=2, passes=100)
            model.update(self.corpus)

            # transform one document
            doc = list(corpus)[0]
            transformed = model[doc]

            vec = matutils.sparse2full(transformed, 2) # convert to dense vector, for easier equality tests
            expected = [0.13, 0.87]
            passed = numpy.allclose(sorted(vec), sorted(expected), atol=1e-2) # must contain the same values, up to re-ordering
            if passed:
                break
            logging.warning("LDA failed to converge on attempt %i (got %s, expected %s)" %
                            (i, sorted(vec), sorted(expected)))
        self.assertTrue(passed)

    def testAlphaAuto(self):
        model1 = self.class_(corpus, id2word=dictionary, alpha='symmetric', passes=10)
        modelauto = self.class_(corpus, id2word=dictionary, alpha='auto', passes=10)

        # did we learn something?
        self.assertFalse(all(numpy.equal(model1.alpha, modelauto.alpha)))

    def testAlpha(self):
        kwargs = dict(
            id2word=dictionary,
            num_topics=2,
            alpha=None
        )
        expected_shape = (2,)

        # should not raise anything
        self.class_(**kwargs)

        kwargs['alpha'] = 'symmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([0.5, 0.5])))

        kwargs['alpha'] = 'asymmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(numpy.allclose(model.alpha, [0.630602, 0.369398]))

        kwargs['alpha'] = 0.3
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([0.3, 0.3])))

        kwargs['alpha'] = 3
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([3, 3])))

        kwargs['alpha'] = [0.3, 0.3]
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([0.3, 0.3])))

        # all should raise an exception for being wrong shape
        kwargs['alpha'] = [0.3, 0.3, 0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['alpha'] = [[0.3], [0.3]]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['alpha'] = [0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['alpha'] = "gensim is cool"
        self.assertRaises(ValueError, self.class_, **kwargs)

    def testEtaAuto(self):
        model1 = self.class_(corpus, id2word=dictionary, eta='symmetric', passes=10)
        modelauto = self.class_(corpus, id2word=dictionary, eta='auto', passes=10)

        # did we learn something?
        self.assertFalse(all(numpy.equal(model1.eta, modelauto.eta)))

    def testEta(self):
        kwargs = dict(
            id2word=dictionary,
            num_topics=2,
            eta=None
        )
        expected_shape = (2, 1)

        # should not raise anything
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.5], [0.5]])))

        kwargs['eta'] = 'symmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.5], [0.5]])))

        kwargs['eta'] = 'asymmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(numpy.allclose(model.eta, [[0.630602], [0.369398]]))

        kwargs['eta'] = 0.3
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.3], [0.3]])))

        kwargs['eta'] = 3
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[3], [3]])))

        kwargs['eta'] = [[0.3], [0.3]]
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.3], [0.3]])))

        kwargs['eta'] = [0.3, 0.3]
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.3], [0.3]])))

        # should be ok with num_topics x num_terms
        testeta = numpy.array([[0.5] * len(dictionary)] * 2)
        kwargs['eta'] = testeta
        self.class_(**kwargs)

        # all should raise an exception for being wrong shape
        kwargs['eta'] = testeta.reshape(tuple(reversed(testeta.shape)))
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['eta'] = [0.3, 0.3, 0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['eta'] = [0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['eta'] = "gensim is cool"
        self.assertRaises(ValueError, self.class_, **kwargs)

    def testTopTopics(self):
        top_topics = self.model.top_topics(self.corpus)

        for topic, score in top_topics:
            self.assertTrue(isinstance(topic, list))
            self.assertTrue(isinstance(score, float))

            for v, k in topic:
                self.assertTrue(isinstance(k, six.string_types))
                self.assertTrue(isinstance(v, float))

    def testGetTopicTerms(self):
        topic_terms = self.model.get_topic_terms(1)

        for k, v in topic_terms:
            self.assertTrue(isinstance(k, int) or isinstance(k, numpy.int32) or isinstance(k,numpy.int64))
            self.assertTrue(isinstance(v, float))

    def testShowTopic(self):
        topic = self.model.show_topic(1)

        for k, v in topic:
            self.assertTrue(isinstance(k, six.string_types))
            self.assertTrue(isinstance(v, float))

    def testShowTopics(self):
        topics = self.model.show_topics(formatted=False)

        for topic_no, topic in topics:
            self.assertTrue(isinstance(topic_no, int))
            self.assertTrue(isinstance(topic, list))
            for k, v in topic:
                self.assertTrue(isinstance(k, six.string_types))
                self.assertTrue(isinstance(v, float))

    def testGetDocumentTopics(self):
        doc_topics = self.model.get_document_topics(self.corpus)

        for topic in doc_topics:
            self.assertTrue(isinstance(topic, list))
            for k, v in topic:
                self.assertTrue(isinstance(k, int))
                self.assertTrue(isinstance(v, float))

        doc_topics = self.model.get_document_topics(self.corpus[0])

        for k, v in doc_topics:
            self.assertTrue(isinstance(k, int))
            self.assertTrue(isinstance(v, float))

    def testPasses(self):
        # long message includes the original error message with a custom one
        self.longMessage = True

        # construct what we expect when passes aren't involved
        test_rhots = list()
        model = self.class_(id2word=dictionary, chunksize=1, num_topics=2)
        final_rhot = lambda: pow(model.offset + (1 * model.num_updates) / model.chunksize, -model.decay)

        # generate 5 updates to test rhot on
        for x in range(5):
            model.update(self.corpus)
            test_rhots.append(final_rhot())

        for passes in [1, 5, 10, 50, 100]:
            model = self.class_(id2word=dictionary, chunksize=1, num_topics=2, passes=passes)
            self.assertEqual(final_rhot(), 1.0)
            # make sure the rhot matches the test after each update
            for test_rhot in test_rhots:
                model.update(self.corpus)

                msg = ", ".join(map(str, [passes, model.num_updates, model.state.numdocs]))
                self.assertAlmostEqual(final_rhot(), test_rhot, msg=msg)

            self.assertEqual(model.state.numdocs, len(corpus) * len(test_rhots))
            self.assertEqual(model.num_updates, len(corpus) * len(test_rhots))

    # def testTopicSeeding(self):
    #     for topic in range(2):
    #         passed = False
    #         for i in range(5):  # restart at most this many times, to mitigate LDA randomness
    #             # try seeding it both ways round, check you get the same
    #             # topics out but with which way round they are depending
    #             # on the way round they're seeded
    #             eta = numpy.ones((2, len(dictionary))) * 0.5
    #             system = dictionary.token2id[u'system']
    #             trees = dictionary.token2id[u'trees']

    #             # aggressively seed the word 'system', in one of the
    #             # two topics, 10 times higher than the other words
    #             eta[topic, system] *= 10.0

    #             model = self.class_(id2word=dictionary, num_topics=2, passes=200, eta=eta)
    #             model.update(self.corpus)

    #             topics = [dict((word, p) for p, word in model.show_topic(j, topn=None)) for j in range(2)]

    #             # check that the word 'system' in the topic we seeded got a high weight,
    #             # and the word 'trees' (the main word in the other topic) a low weight --
    #             # and vice versa for the other topic (which we didn't seed with 'system')
    #             passed = (
    #                 (topics[topic][u'system'] > topics[topic][u'trees'])
    #                 and
    #                 (topics[1 - topic][u'system'] < topics[1 - topic][u'trees'])
    #             )
    #             if passed:
    #                 break
    #             logging.warning("LDA failed to converge on attempt %i (got %s)", i, topics)
    #         self.assertTrue(passed)

    def testPersistence(self):
        fname = testfile()
        model = self.model
        model.save(fname)
        model2 = self.class_.load(fname)
        self.assertEqual(model.num_topics, model2.num_topics)
        self.assertTrue(numpy.allclose(model.expElogbeta, model2.expElogbeta))
        tstvec = []
        self.assertTrue(numpy.allclose(model[tstvec], model2[tstvec])) # try projecting an empty vector

    def testPersistenceIgnore(self):
        fname = testfile()
        model = ldamodel.LdaModel(self.corpus, num_topics=2)
        model.save(fname, ignore='id2word')
        model2 = ldamodel.LdaModel.load(fname)
        self.assertTrue(model2.id2word is None)

        model.save(fname, ignore=['id2word'])
        model2 = ldamodel.LdaModel.load(fname)
        self.assertTrue(model2.id2word is None)

    def testPersistenceCompressed(self):
        fname = testfile() + '.gz'
        model = self.model
        model.save(fname)
        model2 = self.class_.load(fname, mmap=None)
        self.assertEqual(model.num_topics, model2.num_topics)
        self.assertTrue(numpy.allclose(model.expElogbeta, model2.expElogbeta))
        tstvec = []
        self.assertTrue(numpy.allclose(model[tstvec], model2[tstvec])) # try projecting an empty vector

    def testLargeMmap(self):
        fname = testfile()
        model = self.model

        # simulate storing large arrays separately
        model.save(testfile(), sep_limit=0)

        # test loading the large model arrays with mmap
        model2 = self.class_.load(testfile(), mmap='r')
        self.assertEqual(model.num_topics, model2.num_topics)
        self.assertTrue(isinstance(model2.expElogbeta, numpy.memmap))
        self.assertTrue(numpy.allclose(model.expElogbeta, model2.expElogbeta))
        tstvec = []
        self.assertTrue(numpy.allclose(model[tstvec], model2[tstvec])) # try projecting an empty vector

    def testLargeMmapCompressed(self):
        fname = testfile() + '.gz'
        model = self.model

        # simulate storing large arrays separately
        model.save(fname, sep_limit=0)

        # test loading the large model arrays with mmap
        self.assertRaises(IOError, self.class_.load, fname, mmap='r')

#endclass TestLdaModel

class TestLdaMulticore(TestLdaModel):
    def setUp(self):
        self.corpus = mmcorpus.MmCorpus(datapath('testcorpus.mm'))
        self.class_ = ldamulticore.LdaMulticore
        self.model = self.class_(corpus, id2word=dictionary, num_topics=2, passes=100)

    # override LdaModel because multicore does not allow alpha=auto
    def testAlphaAuto(self):
        self.assertRaises(RuntimeError, self.class_, alpha='auto')

#endclass TestLdaMulticore

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.DEBUG)
    unittest.main()
<commit_msg>lda print type of k<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2010 Radim Rehurek <radimrehurek@seznam.cz>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html

"""
Automated tests for checking transformation algorithms (the models package).
"""

import logging
import unittest
import os
import os.path
import tempfile

import six
import numpy
import scipy.linalg

from gensim.corpora import mmcorpus, Dictionary
from gensim.models import ldamodel, ldamulticore
from gensim import matutils

module_path = os.path.dirname(__file__) # needed because sample data files are located in the same folder
datapath = lambda fname: os.path.join(module_path, 'test_data', fname)

# set up vars used in testing ("Deerwester" from the web tutorial)
texts = [['human', 'interface', 'computer'],
 ['survey', 'user', 'computer', 'system', 'response', 'time'],
 ['eps', 'user', 'interface', 'system'],
 ['system', 'human', 'system', 'eps'],
 ['user', 'response', 'time'],
 ['trees'],
 ['graph', 'trees'],
 ['graph', 'minors', 'trees'],
 ['graph', 'minors', 'survey']]
dictionary = Dictionary(texts)
corpus = [dictionary.doc2bow(text) for text in texts]

def testfile():
    # temporary data will be stored to this file
    return os.path.join(tempfile.gettempdir(), 'gensim_models.tst')

class TestLdaModel(unittest.TestCase):
    def setUp(self):
        self.corpus = mmcorpus.MmCorpus(datapath('testcorpus.mm'))
        self.class_ = ldamodel.LdaModel
        self.model = self.class_(corpus, id2word=dictionary, num_topics=2, passes=100)

    def testTransform(self):
        passed = False
        # sometimes, LDA training gets stuck at a local minimum
        # in that case try re-training the model from scratch, hoping for a
        # better random initialization
        for i in range(5): # restart at most 5 times
            # create the transformation model
            model = self.class_(id2word=dictionary, num_topics=2, passes=100)
            model.update(self.corpus)

            # transform one document
            doc = list(corpus)[0]
            transformed = model[doc]

            vec = matutils.sparse2full(transformed, 2) # convert to dense vector, for easier equality tests
            expected = [0.13, 0.87]
            passed = numpy.allclose(sorted(vec), sorted(expected), atol=1e-2) # must contain the same values, up to re-ordering
            if passed:
                break
            logging.warning("LDA failed to converge on attempt %i (got %s, expected %s)" %
                            (i, sorted(vec), sorted(expected)))
        self.assertTrue(passed)

    def testAlphaAuto(self):
        model1 = self.class_(corpus, id2word=dictionary, alpha='symmetric', passes=10)
        modelauto = self.class_(corpus, id2word=dictionary, alpha='auto', passes=10)

        # did we learn something?
        self.assertFalse(all(numpy.equal(model1.alpha, modelauto.alpha)))

    def testAlpha(self):
        kwargs = dict(
            id2word=dictionary,
            num_topics=2,
            alpha=None
        )
        expected_shape = (2,)

        # should not raise anything
        self.class_(**kwargs)

        kwargs['alpha'] = 'symmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([0.5, 0.5])))

        kwargs['alpha'] = 'asymmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(numpy.allclose(model.alpha, [0.630602, 0.369398]))

        kwargs['alpha'] = 0.3
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([0.3, 0.3])))

        kwargs['alpha'] = 3
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([3, 3])))

        kwargs['alpha'] = [0.3, 0.3]
        model = self.class_(**kwargs)
        self.assertEqual(model.alpha.shape, expected_shape)
        self.assertTrue(all(model.alpha == numpy.array([0.3, 0.3])))

        # all should raise an exception for being wrong shape
        kwargs['alpha'] = [0.3, 0.3, 0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['alpha'] = [[0.3], [0.3]]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['alpha'] = [0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['alpha'] = "gensim is cool"
        self.assertRaises(ValueError, self.class_, **kwargs)

    def testEtaAuto(self):
        model1 = self.class_(corpus, id2word=dictionary, eta='symmetric', passes=10)
        modelauto = self.class_(corpus, id2word=dictionary, eta='auto', passes=10)

        # did we learn something?
        self.assertFalse(all(numpy.equal(model1.eta, modelauto.eta)))

    def testEta(self):
        kwargs = dict(
            id2word=dictionary,
            num_topics=2,
            eta=None
        )
        expected_shape = (2, 1)

        # should not raise anything
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.5], [0.5]])))

        kwargs['eta'] = 'symmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.5], [0.5]])))

        kwargs['eta'] = 'asymmetric'
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(numpy.allclose(model.eta, [[0.630602], [0.369398]]))

        kwargs['eta'] = 0.3
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.3], [0.3]])))

        kwargs['eta'] = 3
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[3], [3]])))

        kwargs['eta'] = [[0.3], [0.3]]
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.3], [0.3]])))

        kwargs['eta'] = [0.3, 0.3]
        model = self.class_(**kwargs)
        self.assertEqual(model.eta.shape, expected_shape)
        self.assertTrue(all(model.eta == numpy.array([[0.3], [0.3]])))

        # should be ok with num_topics x num_terms
        testeta = numpy.array([[0.5] * len(dictionary)] * 2)
        kwargs['eta'] = testeta
        self.class_(**kwargs)

        # all should raise an exception for being wrong shape
        kwargs['eta'] = testeta.reshape(tuple(reversed(testeta.shape)))
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['eta'] = [0.3, 0.3, 0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['eta'] = [0.3]
        self.assertRaises(AssertionError, self.class_, **kwargs)

        kwargs['eta'] = "gensim is cool"
        self.assertRaises(ValueError, self.class_, **kwargs)

    def testTopTopics(self):
        top_topics = self.model.top_topics(self.corpus)

        for topic, score in top_topics:
            self.assertTrue(isinstance(topic, list))
            self.assertTrue(isinstance(score, float))

            for v, k in topic:
                self.assertTrue(isinstance(k, six.string_types))
                self.assertTrue(isinstance(v, float))

    def testGetTopicTerms(self):
        topic_terms = self.model.get_topic_terms(1)

        for k, v in topic_terms:
            print type(k)
            self.assertTrue(isinstance(k, int) or isinstance(k, numpy.int32) or isinstance(k,numpy.int64))
            self.assertTrue(isinstance(v, float))

    def testShowTopic(self):
        topic = self.model.show_topic(1)

        for k, v in topic:
            self.assertTrue(isinstance(k, six.string_types))
            self.assertTrue(isinstance(v, float))

    def testShowTopics(self):
        topics = self.model.show_topics(formatted=False)

        for topic_no, topic in topics:
            self.assertTrue(isinstance(topic_no, int))
            self.assertTrue(isinstance(topic, list))
            for k, v in topic:
                self.assertTrue(isinstance(k, six.string_types))
                self.assertTrue(isinstance(v, float))

    def testGetDocumentTopics(self):
        doc_topics = self.model.get_document_topics(self.corpus)

        for topic in doc_topics:
            self.assertTrue(isinstance(topic, list))
            for k, v in topic:
                self.assertTrue(isinstance(k, int))
                self.assertTrue(isinstance(v, float))

        doc_topics = self.model.get_document_topics(self.corpus[0])

        for k, v in doc_topics:
            self.assertTrue(isinstance(k, int))
            self.assertTrue(isinstance(v, float))

    def testPasses(self):
        # long message includes the original error message with a custom one
        self.longMessage = True

        # construct what we expect when passes aren't involved
        test_rhots = list()
        model = self.class_(id2word=dictionary, chunksize=1, num_topics=2)
        final_rhot = lambda: pow(model.offset + (1 * model.num_updates) / model.chunksize, -model.decay)

        # generate 5 updates to test rhot on
        for x in range(5):
            model.update(self.corpus)
            test_rhots.append(final_rhot())

        for passes in [1, 5, 10, 50, 100]:
            model = self.class_(id2word=dictionary, chunksize=1, num_topics=2, passes=passes)
            self.assertEqual(final_rhot(), 1.0)
            # make sure the rhot matches the test after each update
            for test_rhot in test_rhots:
                model.update(self.corpus)

                msg = ", ".join(map(str, [passes, model.num_updates, model.state.numdocs]))
                self.assertAlmostEqual(final_rhot(), test_rhot, msg=msg)

            self.assertEqual(model.state.numdocs, len(corpus) * len(test_rhots))
            self.assertEqual(model.num_updates, len(corpus) * len(test_rhots))

    # def testTopicSeeding(self):
    #     for topic in range(2):
    #         passed = False
    #         for i in range(5):  # restart at most this many times, to mitigate LDA randomness
    #             # try seeding it both ways round, check you get the same
    #             # topics out but with which way round they are depending
    #             # on the way round they're seeded
    #             eta = numpy.ones((2, len(dictionary))) * 0.5
    #             system = dictionary.token2id[u'system']
    #             trees = dictionary.token2id[u'trees']

    #             # aggressively seed the word 'system', in one of the
    #             # two topics, 10 times higher than the other words
    #             eta[topic, system] *= 10.0

    #             model = self.class_(id2word=dictionary, num_topics=2, passes=200, eta=eta)
    #             model.update(self.corpus)

    #             topics = [dict((word, p) for p, word in model.show_topic(j, topn=None)) for j in range(2)]

    #             # check that the word 'system' in the topic we seeded got a high weight,
    #             # and the word 'trees' (the main word in the other topic) a low weight --
    #             # and vice versa for the other topic (which we didn't seed with 'system')
    #             passed = (
    #                 (topics[topic][u'system'] > topics[topic][u'trees'])
    #                 and
    #                 (topics[1 - topic][u'system'] < topics[1 - topic][u'trees'])
    #             )
    #             if passed:
    #                 break
    #             logging.warning("LDA failed to converge on attempt %i (got %s)", i, topics)
    #         self.assertTrue(passed)

    def testPersistence(self):
        fname = testfile()
        model = self.model
        model.save(fname)
        model2 = self.class_.load(fname)
        self.assertEqual(model.num_topics, model2.num_topics)
        self.assertTrue(numpy.allclose(model.expElogbeta, model2.expElogbeta))
        tstvec = []
        self.assertTrue(numpy.allclose(model[tstvec], model2[tstvec])) # try projecting an empty vector

    def testPersistenceIgnore(self):
        fname = testfile()
        model = ldamodel.LdaModel(self.corpus, num_topics=2)
        model.save(fname, ignore='id2word')
        model2 = ldamodel.LdaModel.load(fname)
        self.assertTrue(model2.id2word is None)

        model.save(fname, ignore=['id2word'])
        model2 = ldamodel.LdaModel.load(fname)
        self.assertTrue(model2.id2word is None)

    def testPersistenceCompressed(self):
        fname = testfile() + '.gz'
        model = self.model
        model.save(fname)
        model2 = self.class_.load(fname, mmap=None)
        self.assertEqual(model.num_topics, model2.num_topics)
        self.assertTrue(numpy.allclose(model.expElogbeta, model2.expElogbeta))
        tstvec = []
        self.assertTrue(numpy.allclose(model[tstvec], model2[tstvec])) # try projecting an empty vector

    def testLargeMmap(self):
        fname = testfile()
        model = self.model

        # simulate storing large arrays separately
        model.save(testfile(), sep_limit=0)

        # test loading the large model arrays with mmap
        model2 = self.class_.load(testfile(), mmap='r')
        self.assertEqual(model.num_topics, model2.num_topics)
        self.assertTrue(isinstance(model2.expElogbeta, numpy.memmap))
        self.assertTrue(numpy.allclose(model.expElogbeta, model2.expElogbeta))
        tstvec = []
        self.assertTrue(numpy.allclose(model[tstvec], model2[tstvec])) # try projecting an empty vector

    def testLargeMmapCompressed(self):
        fname = testfile() + '.gz'
        model = self.model

        # simulate storing large arrays separately
        model.save(fname, sep_limit=0)

        # test loading the large model arrays with mmap
        self.assertRaises(IOError, self.class_.load, fname, mmap='r')

#endclass TestLdaModel

class TestLdaMulticore(TestLdaModel):
    def setUp(self):
        self.corpus = mmcorpus.MmCorpus(datapath('testcorpus.mm'))
        self.class_ = ldamulticore.LdaMulticore
        self.model = self.class_(corpus, id2word=dictionary, num_topics=2, passes=100)

    # override LdaModel because multicore does not allow alpha=auto
    def testAlphaAuto(self):
        self.assertRaises(RuntimeError, self.class_, alpha='auto')

#endclass TestLdaMulticore

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.DEBUG)
    unittest.main()
<|endoftext|>


<commit_before>import sys
import timeit
import functools
import operator as op

from cStringIO import StringIO

def disable_test(f):
    return 'disabled'

def baseline(loop_count):
    for num in xrange(loop_count):
        str(num)

def baseline2(loop_count):
    lst = []
    for num in xrange(loop_count):
        lst.append(str(num))
    return lst

def baseline3(loop_count):
    return [str(int) for int in xrange(loop_count)]

def str_sum(loop_count):
    out_str = ''
    for num in xrange(loop_count):
        out_str += str(num)
    return out_str

def str_join(loop_count):
    str_list = []
    for num in xrange(loop_count):
        str_list.append(str(num))
    out_str = ''.join(str_list)
    return out_str

def str_join_lc(loop_count):
    return ''.join(str(num) for num in xrange(loop_count))

def string_io(loop_count):
    file_str = StringIO()
    for num in xrange(loop_count):
        file_str.write(str(num))
    out_str = file_str.getvalue()
    return out_str

def evaluate_range(range_str):
    def aux(range_str):
        return range(*(int(p) for p in range_str.split(':')))

    if 'e' in range_str:
        base, range_str = range_str.split('e')
        base = int(base)
        return [base**x for x in aux(range_str)]
    else:
        return aux(range_str)

def bench(loop_count, functions):
    for function in functions:
        #t = timeit.Timer(functools.partial(function, loop_count))
        t = timeit.Timer('%s(%d)' % (function.func_name, loop_count),
                         setup='from __main__ import %s' % function.func_name)
        sys.stdout.write('%f\t' % t.timeit(number = 1))

def main():
    input_sizes = []
    selected_function_names = []
    for arg in sys.argv[1:]:
        try:
            input_size = int(arg)
        except ValueError:
            try:
                range_ = evaluate_range(arg)
                input_sizes.extend(range_)
            except Exception:
                selected_function_names.append(arg)
        else:
            input_sizes.append(input_size)

    if selected_function_names:
        functions = [function for (name, function) in globals().items()
                   if name in selected_function_names and callable(function)]
    else:
        functions = tuple([baseline] +
                        [f for (name, f) in globals().items()
                         if ('str' in name) and callable(f)])

    functions = sorted(functions, key=op.attrgetter('func_name'))
    if not input_sizes:
        input_sizes.append(100000)

    sys.stdout.write('its\t')
    for function in functions:
        sys.stdout.write('%s\t' % function.func_name.ljust(8))
    sys.stdout.write('\n')
    for loop_count in sorted(input_sizes):
        sys.stdout.write('%d\t' % loop_count)
        bench(loop_count, functions)
        print

if __name__ == '__main__':
    main()<commit_msg>Added variant of tested function.<commit_after>import sys
import timeit
import functools
import operator as op

from cStringIO import StringIO

def disable_test(f):
    return 'disabled'

def baseline(loop_count):
    for num in xrange(loop_count):
        str(num)

def baseline2(loop_count):
    lst = []
    for num in xrange(loop_count):
        lst.append(str(num))
    return lst

def baseline3(loop_count):
    return [str(int) for int in xrange(loop_count)]

def str_sum(loop_count):
    out_str = ''
    for num in xrange(loop_count):
        out_str += str(num)
    return out_str

def str_sum_b(loop_count):
    out_str = ''
    for num in xrange(loop_count):
        out_str = out_str + str(num)
    return out_str

def str_join(loop_count):
    str_list = []
    for num in xrange(loop_count):
        str_list.append(str(num))
    out_str = ''.join(str_list)
    return out_str

def str_join_lc(loop_count):
    return ''.join(str(num) for num in xrange(loop_count))

def string_io(loop_count):
    file_str = StringIO()
    for num in xrange(loop_count):
        file_str.write(str(num))
    out_str = file_str.getvalue()
    return out_str

def evaluate_range(range_str):
    def aux(range_str):
        return range(*(int(p) for p in range_str.split(':')))

    if 'e' in range_str:
        base, range_str = range_str.split('e')
        base = int(base)
        return [base**x for x in aux(range_str)]
    else:
        return aux(range_str)

def bench(loop_count, functions):
    for function in functions:
        #t = timeit.Timer(functools.partial(function, loop_count))
        t = timeit.Timer('%s(%d)' % (function.func_name, loop_count),
                         setup='from __main__ import %s' % function.func_name)
        sys.stdout.write('%f\t' % t.timeit(number = 1))

def main():
    input_sizes = []
    selected_function_names = []
    for arg in sys.argv[1:]:
        try:
            input_size = int(arg)
        except ValueError:
            try:
                range_ = evaluate_range(arg)
                input_sizes.extend(range_)
            except Exception:
                selected_function_names.append(arg)
        else:
            input_sizes.append(input_size)

    if selected_function_names:
        functions = [function for (name, function) in globals().items()
                   if name in selected_function_names and callable(function)]
    else:
        functions = tuple([baseline] +
                        [f for (name, f) in globals().items()
                         if ('str' in name) and callable(f)])

    functions = sorted(functions, key=op.attrgetter('func_name'))
    if not input_sizes:
        input_sizes.append(100000)

    sys.stdout.write('its\t')
    for function in functions:
        sys.stdout.write('%s\t' % function.func_name.ljust(8))
    sys.stdout.write('\n')
    for loop_count in sorted(input_sizes):
        sys.stdout.write('%d\t' % loop_count)
        bench(loop_count, functions)
        print

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


<commit_before>"""
Custom Authenticator to use GitHub OAuth with JupyterHub

Most of the code c/o Kyle Kelley (@rgbkrk)

Extended use of GH attributes (uid/gid, email, save-token) by Adam Thornton
 (athornton@lsst.org)
"""

import json
import os
import re
import string

from tornado.auth import OAuth2Mixin
from tornado import gen, web

from tornado.httputil import url_concat
from tornado.httpclient import HTTPRequest, AsyncHTTPClient

from jupyterhub.auth import LocalAuthenticator

from traitlets import Unicode, Set

from .common import next_page_from_links
from .oauth2 import OAuthLoginHandler, OAuthenticator

# Support github.com and github enterprise installations
GITHUB_HOST = os.environ.get('GITHUB_HOST') or 'github.com'
if GITHUB_HOST == 'github.com':
    GITHUB_API = 'api.github.com'
else:
    GITHUB_API = '%s/api/v3' % GITHUB_HOST

def _api_headers(access_token):
    return {"Accept": "application/json",
            "User-Agent": "JupyterHub",
            "Authorization": "token {}".format(access_token)
            }

class GitHubMixin(OAuth2Mixin):
    _OAUTH_AUTHORIZE_URL = "https://%s/login/oauth/authorize" % GITHUB_HOST
    _OAUTH_ACCESS_TOKEN_URL = "https://%s/login/oauth/access_token" % GITHUB_HOST

class GitHubLoginHandler(OAuthLoginHandler, GitHubMixin):
    """The `scope` attribute is inherited from OAuthLoginHandler and is a
    list of scopes requested when we acquire a GitHub token:

    See github_scope.md for details.
    """

class GitHubOAuthenticator(OAuthenticator):

    login_service = "GitHub"

    # deprecated names
    github_client_id = Unicode(config=True, help="DEPRECATED")

    def _github_client_id_changed(self, name, old, new):
        self.log.warn("github_client_id is deprecated, use client_id")
        self.client_id = new
    github_client_secret = Unicode(config=True, help="DEPRECATED")

    def _github_client_secret_changed(self, name, old, new):
        self.log.warn("github_client_secret is deprecated, use client_secret")
        self.client_secret = new

    client_id_env = 'GITHUB_CLIENT_ID'
    client_secret_env = 'GITHUB_CLIENT_SECRET'
    login_handler = GitHubLoginHandler

    github_organization_whitelist = Set(
        config=True,
        help="Automatically whitelist members of selected organizations",
    )

    @gen.coroutine
    def authenticate(self, handler, data=None):
        """We set up auth_state based on additional GitHub info if we
        receive it.
        """
        code = handler.get_argument("code")
        # TODO: Configure the curl_httpclient for tornado
        http_client = AsyncHTTPClient()

        # Exchange the OAuth code for a GitHub Access Token
        #
        # See: https://developer.github.com/v3/oauth/

        # GitHub specifies a POST request yet requires URL parameters
        params = dict(
            client_id=self.client_id,
            client_secret=self.client_secret,
            code=code
        )

        url = url_concat("https://%s/login/oauth/access_token" % GITHUB_HOST,
                         params)

        req = HTTPRequest(url,
                          method="POST",
                          headers={"Accept": "application/json"},
                          body=''  # Body is required for a POST...
                          )

        resp = yield http_client.fetch(req)
        resp_json = json.loads(resp.body.decode('utf8', 'replace'))

        access_token = resp_json['access_token']

        # Determine who the logged in user is
        req = HTTPRequest("https://%s/user" % GITHUB_API,
                          method="GET",
                          headers=_api_headers(access_token)
                          )
        resp = yield http_client.fetch(req)
        resp_json = json.loads(resp.body.decode('utf8', 'replace'))

        username = resp_json["login"]
        # Check if user is a member of any whitelisted organizations.
        # This check is performed here, as it requires `access_token`.
        if self.github_organization_whitelist:
            for org in self.github_organization_whitelist:
                user_in_org = yield self._check_organization_whitelist(org, username, access_token)
                if not user_in_org:
                    # User not found in member list for any organisation
                    return None
        # username is now the GitHub userid.
        if not username:
            return None
        # Check if user is a member of any whitelisted organizations.
        # This check is performed here, as it requires `access_token`.
        if self.github_organization_whitelist:
            for org in self.github_organization_whitelist:
                user_in_org = yield self._check_organization_whitelist(org, username, access_token)
                if user_in_org:
                    break
            else:  # User not found in member list for any organisation
                return None
        userdict = {"name": username}
        # Now we set up auth_state
        auth_state = {}
        auth_state["username"] = username
        # We may want to do user provisioning in the Lab/Notebook environment.
        #  This next bit is about that.
        #  1) stash the access token
        #  2) use the GitHub ID as the uid
        #  3) set up map of orgs/gids
        #  4) set up name/email for .gitconfig
        # Store the resulting structure in auth_state
        auth_state["access_token"] = access_token
        auth_state["uid"] = resp_json["id"]
        orgs = yield self._get_user_organizations(access_token)
        if orgs:
            auth_state["organization_map"] = orgs
        auth_state["name"] = resp_json["name"]
        # A public email will return in the base query
        if "email" in resp_json and resp_json["email"]:
            auth_state["email"] = resp_json["email"]
        else:
            # Or if you gave
            email = yield self._get_user_email(access_token)
            if email:
                auth_state["email"] = email
        userdict["auth_state"] = auth_state
        return userdict

    @gen.coroutine
    def _check_organization_whitelist(self, org, username, access_token):
        http_client = AsyncHTTPClient()
        headers = _api_headers(access_token)
        # Get all the members for organization 'org'
        next_page = "https://%s/orgs/%s/members" % (GITHUB_API, org)
        while next_page:
            req = HTTPRequest(next_page, method="GET", headers=headers)
            resp = yield http_client.fetch(req)
            resp_json = json.loads(resp.body.decode('utf8', 'replace'))
            next_page = next_page_from_links(resp)
            org_members = set(entry["login"] for entry in resp_json)
            # check if any of the organizations seen so far are in whitelist
            if username in org_members:
                return True
        return False

    @gen.coroutine
    def _get_user_organizations(self, access_token):
        http_client = AsyncHTTPClient()
        headers = _api_headers(access_token)
        next_page = "https://%s/user/orgs" % (GITHUB_API)
        orgmap = {}
        while next_page:
            req = HTTPRequest(next_page, method="GET", headers=headers)
            resp = yield http_client.fetch(req)
            resp_json = json.loads(resp.body.decode('utf8', 'replace'))
            next_page = next_page_from_links(resp)
            for entry in resp_json:
                orgmap[entry["login"]] = entry["id"]
        return orgmap

    @gen.coroutine
    def _get_user_email(self, access_token):
        http_client = AsyncHTTPClient()
        headers = _api_headers(access_token)
        next_page = "https://%s/user/emails" % (GITHUB_API)
        while next_page:
            req = HTTPRequest(next_page, method="GET", headers=headers)
            resp = yield http_client.fetch(req)
            resp_json = json.loads(resp.body.decode('utf8', 'replace'))
            next_page = next_page_from_links(resp)
            for entry in resp_json:
                if "email" in entry:
                    if "primary" in entry and entry["primary"]:
                        return entry["email"]
        return None

class LocalGitHubOAuthenticator(LocalAuthenticator, GitHubOAuthenticator):

    """A version that mixes in local system user creation"""
    pass
<commit_msg>Clean up for PR discussion.<commit_after>"""
Custom Authenticator to use GitHub OAuth with JupyterHub

Most of the code c/o Kyle Kelley (@rgbkrk)

Extended use of GH attributes by Adam Thornton (athornton@lsst.org)
"""

import json
import os
import re
import string

from tornado.auth import OAuth2Mixin
from tornado import gen, web

from tornado.httputil import url_concat
from tornado.httpclient import HTTPRequest, AsyncHTTPClient

from jupyterhub.auth import LocalAuthenticator

from traitlets import Unicode, Set

from .common import next_page_from_links
from .oauth2 import OAuthLoginHandler, OAuthenticator

# Support github.com and github enterprise installations
GITHUB_HOST = os.environ.get('GITHUB_HOST') or 'github.com'
if GITHUB_HOST == 'github.com':
    GITHUB_API = 'api.github.com'
else:
    GITHUB_API = '%s/api/v3' % GITHUB_HOST

def _api_headers(access_token):
    return {"Accept": "application/json",
            "User-Agent": "JupyterHub",
            "Authorization": "token {}".format(access_token)
            }

class GitHubMixin(OAuth2Mixin):
    _OAUTH_AUTHORIZE_URL = "https://%s/login/oauth/authorize" % GITHUB_HOST
    _OAUTH_ACCESS_TOKEN_URL = "https://%s/login/oauth/access_token" % GITHUB_HOST

class GitHubLoginHandler(OAuthLoginHandler, GitHubMixin):
    """The `scope` attribute is inherited from OAuthLoginHandler and is a
    list of scopes requested when we acquire a GitHub token:

    See github_scope.md for details.
    """

class GitHubOAuthenticator(OAuthenticator):

    login_service = "GitHub"

    # deprecated names
    github_client_id = Unicode(config=True, help="DEPRECATED")

    def _github_client_id_changed(self, name, old, new):
        self.log.warn("github_client_id is deprecated, use client_id")
        self.client_id = new
    github_client_secret = Unicode(config=True, help="DEPRECATED")

    def _github_client_secret_changed(self, name, old, new):
        self.log.warn("github_client_secret is deprecated, use client_secret")
        self.client_secret = new

    client_id_env = 'GITHUB_CLIENT_ID'
    client_secret_env = 'GITHUB_CLIENT_SECRET'
    login_handler = GitHubLoginHandler

    github_organization_whitelist = Set(
        config=True,
        help="Automatically whitelist members of selected organizations",
    )

    @gen.coroutine
    def authenticate(self, handler, data=None):
        """We set up auth_state based on additional GitHub info if we
        receive it.
        """
        code = handler.get_argument("code")
        # TODO: Configure the curl_httpclient for tornado
        http_client = AsyncHTTPClient()

        # Exchange the OAuth code for a GitHub Access Token
        #
        # See: https://developer.github.com/v3/oauth/

        # GitHub specifies a POST request yet requires URL parameters
        params = dict(
            client_id=self.client_id,
            client_secret=self.client_secret,
            code=code
        )

        url = url_concat("https://%s/login/oauth/access_token" % GITHUB_HOST,
                         params)

        req = HTTPRequest(url,
                          method="POST",
                          headers={"Accept": "application/json"},
                          body=''  # Body is required for a POST...
                          )

        resp = yield http_client.fetch(req)
        resp_json = json.loads(resp.body.decode('utf8', 'replace'))

        access_token = resp_json['access_token']

        # Determine who the logged in user is
        req = HTTPRequest("https://%s/user" % GITHUB_API,
                          method="GET",
                          headers=_api_headers(access_token)
                          )
        resp = yield http_client.fetch(req)
        resp_json = json.loads(resp.body.decode('utf8', 'replace'))

        username = resp_json["login"]
        # username is now the GitHub userid.
        if not username:
            return None
        # Check if user is a member of any whitelisted organizations.
        # This check is performed here, as it requires `access_token`.
        if self.github_organization_whitelist:
            for org in self.github_organization_whitelist:
                user_in_org = yield self._check_organization_whitelist(org, username, access_token)
                if user_in_org:
                    break
            else:  # User not found in member list for any organisation
                return None
        userdict = {"name": username}
        # Now we set up auth_state
        auth_state = {}
        auth_state["username"] = username
        # We may want to do user provisioning in the Lab/Notebook environment.
        #  This next bit is about that.
        #  1) stash the access token
        #  2) use the GitHub ID as the uid
        #  3) set up name/email for .gitconfig
        # Store the resulting structure in auth_state
        #
        # Once you have the access token (and the appropriate scope set in the
        #  handler), you can use that in your subclassed authenticator to
        #  do nifty tricks by making more API calls to get more information
        #  to pass to the spawned Notebook/Lab.
        auth_state["access_token"] = access_token
        auth_state["uid"] = resp_json["id"]
        auth_state["name"] = resp_json["name"]
        # A public email will return in the initial query (assuming default
        #  scope).  Private will not.
        if "email" in resp_json and resp_json["email"]:
            auth_state["email"] = resp_json["email"]
        userdict["auth_state"] = auth_state
        return userdict

    @gen.coroutine
    def _check_organization_whitelist(self, org, username, access_token):
        http_client = AsyncHTTPClient()
        headers = _api_headers(access_token)
        # Get all the members for organization 'org'
        # With empty scope (even if authenticated by an org member), this
        #  will only yield public org members.  You want 'read:org' in order
        #  to be able to iterate through all members.
        next_page = "https://%s/orgs/%s/members" % (GITHUB_API, org)
        while next_page:
            req = HTTPRequest(next_page, method="GET", headers=headers)
            resp = yield http_client.fetch(req)
            resp_json = json.loads(resp.body.decode('utf8', 'replace'))
            next_page = next_page_from_links(resp)
            org_members = set(entry["login"] for entry in resp_json)
            # check if any of the organizations seen so far are in whitelist
            if username in org_members:
                return True
        return False

class LocalGitHubOAuthenticator(LocalAuthenticator, GitHubOAuthenticator):

    """A version that mixes in local system user creation"""
    pass
<|endoftext|>


<commit_before>__author__ = 'chris'

import ast
import json
import os
import time
from binascii import unhexlify
from random import shuffle

import bleach
import nacl.encoding
import nacl.signing
from twisted.internet.protocol import Protocol, Factory, connectionDone
from txws import WebSocketProtocol, WebSocketFactory

from api.utils import smart_unicode
from config import DATA_FOLDER, str_to_bool
from dht.node import Node
from keys.keychain import KeyChain
from log import Logger
from market.profile import Profile
from protos import objects
from protos.countries import CountryCode
from protos.objects import PlaintextMessage, Value, Listings

ALLOWED_TAGS = ('h2', 'h3', 'h4', 'h5', 'h6', 'p', 'a', 'u', 'ul', 'ol', 'nl', 'li', 'b', 'i', 'strong',
                'em', 'strike', 'hr', 'br', 'img', 'blockquote', 'span')

# pylint: disable=W0232
class WSProtocol(Protocol):
    """
    Handles new incoming requests coming from a websocket.
    """

    def __init__(self):
        self.log = Logger(system=self)

    def connectionMade(self):
        self.factory.register(self)

    def connectionLost(self, reason=connectionDone):
        self.factory.unregister(self)

    def get_vendors(self, message_id):
        if message_id in self.factory.outstanding_vendors:
            queried = self.factory.outstanding_vendors[message_id]
        else:
            queried = []
            self.factory.outstanding_vendors = {}
            self.factory.outstanding_vendors[message_id] = queried

        vendors = self.factory.mserver.protocol.multiplexer.vendors.values()

        shuffle(vendors)
        to_query = []
        for vendor in vendors:
            if vendor.id not in queried:
                to_query.append(vendor)

        def handle_response(metadata, node):
            to_query.remove(node)
            if metadata is not None:
                vendor = {
                    "id": message_id,
                    "vendor":
                        {
                            "guid": node.id.encode("hex"),
                            "name": metadata.name,
                            "short_description": metadata.short_description,
                            "handle": metadata.handle,
                            "avatar_hash": metadata.avatar_hash.encode("hex"),
                            "nsfw": metadata.nsfw
                        }
                }
                self.transport.write(str(bleach.clean(json.dumps(vendor, indent=4), tags=ALLOWED_TAGS)))
                queried.append(node.id)
                return True
            else:
                if node.id in self.factory.mserver.protocol.multiplexer.vendors:
                    del self.factory.mserver.protocol.multiplexer.vendors[node.id]
                self.factory.db.vendors.delete_vendor(node.id.encode("hex"))
                return False

        for node in to_query[:30]:
            self.factory.mserver.get_user_metadata(node).addCallback(handle_response, node)

    def get_moderators(self, message_id):

        def parse_response(moderators):
            if moderators is not None:
                current_mods = self.factory.db.settings.get()[11]
                self.factory.db.moderators.clear_all(except_guids=current_mods)

                def parse_profile(profile, node):
                    if profile is not None:
                        # TODO: should check signatures here before entering in database
                        self.factory.db.moderators.save_moderator(node.id.encode("hex"), node.pubkey,
                                                                  profile.bitcoin_key.public_key,
                                                                  profile.bitcoin_key.signature, profile.name,
                                                                  profile.avatar_hash, profile.moderation_fee,
                                                                  profile.handle, profile.short_description)
                        moderator = {
                            "id": message_id,
                            "moderator":
                                {
                                    "guid": node.id.encode("hex"),
                                    "name": profile.name,
                                    "handle": profile.handle,
                                    "short_description": profile.short_description,
                                    "avatar_hash": profile.avatar_hash.encode("hex"),
                                    "about": profile.about,
                                    "fee": profile.moderation_fee
                                }
                        }
                        self.transport.write(str(bleach.clean(json.dumps(moderator, indent=4), tags=ALLOWED_TAGS)))
                    else:
                        self.factory.db.moderators.delete_moderator(node.id)
                for mod in moderators:
                    try:
                        val = objects.Value()
                        val.ParseFromString(mod)
                        n = objects.Node()
                        n.ParseFromString(val.serializedData)
                        node_to_ask = Node(n.guid, n.nodeAddress.ip, n.nodeAddress.port, n.publicKey,
                                           None if not n.HasField("relayAddress") else
                                           (n.relayAddress.ip, n.relayAddress.port),
                                           n.natType, n.vendor)
                        if n.guid == KeyChain(self.factory.db).guid:
                            parse_profile(Profile(self.factory.db).get(), node_to_ask)
                        else:
                            self.factory.mserver.get_profile(node_to_ask)\
                                .addCallback(parse_profile, node_to_ask)
                    except Exception:
                        pass
        self.factory.kserver.get("moderators").addCallback(parse_response)

    def get_homepage_listings(self, message_id, only_following=False):
        if message_id not in self.factory.outstanding_listings:
            self.factory.outstanding_listings = {}
            self.factory.outstanding_listings[message_id] = []

        vendors = dict(self.factory.mserver.protocol.multiplexer.vendors)
        self.log.info("Fetching listings from %s vendors" % len(vendors))

        def get_following_from_vendors(vendors):
            follow_data = self.factory.mserver.db.follow.get_following()
            following_guids = []
            if follow_data is not None:
                f = objects.Following()
                f.ParseFromString(follow_data)
                for user in f.users:
                    following_guids.append(user.guid)
            vendor_list = []
            for k, v in vendors.items():
                if k in following_guids:
                    vendor_list.append(v)
            return vendor_list

        def handle_response(listings, node):
            count = 0
            if listings is not None:
                for l in listings.listing:
                    try:
                        if l.contract_hash not in self.factory.outstanding_listings[message_id]:
                            listing_json = {
                                "id": message_id,
                                "listing":
                                    {
                                        "guid": node.id.encode("hex"),
                                        "handle": listings.handle,
                                        "avatar_hash": listings.avatar_hash.encode("hex"),
                                        "title": l.title,
                                        "contract_hash": l.contract_hash.encode("hex"),
                                        "thumbnail_hash": l.thumbnail_hash.encode("hex"),
                                        "category": l.category,
                                        "price": l.price,
                                        "currency_code": l.currency_code,
                                        "nsfw": l.nsfw,
                                        "origin": str(CountryCode.Name(l.origin)),
                                        "ships_to": []
                                    }
                            }
                            for country in l.ships_to:
                                listing_json["listing"]["ships_to"].append(str(CountryCode.Name(country)))
                            if not os.path.isfile(os.path.join( \
                                    DATA_FOLDER, 'cache', l.thumbnail_hash.encode("hex"))):
                                self.factory.mserver.get_image(node, l.thumbnail_hash)
                            if not os.path.isfile(os.path.join( \
                                    DATA_FOLDER, 'cache', listings.avatar_hash.encode("hex"))):
                                self.factory.mserver.get_image(node, listings.avatar_hash)
                            self.transport.write(str(bleach.clean(
                                json.dumps(listing_json, indent=4), tags=ALLOWED_TAGS)))
                            count += 1
                            self.factory.outstanding_listings[message_id].append(l.contract_hash)
                            if count == 3:
                                break
                    except Exception:
                        pass
                if node.id in vendors:
                    del vendors[node.id]
            else:
                if node.id in vendors:
                    del vendors[node.id]
                if node.id in self.factory.mserver.protocol.multiplexer.vendors:
                    del self.factory.mserver.protocol.multiplexer.vendors[node.id]
                    self.factory.db.vendors.delete_vendor(node.id.encode("hex"))
                if only_following:
                    vendor_list = get_following_from_vendors(vendors)
                else:
                    vendor_list = vendors.values()
                if len(vendor_list) > 0:
                    shuffle(vendor_list)
                    node_to_ask = vendor_list[0]
                    if node_to_ask is not None:
                        self.factory.mserver.get_listings(node_to_ask).addCallback(handle_response, node_to_ask)
        if only_following:
            vendor_list = get_following_from_vendors(vendors)
        else:
            vendor_list = vendors.values()
        shuffle(vendor_list)
        for vendor in vendor_list[:15]:
            self.factory.mserver.get_listings(vendor).addCallback(handle_response, vendor)

    def send_message(self, message_id, guid, handle, message, subject, message_type, recipient_key):

        enc_key = nacl.signing.VerifyKey(unhexlify(recipient_key)).to_curve25519_public_key().encode()

        self.factory.db.messages.save_message(guid, handle, unhexlify(recipient_key), subject,
                                              message_type.upper(), message, time.time(), "", "", True,
                                              message_id)

        def send(node_to_send):
            n = node_to_send if node_to_send is not None else Node(unhexlify(guid))
            self.factory.mserver.send_message(n, enc_key,
                                              PlaintextMessage.Type.Value(message_type.upper()),
                                              message, subject,
                                              store_only=True if node_to_send is None else False)
        self.factory.kserver.resolve(unhexlify(guid)).addCallback(send)

    def search(self, message_id, keyword):
        def respond(l, node):
            if l is not None:
                listing_json = {
                    "id": message_id,
                    "listing":
                        {
                            "guid": node.id.encode("hex"),
                            "title": l.title,
                            "contract_hash": l.contract_hash.encode("hex"),
                            "thumbnail_hash": l.thumbnail_hash.encode("hex"),
                            "category": l.category,
                            "price": l.price,
                            "currency_code": l.currency_code,
                            "nsfw": l.nsfw,
                            "origin": str(CountryCode.Name(l.origin)),
                            "ships_to": [],
                            "avatar_hash": l.avatar_hash.encode("hex"),
                            "handle": l.handle
                        }
                }
                for country in l.ships_to:
                    listing_json["listing"]["ships_to"].append(str(CountryCode.Name(country)))
                self.transport.write(str(bleach.clean(json.dumps(listing_json, indent=4), tags=ALLOWED_TAGS)))

        def parse_results(values):
            if values is not None:
                for v in values:
                    try:
                        val = Value()
                        val.ParseFromString(v)
                        n = objects.Node()
                        n.ParseFromString(val.serializedData)
                        node_to_ask = Node(n.guid, n.nodeAddress.ip, n.nodeAddress.port, n.publicKey,
                                           None if not n.HasField("relayAddress") else
                                           (n.relayAddress.ip, n.relayAddress.port),
                                           n.natType, n.vendor)
                        if n.guid == KeyChain(self.factory.db).guid:
                            proto = self.factory.db.listings.get_proto()
                            l = Listings()
                            l.ParseFromString(proto)
                            for listing in l.listing:
                                if listing.contract_hash == val.valueKey:
                                    respond(listing, node_to_ask)
                        else:
                            self.factory.mserver.get_contract_metadata(node_to_ask, val.valueKey)\
                                .addCallback(respond, node_to_ask)
                    except Exception:
                        pass
        self.factory.kserver.get(keyword.lower()).addCallback(parse_results)

    def dataReceived(self, payload):
        try:
            request_json = json.loads(payload)
            if isinstance(request_json, unicode):
                payload = ast.literal_eval(payload)
                request_json = json.loads(payload)

            message_id = str(request_json["request"]["id"])

            if request_json["request"]["command"] == "get_vendors":
                self.get_vendors(message_id)

            if request_json["request"]["command"] == "get_moderators":
                self.get_moderators(message_id)

            elif request_json["request"]["command"] == "get_homepage_listings":
                self.get_homepage_listings(message_id,
                                           str_to_bool(request_json["request"]["only_following"])
                                           if "only_following" in request_json["request"] else False)

            elif request_json["request"]["command"] == "search":
                self.search(message_id, request_json["request"]["keyword"].lower())

            elif request_json["request"]["command"] == "send_message":
                self.send_message(message_id, request_json["request"]["guid"],
                                  request_json["request"]["handle"],
                                  smart_unicode(request_json["request"]["message"]),
                                  request_json["request"]["subject"],
                                  request_json["request"]["message_type"],
                                  request_json["request"]["public_key"])

        except Exception as e:
            print 'Exception occurred: %s' % e

class WSFactory(Factory):

    def __init__(self, mserver, kserver, only_ip=None):
        if only_ip == None:
            only_ip = ["127.0.0.1"]
        self.mserver = mserver
        self.kserver = kserver
        self.db = mserver.db
        self.outstanding_listings = {}
        self.outstanding_vendors = {}
        self.protocol = WSProtocol
        self.only_ip = only_ip
        self.clients = []

    def buildProtocol(self, addr):
        if addr.host not in self.only_ip and "0.0.0.0" not in self.only_ip:
            return
        return Factory.buildProtocol(self, addr)

    def register(self, client):
        if client not in self.clients:
            self.clients.append(client)

    def unregister(self, client):
        if client in self.clients:
            self.clients.remove(client)

    def push(self, msg):
        for c in self.clients:
            c.transport.write(msg)

class AuthenticatedWebSocketProtocol(WebSocketProtocol):

    def validateHeaders(self):
        if "Cookie" in self.headers:
            for session in self.factory.authenticated_sessions:
                if "TWISTED_SESSION=" + session.uid in self.headers["Cookie"]:
                    return WebSocketProtocol.validateHeaders(self)
        return False

class AuthenticatedWebSocketFactory(WebSocketFactory):

    authenticated_sessions = None
<commit_msg>Unmarshal json moderators<commit_after>__author__ = 'chris'

import ast
import json
import os
import time
from binascii import unhexlify
from random import shuffle

import bleach
import nacl.encoding
import nacl.signing
from twisted.internet.protocol import Protocol, Factory, connectionDone
from txws import WebSocketProtocol, WebSocketFactory

from api.utils import smart_unicode
from config import DATA_FOLDER, str_to_bool
from dht.node import Node
from keys.keychain import KeyChain
from log import Logger
from market.profile import Profile
from protos import objects
from protos.countries import CountryCode
from protos.objects import PlaintextMessage, Value, Listings

ALLOWED_TAGS = ('h2', 'h3', 'h4', 'h5', 'h6', 'p', 'a', 'u', 'ul', 'ol', 'nl', 'li', 'b', 'i', 'strong',
                'em', 'strike', 'hr', 'br', 'img', 'blockquote', 'span')

# pylint: disable=W0232
class WSProtocol(Protocol):
    """
    Handles new incoming requests coming from a websocket.
    """

    def __init__(self):
        self.log = Logger(system=self)

    def connectionMade(self):
        self.factory.register(self)

    def connectionLost(self, reason=connectionDone):
        self.factory.unregister(self)

    def get_vendors(self, message_id):
        if message_id in self.factory.outstanding_vendors:
            queried = self.factory.outstanding_vendors[message_id]
        else:
            queried = []
            self.factory.outstanding_vendors = {}
            self.factory.outstanding_vendors[message_id] = queried

        vendors = self.factory.mserver.protocol.multiplexer.vendors.values()

        shuffle(vendors)
        to_query = []
        for vendor in vendors:
            if vendor.id not in queried:
                to_query.append(vendor)

        def handle_response(metadata, node):
            to_query.remove(node)
            if metadata is not None:
                vendor = {
                    "id": message_id,
                    "vendor":
                        {
                            "guid": node.id.encode("hex"),
                            "name": metadata.name,
                            "short_description": metadata.short_description,
                            "handle": metadata.handle,
                            "avatar_hash": metadata.avatar_hash.encode("hex"),
                            "nsfw": metadata.nsfw
                        }
                }
                self.transport.write(str(bleach.clean(json.dumps(vendor, indent=4), tags=ALLOWED_TAGS)))
                queried.append(node.id)
                return True
            else:
                if node.id in self.factory.mserver.protocol.multiplexer.vendors:
                    del self.factory.mserver.protocol.multiplexer.vendors[node.id]
                self.factory.db.vendors.delete_vendor(node.id.encode("hex"))
                return False

        for node in to_query[:30]:
            self.factory.mserver.get_user_metadata(node).addCallback(handle_response, node)

    def get_moderators(self, message_id):

        def parse_response(moderators):
            if moderators is not None:
                current_mods = json.loads(self.factory.db.settings.get()[11])
                self.factory.db.moderators.clear_all(except_guids=current_mods)

                def parse_profile(profile, node):
                    if profile is not None:
                        # TODO: should check signatures here before entering in database
                        self.factory.db.moderators.save_moderator(node.id.encode("hex"), node.pubkey,
                                                                  profile.bitcoin_key.public_key,
                                                                  profile.bitcoin_key.signature, profile.name,
                                                                  profile.avatar_hash, profile.moderation_fee,
                                                                  profile.handle, profile.short_description)
                        moderator = {
                            "id": message_id,
                            "moderator":
                                {
                                    "guid": node.id.encode("hex"),
                                    "name": profile.name,
                                    "handle": profile.handle,
                                    "short_description": profile.short_description,
                                    "avatar_hash": profile.avatar_hash.encode("hex"),
                                    "about": profile.about,
                                    "fee": profile.moderation_fee
                                }
                        }
                        self.transport.write(str(bleach.clean(json.dumps(moderator, indent=4), tags=ALLOWED_TAGS)))
                    else:
                        self.factory.db.moderators.delete_moderator(node.id)
                for mod in moderators:
                    try:
                        val = objects.Value()
                        val.ParseFromString(mod)
                        n = objects.Node()
                        n.ParseFromString(val.serializedData)
                        node_to_ask = Node(n.guid, n.nodeAddress.ip, n.nodeAddress.port, n.publicKey,
                                           None if not n.HasField("relayAddress") else
                                           (n.relayAddress.ip, n.relayAddress.port),
                                           n.natType, n.vendor)
                        if n.guid == KeyChain(self.factory.db).guid:
                            parse_profile(Profile(self.factory.db).get(), node_to_ask)
                        else:
                            self.factory.mserver.get_profile(node_to_ask)\
                                .addCallback(parse_profile, node_to_ask)
                    except Exception:
                        pass
        self.factory.kserver.get("moderators").addCallback(parse_response)

    def get_homepage_listings(self, message_id, only_following=False):
        if message_id not in self.factory.outstanding_listings:
            self.factory.outstanding_listings = {}
            self.factory.outstanding_listings[message_id] = []

        vendors = dict(self.factory.mserver.protocol.multiplexer.vendors)
        self.log.info("Fetching listings from %s vendors" % len(vendors))

        def get_following_from_vendors(vendors):
            follow_data = self.factory.mserver.db.follow.get_following()
            following_guids = []
            if follow_data is not None:
                f = objects.Following()
                f.ParseFromString(follow_data)
                for user in f.users:
                    following_guids.append(user.guid)
            vendor_list = []
            for k, v in vendors.items():
                if k in following_guids:
                    vendor_list.append(v)
            return vendor_list

        def handle_response(listings, node):
            count = 0
            if listings is not None:
                for l in listings.listing:
                    try:
                        if l.contract_hash not in self.factory.outstanding_listings[message_id]:
                            listing_json = {
                                "id": message_id,
                                "listing":
                                    {
                                        "guid": node.id.encode("hex"),
                                        "handle": listings.handle,
                                        "avatar_hash": listings.avatar_hash.encode("hex"),
                                        "title": l.title,
                                        "contract_hash": l.contract_hash.encode("hex"),
                                        "thumbnail_hash": l.thumbnail_hash.encode("hex"),
                                        "category": l.category,
                                        "price": l.price,
                                        "currency_code": l.currency_code,
                                        "nsfw": l.nsfw,
                                        "origin": str(CountryCode.Name(l.origin)),
                                        "ships_to": []
                                    }
                            }
                            for country in l.ships_to:
                                listing_json["listing"]["ships_to"].append(str(CountryCode.Name(country)))
                            if not os.path.isfile(os.path.join( \
                                    DATA_FOLDER, 'cache', l.thumbnail_hash.encode("hex"))):
                                self.factory.mserver.get_image(node, l.thumbnail_hash)
                            if not os.path.isfile(os.path.join( \
                                    DATA_FOLDER, 'cache', listings.avatar_hash.encode("hex"))):
                                self.factory.mserver.get_image(node, listings.avatar_hash)
                            self.transport.write(str(bleach.clean(
                                json.dumps(listing_json, indent=4), tags=ALLOWED_TAGS)))
                            count += 1
                            self.factory.outstanding_listings[message_id].append(l.contract_hash)
                            if count == 3:
                                break
                    except Exception:
                        pass
                if node.id in vendors:
                    del vendors[node.id]
            else:
                if node.id in vendors:
                    del vendors[node.id]
                if node.id in self.factory.mserver.protocol.multiplexer.vendors:
                    del self.factory.mserver.protocol.multiplexer.vendors[node.id]
                    self.factory.db.vendors.delete_vendor(node.id.encode("hex"))
                if only_following:
                    vendor_list = get_following_from_vendors(vendors)
                else:
                    vendor_list = vendors.values()
                if len(vendor_list) > 0:
                    shuffle(vendor_list)
                    node_to_ask = vendor_list[0]
                    if node_to_ask is not None:
                        self.factory.mserver.get_listings(node_to_ask).addCallback(handle_response, node_to_ask)
        if only_following:
            vendor_list = get_following_from_vendors(vendors)
        else:
            vendor_list = vendors.values()
        shuffle(vendor_list)
        for vendor in vendor_list[:15]:
            self.factory.mserver.get_listings(vendor).addCallback(handle_response, vendor)

    def send_message(self, message_id, guid, handle, message, subject, message_type, recipient_key):

        enc_key = nacl.signing.VerifyKey(unhexlify(recipient_key)).to_curve25519_public_key().encode()

        self.factory.db.messages.save_message(guid, handle, unhexlify(recipient_key), subject,
                                              message_type.upper(), message, time.time(), "", "", True,
                                              message_id)

        def send(node_to_send):
            n = node_to_send if node_to_send is not None else Node(unhexlify(guid))
            self.factory.mserver.send_message(n, enc_key,
                                              PlaintextMessage.Type.Value(message_type.upper()),
                                              message, subject,
                                              store_only=True if node_to_send is None else False)
        self.factory.kserver.resolve(unhexlify(guid)).addCallback(send)

    def search(self, message_id, keyword):
        def respond(l, node):
            if l is not None:
                listing_json = {
                    "id": message_id,
                    "listing":
                        {
                            "guid": node.id.encode("hex"),
                            "title": l.title,
                            "contract_hash": l.contract_hash.encode("hex"),
                            "thumbnail_hash": l.thumbnail_hash.encode("hex"),
                            "category": l.category,
                            "price": l.price,
                            "currency_code": l.currency_code,
                            "nsfw": l.nsfw,
                            "origin": str(CountryCode.Name(l.origin)),
                            "ships_to": [],
                            "avatar_hash": l.avatar_hash.encode("hex"),
                            "handle": l.handle
                        }
                }
                for country in l.ships_to:
                    listing_json["listing"]["ships_to"].append(str(CountryCode.Name(country)))
                self.transport.write(str(bleach.clean(json.dumps(listing_json, indent=4), tags=ALLOWED_TAGS)))

        def parse_results(values):
            if values is not None:
                for v in values:
                    try:
                        val = Value()
                        val.ParseFromString(v)
                        n = objects.Node()
                        n.ParseFromString(val.serializedData)
                        node_to_ask = Node(n.guid, n.nodeAddress.ip, n.nodeAddress.port, n.publicKey,
                                           None if not n.HasField("relayAddress") else
                                           (n.relayAddress.ip, n.relayAddress.port),
                                           n.natType, n.vendor)
                        if n.guid == KeyChain(self.factory.db).guid:
                            proto = self.factory.db.listings.get_proto()
                            l = Listings()
                            l.ParseFromString(proto)
                            for listing in l.listing:
                                if listing.contract_hash == val.valueKey:
                                    respond(listing, node_to_ask)
                        else:
                            self.factory.mserver.get_contract_metadata(node_to_ask, val.valueKey)\
                                .addCallback(respond, node_to_ask)
                    except Exception:
                        pass
        self.factory.kserver.get(keyword.lower()).addCallback(parse_results)

    def dataReceived(self, payload):
        try:
            request_json = json.loads(payload)
            if isinstance(request_json, unicode):
                payload = ast.literal_eval(payload)
                request_json = json.loads(payload)

            message_id = str(request_json["request"]["id"])

            if request_json["request"]["command"] == "get_vendors":
                self.get_vendors(message_id)

            if request_json["request"]["command"] == "get_moderators":
                self.get_moderators(message_id)

            elif request_json["request"]["command"] == "get_homepage_listings":
                self.get_homepage_listings(message_id,
                                           str_to_bool(request_json["request"]["only_following"])
                                           if "only_following" in request_json["request"] else False)

            elif request_json["request"]["command"] == "search":
                self.search(message_id, request_json["request"]["keyword"].lower())

            elif request_json["request"]["command"] == "send_message":
                self.send_message(message_id, request_json["request"]["guid"],
                                  request_json["request"]["handle"],
                                  smart_unicode(request_json["request"]["message"]),
                                  request_json["request"]["subject"],
                                  request_json["request"]["message_type"],
                                  request_json["request"]["public_key"])

        except Exception as e:
            print 'Exception occurred: %s' % e

class WSFactory(Factory):

    def __init__(self, mserver, kserver, only_ip=None):
        if only_ip == None:
            only_ip = ["127.0.0.1"]
        self.mserver = mserver
        self.kserver = kserver
        self.db = mserver.db
        self.outstanding_listings = {}
        self.outstanding_vendors = {}
        self.protocol = WSProtocol
        self.only_ip = only_ip
        self.clients = []

    def buildProtocol(self, addr):
        if addr.host not in self.only_ip and "0.0.0.0" not in self.only_ip:
            return
        return Factory.buildProtocol(self, addr)

    def register(self, client):
        if client not in self.clients:
            self.clients.append(client)

    def unregister(self, client):
        if client in self.clients:
            self.clients.remove(client)

    def push(self, msg):
        for c in self.clients:
            c.transport.write(msg)

class AuthenticatedWebSocketProtocol(WebSocketProtocol):

    def validateHeaders(self):
        if "Cookie" in self.headers:
            for session in self.factory.authenticated_sessions:
                if "TWISTED_SESSION=" + session.uid in self.headers["Cookie"]:
                    return WebSocketProtocol.validateHeaders(self)
        return False

class AuthenticatedWebSocketFactory(WebSocketFactory):

    authenticated_sessions = None
<|endoftext|>