<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals
import uuid

from django.test import SimpleTestCase
from elasticsearch.exceptions import ConnectionError
from mock import patch, MagicMock

from corehq.apps.reports.analytics.esaccessors import get_user_stubs
from corehq.elastic import doc_exists_in_es, get_es_new
from corehq.pillows.mappings.user_mapping import USER_INDEX_INFO
from corehq.util.test_utils import trap_extra_setup, mock_out_couch
from dimagi.utils.couch.undo import DELETED_SUFFIX

from ..models import CommCareUser, WebUser

# Note that you can't directly patch the signal handler, as that code has
# already been called.  It's easier to patch something that the handler calls.

# Also, you need to patch the path to the function in the file where the signal
# handler uses it, not where it's actually defined.  That's quite a gotcha.
from pillowtop.es_utils import initialize_index_and_mapping

@mock_out_couch()
@patch('corehq.apps.users.models.CouchUser.sync_to_django_user', new=MagicMock)
class TestUserSignals(SimpleTestCase):

    @patch('corehq.apps.analytics.signals.update_hubspot_properties_v2')
    @patch('corehq.apps.callcenter.signals.sync_call_center_user_case')
    @patch('corehq.apps.cachehq.signals.invalidate_document')
    @patch('corehq.apps.users.signals.send_to_elasticsearch')
    def test_commcareuser_save(self, send_to_es, invalidate, sync_call_center,
                               update_hubspot_properties_v2):
        CommCareUser(username='test').save()

        self.assertTrue(send_to_es.called)
        self.assertTrue(invalidate.called)
        self.assertTrue(sync_call_center.called)
        self.assertFalse(update_hubspot_properties_v2.called)

    @patch('corehq.apps.analytics.signals.update_hubspot_properties_v2')
    @patch('corehq.apps.callcenter.signals.sync_call_center_user_case')
    @patch('corehq.apps.cachehq.signals.invalidate_document')
    @patch('corehq.apps.users.signals.send_to_elasticsearch')
    def test_webuser_save(self, send_to_es, invalidate, sync_call_center,
                          update_hubspot_properties_v2):
        WebUser().save()

        self.assertTrue(send_to_es.called)
        self.assertTrue(invalidate.called)
        self.assertFalse(sync_call_center.called)
        self.assertTrue(update_hubspot_properties_v2.called)

@mock_out_couch()
@patch('corehq.apps.users.models.CouchUser.sync_to_django_user', new=MagicMock)
@patch('corehq.apps.analytics.signals.update_hubspot_properties_v2')
@patch('corehq.apps.callcenter.signals.sync_call_center_user_case')
@patch('corehq.apps.cachehq.signals.invalidate_document')
class TestUserSyncToEs(SimpleTestCase):

    @classmethod
    def setUpClass(cls):
        super(TestUserSyncToEs, cls).setUpClass()

        # create the index
        cls.es = get_es_new()
        with trap_extra_setup(ConnectionError):
            initialize_index_and_mapping(cls.es, USER_INDEX_INFO)

    def test_sync_to_es_create_update_delete(self, *mocks):
        domain = 'user_es_domain'
        user = CommCareUser(
            domain=domain,
            username='user1',
            _id=uuid.uuid4().hex,
            is_active=True,
            first_name='user1 first name',
            last_name='user1 last name',
            location_id='location1'
        )
        user.save()

        self.check_user(user)

        user.first_name = 'new first name'
        user.save()
        self.check_user(user)

        # simulate retire without needing couch
        user.base_doc += DELETED_SUFFIX
        user.save()
        self.es.indices.refresh(USER_INDEX_INFO.index)
        self.assertFalse(doc_exists_in_es(USER_INDEX_INFO, user._id))

    def check_user(self, user):
        self.es.indices.refresh(USER_INDEX_INFO.index)
        results = get_user_stubs([user._id])
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0], {
            '_id': user._id,
            'username': user.username,
            'is_active': True,
            'first_name': user.first_name,
            'last_name': user.last_name,
            'doc_type': user.doc_type,
            'location_id': 'location1',
            '__group_ids': []
        })
<commit_msg>fix tests<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals
import uuid

from django.test import SimpleTestCase
from elasticsearch.exceptions import ConnectionError
from mock import patch, MagicMock

from corehq.apps.reports.analytics.esaccessors import get_user_stubs
from corehq.elastic import doc_exists_in_es, get_es_new
from corehq.pillows.mappings.user_mapping import USER_INDEX_INFO
from corehq.util.test_utils import trap_extra_setup, mock_out_couch
from dimagi.utils.couch.undo import DELETED_SUFFIX

from ..models import CommCareUser, WebUser

# Note that you can't directly patch the signal handler, as that code has
# already been called.  It's easier to patch something that the handler calls.

# Also, you need to patch the path to the function in the file where the signal
# handler uses it, not where it's actually defined.  That's quite a gotcha.
from pillowtop.es_utils import initialize_index_and_mapping

@mock_out_couch()
@patch('corehq.apps.users.models.CouchUser.sync_to_django_user', new=MagicMock)
class TestUserSignals(SimpleTestCase):

    @patch('corehq.apps.analytics.signals.update_hubspot_properties_v2.delay')
    @patch('corehq.apps.callcenter.signals.sync_call_center_user_case')
    @patch('corehq.apps.cachehq.signals.invalidate_document')
    @patch('corehq.apps.users.signals.send_to_elasticsearch')
    def test_commcareuser_save(self, send_to_es, invalidate, sync_call_center,
                               update_hubspot_properties_v2):
        CommCareUser(username='test').save()

        self.assertTrue(send_to_es.called)
        self.assertTrue(invalidate.called)
        self.assertTrue(sync_call_center.called)
        self.assertFalse(update_hubspot_properties_v2.called)

    @patch('corehq.apps.analytics.signals.update_hubspot_properties_v2.delay')
    @patch('corehq.apps.callcenter.signals.sync_call_center_user_case')
    @patch('corehq.apps.cachehq.signals.invalidate_document')
    @patch('corehq.apps.users.signals.send_to_elasticsearch')
    def test_webuser_save(self, send_to_es, invalidate, sync_call_center,
                          update_hubspot_properties_v2):
        WebUser().save()

        self.assertTrue(send_to_es.called)
        self.assertTrue(invalidate.called)
        self.assertFalse(sync_call_center.called)
        self.assertTrue(update_hubspot_properties_v2.called)

@mock_out_couch()
@patch('corehq.apps.users.models.CouchUser.sync_to_django_user', new=MagicMock)
@patch('corehq.apps.analytics.signals.update_hubspot_properties_v2')
@patch('corehq.apps.callcenter.signals.sync_call_center_user_case')
@patch('corehq.apps.cachehq.signals.invalidate_document')
class TestUserSyncToEs(SimpleTestCase):

    @classmethod
    def setUpClass(cls):
        super(TestUserSyncToEs, cls).setUpClass()

        # create the index
        cls.es = get_es_new()
        with trap_extra_setup(ConnectionError):
            initialize_index_and_mapping(cls.es, USER_INDEX_INFO)

    def test_sync_to_es_create_update_delete(self, *mocks):
        domain = 'user_es_domain'
        user = CommCareUser(
            domain=domain,
            username='user1',
            _id=uuid.uuid4().hex,
            is_active=True,
            first_name='user1 first name',
            last_name='user1 last name',
            location_id='location1'
        )
        user.save()

        self.check_user(user)

        user.first_name = 'new first name'
        user.save()
        self.check_user(user)

        # simulate retire without needing couch
        user.base_doc += DELETED_SUFFIX
        user.save()
        self.es.indices.refresh(USER_INDEX_INFO.index)
        self.assertFalse(doc_exists_in_es(USER_INDEX_INFO, user._id))

    def check_user(self, user):
        self.es.indices.refresh(USER_INDEX_INFO.index)
        results = get_user_stubs([user._id])
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0], {
            '_id': user._id,
            'username': user.username,
            'is_active': True,
            'first_name': user.first_name,
            'last_name': user.last_name,
            'doc_type': user.doc_type,
            'location_id': 'location1',
            '__group_ids': []
        })
<|endoftext|>


<commit_before>from celery.schedules import crontab
from celery.task import periodic_task
from django.conf import settings

from corehq.util.datadog.gauges import datadog_gauge
from corehq.util.decorators import serial_task
from pillowtop.utils import get_all_pillows_json

@periodic_task(run_every=crontab(minute="*/5"), queue=settings.CELERY_PERIODIC_QUEUE)
def run_pillow_datadog_metrics():
    pillow_datadog_metrics.delay()

@serial_task('pillow-datadog-metrics', timeout=30 * 60, queue=settings.CELERY_PERIODIC_QUEUE, max_retries=0)
def pillow_datadog_metrics():
    def _is_couch(pillow):
        # text is couch, json is kafka
        return pillow['seq_format'] == 'text'

    pillow_meta = get_all_pillows_json()

    for pillow in pillow_meta:
        tags = [
            'pillow_name:{}'.format(pillow['name']),
            'feed_type:{}'.format('couch' if _is_couch(pillow) else 'kafka')
        ]

        for topic_name, offset in pillow['offsets'].items():
            if _is_couch(pillow):
                assert isinstance(pillow['seq'], int)
                assert len(pillow['offsets']) == 1
                tags_with_topic = tags + ['topic:{}'.format(topic_name)]
                processed_offset = pillow['seq']
            else:
                assert isinstance(pillow['seq'], dict)
                if not pillow['seq']:
                    # this pillow has never been initialized.
                    # (custom pillows on most environments)
                    continue
                assert len(pillow['offsets']) == len(pillow['seq'])
                topic, partition = topic_name.split(',')
                tags_with_topic = tags + ['topic:{}-{}'.format(topic, partition)]
                processed_offset = pillow['seq'][topic_name]

            if processed_offset == 0:
                # assume if nothing has been processed that this pillow is not
                # supposed to be running
                continue

            datadog_gauge(
                'commcare.change_feed.current_offsets',
                offset, tags=tags_with_topic
            )
            datadog_gauge(
                'commcare.change_feed.processed_offsets',
                processed_offset, tags=tags_with_topic
            )
            needs_processing = offset - processed_offset
            datadog_gauge(
                'commcare.change_feed.need_processing',
                needs_processing, tags=tags_with_topic
            )
<commit_msg>Switch from assert to soft assert<commit_after>from celery.schedules import crontab
from celery.task import periodic_task
from django.conf import settings

from corehq.util.datadog.gauges import datadog_gauge
from corehq.util.decorators import serial_task
from corehq.util.soft_assert import soft_assert
from pillowtop.utils import get_all_pillows_json

_assert = soft_assert("{}@{}".format('jemord', 'dimagi.com'))

@periodic_task(run_every=crontab(minute="*/5"), queue=settings.CELERY_PERIODIC_QUEUE)
def run_pillow_datadog_metrics():
    pillow_datadog_metrics.delay()

@serial_task('pillow-datadog-metrics', timeout=30 * 60, queue=settings.CELERY_PERIODIC_QUEUE, max_retries=0)
def pillow_datadog_metrics():
    def _is_couch(pillow):
        # text is couch, json is kafka
        return pillow['seq_format'] == 'text'

    pillow_meta = get_all_pillows_json()

    for pillow in pillow_meta:
        tags = [
            'pillow_name:{}'.format(pillow['name']),
            'feed_type:{}'.format('couch' if _is_couch(pillow) else 'kafka')
        ]

        for topic_name, offset in pillow['offsets'].items():
            if _is_couch(pillow):
                if not isinstance(pillow['seq'], int) or len(pillow['offsets']) != 1:
                    _assert(False, "Unexpected couch pillow format {}".format(pillow['name']))
                    continue
                tags_with_topic = tags + ['topic:{}'.format(topic_name)]
                processed_offset = pillow['seq']
            else:
                if not isinstance(pillow['seq'], dict) or len(pillow['offsets']) != len(pillow['seq']):
                    _assert(False, "Unexpected kafka pillow format {}".format(pillow['name']))
                    continue
                if not pillow['seq']:
                    # this pillow has never been initialized.
                    # (custom pillows on most environments)
                    continue
                topic, partition = topic_name.split(',')
                tags_with_topic = tags + ['topic:{}-{}'.format(topic, partition)]
                processed_offset = pillow['seq'][topic_name]

            if processed_offset == 0:
                # assume if nothing has been processed that this pillow is not
                # supposed to be running
                continue

            datadog_gauge(
                'commcare.change_feed.current_offsets',
                offset, tags=tags_with_topic
            )
            datadog_gauge(
                'commcare.change_feed.processed_offsets',
                processed_offset, tags=tags_with_topic
            )
            needs_processing = offset - processed_offset
            datadog_gauge(
                'commcare.change_feed.need_processing',
                needs_processing, tags=tags_with_topic
            )
<|endoftext|>


<commit_before>from collections import namedtuple
from copy import deepcopy
from datetime import datetime
import sys

import json
from django.conf import settings
from kafka import KafkaConsumer

from corehq.util.json import CommCareJSONEncoder
from dimagi.utils.chunked import chunked
from dimagi.utils.modules import to_function

from pillowtop.exceptions import PillowNotFoundError
from pillowtop.logger import pillow_logging
from pillowtop.dao.exceptions import DocumentMismatchError, DocumentMissingError

def _get_pillow_instance(full_class_str):
    pillow_class = _import_class_or_function(full_class_str)
    if pillow_class is None:
        raise ValueError('No pillow class found for {}'.format(full_class_str))
    return pillow_class()

def _import_class_or_function(full_class_str):
    return to_function(full_class_str, failhard=settings.DEBUG)

def get_all_pillow_classes():
    return [config.get_class() for config in get_all_pillow_configs()]

def get_all_pillow_instances():
    return [config.get_instance() for config in get_all_pillow_configs()]

def get_couch_pillow_instances():
    from pillowtop.feed.couch import CouchChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), CouchChangeFeed)
    ]

def get_kafka_pillow_instances():
    from corehq.apps.change_feed.consumer.feed import KafkaChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), KafkaChangeFeed)
    ]

def get_all_pillow_configs():
    return get_pillow_configs_from_settings_dict(getattr(settings, 'PILLOWTOPS', {}))

def get_pillow_configs_from_settings_dict(pillow_settings_dict):
    """
    The pillow_settings_dict is expected to be a dict mapping groups to list of pillow configs
    """
    for section, list_of_pillows in pillow_settings_dict.items():
        for pillow_config in list_of_pillows:
            yield get_pillow_config_from_setting(section, pillow_config)

class PillowConfig(namedtuple('PillowConfig', ['section', 'name', 'class_name', 'instance_generator', 'params'])):
    """
    Helper object for getting pillow classes/instances from settings
    """

    def __hash__(self):
        return hash(self.name)

    def get_class(self):
        return _import_class_or_function(self.class_name)

    def get_instance(self, **kwargs):
        if self.instance_generator:
            instance_generator_fn = _import_class_or_function(self.instance_generator)
            params = deepcopy(self.params)  # parameters defined in settings file
            params.update(kwargs)  # parameters passed in via run_ptop
            return instance_generator_fn(self.name, **params)
        else:
            return _get_pillow_instance(self.class_name)

def get_pillow_config_from_setting(section, pillow_config_string_or_dict):
    if isinstance(pillow_config_string_or_dict, str):
        return PillowConfig(
            section,
            pillow_config_string_or_dict.rsplit('.', 1)[1],
            pillow_config_string_or_dict,
            None,
            {}
        )
    else:
        assert 'class' in pillow_config_string_or_dict
        class_name = pillow_config_string_or_dict['class']
        return PillowConfig(
            section,
            pillow_config_string_or_dict.get('name', class_name),
            class_name,
            pillow_config_string_or_dict.get('instance', None),
            pillow_config_string_or_dict.get('params', {}),
        )

def get_pillow_by_name(pillow_class_name, instantiate=True, **kwargs):
    config = get_pillow_config_by_name(pillow_class_name)
    return config.get_instance(**kwargs) if instantiate else config.get_class()

def get_pillow_config_by_name(pillow_name):
    all_configs = get_all_pillow_configs()
    for config in all_configs:
        if config.name == pillow_name:
            return config
    raise PillowNotFoundError('No pillow found with name {}'.format(pillow_name))

def force_seq_int(seq):
    if seq is None or seq == '':
        return None
    elif isinstance(seq, dict):
        # multi-topic checkpoints don't support a single sequence id
        return None
    elif isinstance(seq, str):
        return int(seq.split('-')[0])
    else:
        assert isinstance(seq, int), type(seq)
        return seq

def safe_force_seq_int(seq, default=None):
    if isinstance(seq, dict):
        return default
    try:
        return force_seq_int(seq)
    except (AssertionError, ValueError):
        return default

def _get_consumer():
    return KafkaConsumer(
        client_id='pillowtop_utils',
        bootstrap_servers=settings.KAFKA_BROKERS,
        request_timeout_ms=1000
    )

def get_all_pillows_json():
    pillow_configs = get_all_pillow_configs()
    consumer = _get_consumer()
    return [get_pillow_json(pillow_config, consumer) for pillow_config in pillow_configs]

def get_pillow_json(pillow_config, consumer=None):
    assert isinstance(pillow_config, PillowConfig)

    def _is_kafka(checkpoint):
        return checkpoint.sequence_format == 'json'

    pillow = pillow_config.get_instance()
    if consumer is None:
        consumer = _get_consumer()

    checkpoint = pillow.get_checkpoint()
    timestamp = checkpoint.timestamp
    if timestamp:
        time_since_last = datetime.utcnow() - timestamp
        seconds_since_last = time_since_last.total_seconds()
        hours_since_last = seconds_since_last // 3600

        try:
            # remove microsecond portion
            prettified_time_since_last = str(time_since_last)
            prettified_time_since_last = prettified_time_since_last[0:prettified_time_since_last.index('.')]
        except ValueError:
            pass
    else:
        seconds_since_last = 0
        prettified_time_since_last = ''
        hours_since_last = None

    if _is_kafka(checkpoint):
        # breaking composition boundaries so that we can use one Kafka connection
        offsets = {
            "{},{}".format(tp.topic, tp.partition): offset
            for tp, offset in consumer.end_offsets(list(checkpoint.wrapped_sequence.keys())).items()
        }
    else:
        offsets = pillow.get_change_feed().get_latest_offsets_json()

    def _seq_to_int(checkpoint, seq):
        from pillowtop.models import kafka_seq_to_str
        if _is_kafka(checkpoint):
            return json.loads(kafka_seq_to_str(seq))
        else:
            return force_seq_int(seq)

    return {
        'name': pillow_config.name,
        'seq_format': checkpoint.sequence_format,
        'seq': _seq_to_int(checkpoint, checkpoint.wrapped_sequence),
        'offsets': offsets,
        'seconds_since_last': seconds_since_last,
        'time_since_last': prettified_time_since_last,
        'hours_since_last': hours_since_last
    }

ChangeError = namedtuple('ChangeError', 'change exception')

class ErrorCollector(object):
    def __init__(self):
        self.errors = []

    def add_error(self, error):
        self.errors.append(error)

def build_bulk_payload(index_info, changes, doc_transform=None, error_collector=None):
    doc_transform = doc_transform or (lambda x: x)
    payload = []
    for change in changes:
        if change.deleted and change.id:
            payload.append({
                "delete": {
                    "_index": index_info.index,
                    "_type": index_info.type,
                    "_id": change.id
                }
            })
        elif not change.deleted:
            try:
                doc = change.get_document()
                doc = doc_transform(doc)
                payload.append({
                    "index": {
                        "_index": index_info.index,
                        "_type": index_info.type,
                        "_id": doc['_id']
                    }
                })
                payload.append(doc)
            except Exception as e:
                if not error_collector:
                    raise
                error_collector.add_error(ChangeError(change, e))
    return payload

def prepare_bulk_payloads(bulk_changes, max_size, chunk_size=100):
    payloads = [b'']
    for bulk_chunk in chunked(bulk_changes, chunk_size):
        current_payload = payloads[-1]
        json_bulk_chunks = [
            json.dumps(obj, cls=CommCareJSONEncoder).encode('utf-8')
            for obj in bulk_chunk
        ]
        payload_chunk = b'\n'.join(json_bulk_chunks) + b'\n'
        appended_payload = current_payload + payload_chunk
        new_payload_size = sys.getsizeof(appended_payload)
        if new_payload_size > max_size:
            payloads.append(payload_chunk)
        else:
            payloads[-1] = appended_payload

    return [_f for _f in payloads if _f]

def ensure_matched_revisions(change, fetched_document):
    """
    This function ensures that the document fetched from a change matches the
    revision at which it was pushed to kafka at.

    See http://manage.dimagi.com/default.asp?237983 for more details

    :raises: DocumentMismatchError - Raised when the revisions of the fetched document
        and the change metadata do not match
    """

    change_has_rev = change.metadata and change.metadata.document_rev is not None
    doc_has_rev = fetched_document and '_rev' in fetched_document
    if doc_has_rev and change_has_rev:

        doc_rev = fetched_document['_rev']
        change_rev = change.metadata.document_rev
        if doc_rev != change_rev:
            fetched_rev = _convert_rev_to_int(doc_rev)
            stored_rev = _convert_rev_to_int(change_rev)
            if fetched_rev < stored_rev or stored_rev == -1:
                message = "Mismatched revs for {}: Cloudant rev {} vs. Changes feed rev {}".format(
                    change.id,
                    doc_rev,
                    change_rev
                )
                pillow_logging.warning(message)
                raise DocumentMismatchError(message)

def _convert_rev_to_int(rev):
    try:
        return int(rev.split('-')[0])
    except (ValueError, AttributeError):
        return -1

def ensure_document_exists(change):
    """
    This is to ensure that the Couch document exists in the Couch database when the
    change is processed. We only care about the scenario where the document is missing.

    :raises: DocumentMissingError - Raised when the document is missing (never existed)
    """
    change.get_document()
    if change.error_raised is not None and isinstance(change.error_raised, DocumentMissingError):
        raise change.error_raised
<commit_msg>close consumer<commit_after>from collections import namedtuple
from copy import deepcopy
from datetime import datetime
import sys

import json
from django.conf import settings
from kafka import KafkaConsumer

from corehq.util.io import ClosingContextProxy
from corehq.util.json import CommCareJSONEncoder
from dimagi.utils.chunked import chunked
from dimagi.utils.modules import to_function

from pillowtop.exceptions import PillowNotFoundError
from pillowtop.logger import pillow_logging
from pillowtop.dao.exceptions import DocumentMismatchError, DocumentMissingError

def _get_pillow_instance(full_class_str):
    pillow_class = _import_class_or_function(full_class_str)
    if pillow_class is None:
        raise ValueError('No pillow class found for {}'.format(full_class_str))
    return pillow_class()

def _import_class_or_function(full_class_str):
    return to_function(full_class_str, failhard=settings.DEBUG)

def get_all_pillow_classes():
    return [config.get_class() for config in get_all_pillow_configs()]

def get_all_pillow_instances():
    return [config.get_instance() for config in get_all_pillow_configs()]

def get_couch_pillow_instances():
    from pillowtop.feed.couch import CouchChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), CouchChangeFeed)
    ]

def get_kafka_pillow_instances():
    from corehq.apps.change_feed.consumer.feed import KafkaChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), KafkaChangeFeed)
    ]

def get_all_pillow_configs():
    return get_pillow_configs_from_settings_dict(getattr(settings, 'PILLOWTOPS', {}))

def get_pillow_configs_from_settings_dict(pillow_settings_dict):
    """
    The pillow_settings_dict is expected to be a dict mapping groups to list of pillow configs
    """
    for section, list_of_pillows in pillow_settings_dict.items():
        for pillow_config in list_of_pillows:
            yield get_pillow_config_from_setting(section, pillow_config)

class PillowConfig(namedtuple('PillowConfig', ['section', 'name', 'class_name', 'instance_generator', 'params'])):
    """
    Helper object for getting pillow classes/instances from settings
    """

    def __hash__(self):
        return hash(self.name)

    def get_class(self):
        return _import_class_or_function(self.class_name)

    def get_instance(self, **kwargs):
        if self.instance_generator:
            instance_generator_fn = _import_class_or_function(self.instance_generator)
            params = deepcopy(self.params)  # parameters defined in settings file
            params.update(kwargs)  # parameters passed in via run_ptop
            return instance_generator_fn(self.name, **params)
        else:
            return _get_pillow_instance(self.class_name)

def get_pillow_config_from_setting(section, pillow_config_string_or_dict):
    if isinstance(pillow_config_string_or_dict, str):
        return PillowConfig(
            section,
            pillow_config_string_or_dict.rsplit('.', 1)[1],
            pillow_config_string_or_dict,
            None,
            {}
        )
    else:
        assert 'class' in pillow_config_string_or_dict
        class_name = pillow_config_string_or_dict['class']
        return PillowConfig(
            section,
            pillow_config_string_or_dict.get('name', class_name),
            class_name,
            pillow_config_string_or_dict.get('instance', None),
            pillow_config_string_or_dict.get('params', {}),
        )

def get_pillow_by_name(pillow_class_name, instantiate=True, **kwargs):
    config = get_pillow_config_by_name(pillow_class_name)
    return config.get_instance(**kwargs) if instantiate else config.get_class()

def get_pillow_config_by_name(pillow_name):
    all_configs = get_all_pillow_configs()
    for config in all_configs:
        if config.name == pillow_name:
            return config
    raise PillowNotFoundError('No pillow found with name {}'.format(pillow_name))

def force_seq_int(seq):
    if seq is None or seq == '':
        return None
    elif isinstance(seq, dict):
        # multi-topic checkpoints don't support a single sequence id
        return None
    elif isinstance(seq, str):
        return int(seq.split('-')[0])
    else:
        assert isinstance(seq, int), type(seq)
        return seq

def safe_force_seq_int(seq, default=None):
    if isinstance(seq, dict):
        return default
    try:
        return force_seq_int(seq)
    except (AssertionError, ValueError):
        return default

def _get_consumer():
    return ClosingContextProxy(KafkaConsumer(
        client_id='pillowtop_utils',
        bootstrap_servers=settings.KAFKA_BROKERS,
        request_timeout_ms=1000
    ))

def get_all_pillows_json():
    pillow_configs = get_all_pillow_configs()
    consumer = _get_consumer()
    with consumer:
        return [get_pillow_json(pillow_config, consumer) for pillow_config in pillow_configs]

def get_pillow_json(pillow_config, consumer=None):
    assert isinstance(pillow_config, PillowConfig)

    def _is_kafka(checkpoint):
        return checkpoint.sequence_format == 'json'

    pillow = pillow_config.get_instance()
    if consumer is None:
        consumer = _get_consumer()

    checkpoint = pillow.get_checkpoint()
    timestamp = checkpoint.timestamp
    if timestamp:
        time_since_last = datetime.utcnow() - timestamp
        seconds_since_last = time_since_last.total_seconds()
        hours_since_last = seconds_since_last // 3600

        try:
            # remove microsecond portion
            prettified_time_since_last = str(time_since_last)
            prettified_time_since_last = prettified_time_since_last[0:prettified_time_since_last.index('.')]
        except ValueError:
            pass
    else:
        seconds_since_last = 0
        prettified_time_since_last = ''
        hours_since_last = None

    if _is_kafka(checkpoint):
        # breaking composition boundaries so that we can use one Kafka connection
        offsets = {
            "{},{}".format(tp.topic, tp.partition): offset
            for tp, offset in consumer.end_offsets(list(checkpoint.wrapped_sequence.keys())).items()
        }
    else:
        offsets = pillow.get_change_feed().get_latest_offsets_json()

    def _seq_to_int(checkpoint, seq):
        from pillowtop.models import kafka_seq_to_str
        if _is_kafka(checkpoint):
            return json.loads(kafka_seq_to_str(seq))
        else:
            return force_seq_int(seq)

    return {
        'name': pillow_config.name,
        'seq_format': checkpoint.sequence_format,
        'seq': _seq_to_int(checkpoint, checkpoint.wrapped_sequence),
        'offsets': offsets,
        'seconds_since_last': seconds_since_last,
        'time_since_last': prettified_time_since_last,
        'hours_since_last': hours_since_last
    }

ChangeError = namedtuple('ChangeError', 'change exception')

class ErrorCollector(object):
    def __init__(self):
        self.errors = []

    def add_error(self, error):
        self.errors.append(error)

def build_bulk_payload(index_info, changes, doc_transform=None, error_collector=None):
    doc_transform = doc_transform or (lambda x: x)
    payload = []
    for change in changes:
        if change.deleted and change.id:
            payload.append({
                "delete": {
                    "_index": index_info.index,
                    "_type": index_info.type,
                    "_id": change.id
                }
            })
        elif not change.deleted:
            try:
                doc = change.get_document()
                doc = doc_transform(doc)
                payload.append({
                    "index": {
                        "_index": index_info.index,
                        "_type": index_info.type,
                        "_id": doc['_id']
                    }
                })
                payload.append(doc)
            except Exception as e:
                if not error_collector:
                    raise
                error_collector.add_error(ChangeError(change, e))
    return payload

def prepare_bulk_payloads(bulk_changes, max_size, chunk_size=100):
    payloads = [b'']
    for bulk_chunk in chunked(bulk_changes, chunk_size):
        current_payload = payloads[-1]
        json_bulk_chunks = [
            json.dumps(obj, cls=CommCareJSONEncoder).encode('utf-8')
            for obj in bulk_chunk
        ]
        payload_chunk = b'\n'.join(json_bulk_chunks) + b'\n'
        appended_payload = current_payload + payload_chunk
        new_payload_size = sys.getsizeof(appended_payload)
        if new_payload_size > max_size:
            payloads.append(payload_chunk)
        else:
            payloads[-1] = appended_payload

    return [_f for _f in payloads if _f]

def ensure_matched_revisions(change, fetched_document):
    """
    This function ensures that the document fetched from a change matches the
    revision at which it was pushed to kafka at.

    See http://manage.dimagi.com/default.asp?237983 for more details

    :raises: DocumentMismatchError - Raised when the revisions of the fetched document
        and the change metadata do not match
    """

    change_has_rev = change.metadata and change.metadata.document_rev is not None
    doc_has_rev = fetched_document and '_rev' in fetched_document
    if doc_has_rev and change_has_rev:

        doc_rev = fetched_document['_rev']
        change_rev = change.metadata.document_rev
        if doc_rev != change_rev:
            fetched_rev = _convert_rev_to_int(doc_rev)
            stored_rev = _convert_rev_to_int(change_rev)
            if fetched_rev < stored_rev or stored_rev == -1:
                message = "Mismatched revs for {}: Cloudant rev {} vs. Changes feed rev {}".format(
                    change.id,
                    doc_rev,
                    change_rev
                )
                pillow_logging.warning(message)
                raise DocumentMismatchError(message)

def _convert_rev_to_int(rev):
    try:
        return int(rev.split('-')[0])
    except (ValueError, AttributeError):
        return -1

def ensure_document_exists(change):
    """
    This is to ensure that the Couch document exists in the Couch database when the
    change is processed. We only care about the scenario where the document is missing.

    :raises: DocumentMissingError - Raised when the document is missing (never existed)
    """
    change.get_document()
    if change.error_raised is not None and isinstance(change.error_raised, DocumentMissingError):
        raise change.error_raised
<|endoftext|>


<commit_before>from corehq.util.elastic import es_index
from pillowtop.es_utils import ElasticsearchIndexInfo

USER_INDEX = es_index("hqusers_2016-02-16_1402")
USER_MAPPING={'_all': {'analyzer': 'standard'},
 '_meta': {'comment': 'Simon modified on 2016-01-20',
           'created': None},
 'date_detection': False,
 'date_formats': ['yyyy-MM-dd',
                  "yyyy-MM-dd'T'HH:mm:ssZZ",
                  "yyyy-MM-dd'T'HH:mm:ss.SSSSSS",
                  "yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'",
                  "yyyy-MM-dd'T'HH:mm:ss'Z'",
                  "yyyy-MM-dd'T'HH:mm:ssZ",
                  "yyyy-MM-dd'T'HH:mm:ssZZ'Z'",
                  "yyyy-MM-dd'T'HH:mm:ss.SSSZZ",
                  "yyyy-MM-dd'T'HH:mm:ss",
                  "yyyy-MM-dd' 'HH:mm:ss",
                  "yyyy-MM-dd' 'HH:mm:ss.SSSSSS",
                  "mm/dd/yy' 'HH:mm:ss"],
 'dynamic': False,
 'properties': {'CURRENT_VERSION': {'type': 'string'},
                'base_doc': {'type': 'string'},
                'created_on': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                               'type': 'date'},
                'date_joined': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                                'type': 'date'},
                'doc_type': {'index': 'not_analyzed', 'type': 'string'},
                'domain': {'fields': {'domain': {'index': 'analyzed',
                                                 'type': 'string'},
                                      'exact': {'index': 'not_analyzed',
                                                'type': 'string'}},
                           'type': 'multi_field'},
                'location_id': {'index': 'not_analyzed', 'type': 'string'},
                'domain_membership': {'dynamic': False,
                                      'properties': {'doc_type': {'index': 'not_analyzed',
                                                                  'type': 'string'},
                                                     'domain': {'fields': {'domain': {'index': 'analyzed',
                                                                                'type': 'string'},
                                                                           'exact': {'index': 'not_analyzed',
                                                                                'type': 'string'}},
                                                                'type': 'multi_field'},
                                                     'is_admin': {'type': 'boolean'},
                                                     'override_global_tz': {'type': 'boolean'},
                                                     'role_id': {'type': 'string'},
                                                     'timezone': {'type': 'string'},
                                                     'location_id': {'index': 'not_analyzed',
                                                                     'type': 'string'}},
                                      'type': 'object'},
                'domain_memberships': {'dynamic': False,
                                      'properties': {'doc_type': {'index': 'not_analyzed',
                                                                  'type': 'string'},
                                                     'domain': {'fields': {'domain': {'index': 'analyzed',
                                                                                'type': 'string'},
                                                                           'exact': {'index': 'not_analyzed',
                                                                                'type': 'string'}},
                                                                'type': 'multi_field'},
                                                     'is_admin': {'type': 'boolean'},
                                                     'override_global_tz': {'type': 'boolean'},
                                                     'role_id': {'type': 'string'},
                                                     'timezone': {'type': 'string'},
                                                     'location_id': {'index': 'not_analyzed',
                                                                     'type': 'string'}},
                                      'type': 'object'},
                'email_opt_out': {'type': 'boolean'},
                'eulas': {'dynamic': False,
                          'properties': {'date': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                                                  'type': 'date'},
                                         'doc_type': {'index': 'not_analyzed',
                                                      'type': 'string'},
                                         'signed': {'type': 'boolean'},
                                         'type': {'type': 'string'},
                                         'user_id': {'type': 'string'},
                                         'user_ip': {'type': 'string'},
                                         'version': {'type': 'string'}},
                          'type': 'object'},
                'first_name': {'type': 'string'},
                'is_active': {'type': 'boolean'},
                'is_staff': {'type': 'boolean'},
                'is_superuser': {'type': 'boolean'},
                'language': {'type': 'string'},
                'last_login': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                               'type': 'date'},
                'last_name': {'type': 'string'},
                'password': {'type': 'string'},
                'registering_device_id': {'type': 'string'},
                'status': {'type': 'string'},
                'user_data': {'type': 'object', 'enabled': False},
                'username': {'fields': {'exact': {'include_in_all': False,
                                                  'index': 'not_analyzed',
                                                  'type': 'string'},
                                        'username': {'analyzer': 'standard',
                                                     'index': 'analyzed',
                                                     'type': 'string'}},
                             'type': 'multi_field'},
                '__group_ids': {'type': 'string'},
                '__group_names': {'fields': {'__group_names': { 'index': 'analyzed',
                                                                'type': 'string'},
                                             'exact': {'index': 'not_analyzed',
                                                                'type': 'string'}},
                                  'type': 'multi_field'}}}

USER_META = {
    "settings": {
        "analysis": {
            "analyzer": {
                "default": {
                    "type": "custom",
                    "tokenizer": "whitespace",
                    "filter": ["lowercase"]
                },
            }
        }
    }
}

USER_INDEX_INFO = ElasticsearchIndexInfo(
    index=USER_INDEX,
    alias='hqusers',
    type='user',
    meta=USER_META,
    mapping=USER_MAPPING,
)
<commit_msg>Add analytics_enabled to user mapping<commit_after>from corehq.util.elastic import es_index
from pillowtop.es_utils import ElasticsearchIndexInfo

USER_INDEX = es_index("hqusers_2016-07-11-14-34")
USER_MAPPING={'_all': {'analyzer': 'standard'},
 '_meta': {'comment': 'Operation Privacy modified on 2016-07-11',
           'created': None},
 'date_detection': False,
 'date_formats': ['yyyy-MM-dd',
                  "yyyy-MM-dd'T'HH:mm:ssZZ",
                  "yyyy-MM-dd'T'HH:mm:ss.SSSSSS",
                  "yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'",
                  "yyyy-MM-dd'T'HH:mm:ss'Z'",
                  "yyyy-MM-dd'T'HH:mm:ssZ",
                  "yyyy-MM-dd'T'HH:mm:ssZZ'Z'",
                  "yyyy-MM-dd'T'HH:mm:ss.SSSZZ",
                  "yyyy-MM-dd'T'HH:mm:ss",
                  "yyyy-MM-dd' 'HH:mm:ss",
                  "yyyy-MM-dd' 'HH:mm:ss.SSSSSS",
                  "mm/dd/yy' 'HH:mm:ss"],
 'dynamic': False,
 'properties': {'CURRENT_VERSION': {'type': 'string'},
                'base_doc': {'type': 'string'},
                'created_on': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                               'type': 'date'},
                'date_joined': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                                'type': 'date'},
                'doc_type': {'index': 'not_analyzed', 'type': 'string'},
                'domain': {'fields': {'domain': {'index': 'analyzed',
                                                 'type': 'string'},
                                      'exact': {'index': 'not_analyzed',
                                                'type': 'string'}},
                           'type': 'multi_field'},
                'location_id': {'index': 'not_analyzed', 'type': 'string'},
                'domain_membership': {'dynamic': False,
                                      'properties': {'doc_type': {'index': 'not_analyzed',
                                                                  'type': 'string'},
                                                     'domain': {'fields': {'domain': {'index': 'analyzed',
                                                                                'type': 'string'},
                                                                           'exact': {'index': 'not_analyzed',
                                                                                'type': 'string'}},
                                                                'type': 'multi_field'},
                                                     'is_admin': {'type': 'boolean'},
                                                     'override_global_tz': {'type': 'boolean'},
                                                     'role_id': {'type': 'string'},
                                                     'timezone': {'type': 'string'},
                                                     'location_id': {'index': 'not_analyzed',
                                                                     'type': 'string'}},
                                      'type': 'object'},
                'domain_memberships': {'dynamic': False,
                                      'properties': {'doc_type': {'index': 'not_analyzed',
                                                                  'type': 'string'},
                                                     'domain': {'fields': {'domain': {'index': 'analyzed',
                                                                                'type': 'string'},
                                                                           'exact': {'index': 'not_analyzed',
                                                                                'type': 'string'}},
                                                                'type': 'multi_field'},
                                                     'is_admin': {'type': 'boolean'},
                                                     'override_global_tz': {'type': 'boolean'},
                                                     'role_id': {'type': 'string'},
                                                     'timezone': {'type': 'string'},
                                                     'location_id': {'index': 'not_analyzed',
                                                                     'type': 'string'}},
                                      'type': 'object'},
                'email_opt_out': {'type': 'boolean'},
                'analytics_enabled': {'type': 'boolean'},
                'eulas': {'dynamic': False,
                          'properties': {'date': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                                                  'type': 'date'},
                                         'doc_type': {'index': 'not_analyzed',
                                                      'type': 'string'},
                                         'signed': {'type': 'boolean'},
                                         'type': {'type': 'string'},
                                         'user_id': {'type': 'string'},
                                         'user_ip': {'type': 'string'},
                                         'version': {'type': 'string'}},
                          'type': 'object'},
                'first_name': {'type': 'string'},
                'is_active': {'type': 'boolean'},
                'is_staff': {'type': 'boolean'},
                'is_superuser': {'type': 'boolean'},
                'language': {'type': 'string'},
                'last_login': {'format': "yyyy-MM-dd||yyyy-MM-dd'T'HH:mm:ssZZ||yyyy-MM-dd'T'HH:mm:ss.SSSSSS||yyyy-MM-dd'T'HH:mm:ss.SSSSSS'Z'||yyyy-MM-dd'T'HH:mm:ss'Z'||yyyy-MM-dd'T'HH:mm:ssZ||yyyy-MM-dd'T'HH:mm:ssZZ'Z'||yyyy-MM-dd'T'HH:mm:ss.SSSZZ||yyyy-MM-dd'T'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss||yyyy-MM-dd' 'HH:mm:ss.SSSSSS||mm/dd/yy' 'HH:mm:ss",
                               'type': 'date'},
                'last_name': {'type': 'string'},
                'password': {'type': 'string'},
                'registering_device_id': {'type': 'string'},
                'status': {'type': 'string'},
                'user_data': {'type': 'object', 'enabled': False},
                'username': {'fields': {'exact': {'include_in_all': False,
                                                  'index': 'not_analyzed',
                                                  'type': 'string'},
                                        'username': {'analyzer': 'standard',
                                                     'index': 'analyzed',
                                                     'type': 'string'}},
                             'type': 'multi_field'},
                '__group_ids': {'type': 'string'},
                '__group_names': {'fields': {'__group_names': { 'index': 'analyzed',
                                                                'type': 'string'},
                                             'exact': {'index': 'not_analyzed',
                                                                'type': 'string'}},
                                  'type': 'multi_field'}}}

USER_META = {
    "settings": {
        "analysis": {
            "analyzer": {
                "default": {
                    "type": "custom",
                    "tokenizer": "whitespace",
                    "filter": ["lowercase"]
                },
            }
        }
    }
}

USER_INDEX_INFO = ElasticsearchIndexInfo(
    index=USER_INDEX,
    alias='hqusers',
    type='user',
    meta=USER_META,
    mapping=USER_MAPPING,
)
<|endoftext|>


<commit_before>from django.conf import settings
from django.core.files.storage import default_storage

from pretix.base.i18n import LazyI18nString
from pretix.base.settings import GlobalSettingsObject

from .signals import footer_link, html_head

def contextprocessor(request):
    """
    Adds data to all template contexts
    """
    if request.path.startswith('/control'):
        return {}

    ctx = {
        'css_file': None,
        'DEBUG': settings.DEBUG,
    }
    _html_head = []
    _footer = []

    if hasattr(request, 'event'):
        pretix_settings = request.event.settings
    elif hasattr(request, 'organizer'):
        pretix_settings = request.organizer.settings
    else:
        pretix_settings = GlobalSettingsObject().settings

    text = pretix_settings.get('footer_text', as_type=LazyI18nString)
    link = pretix_settings.get('footer_link', as_type=LazyI18nString)

    if text:
        if link:
            _footer.append({'url': str(link), 'label': str(text)})
        else:
            ctx['footer_text'] = str(text)

    if hasattr(request, 'event'):
        for receiver, response in html_head.send(request.event, request=request):
            _html_head.append(response)
        for receiver, response in footer_link.send(request.event, request=request):
            _footer.append(response)

        if request.event.settings.presale_css_file:
            ctx['css_file'] = default_storage.url(request.event.settings.presale_css_file)
        ctx['event_logo'] = request.event.settings.get('logo_image', as_type=str, default='')[7:]
        ctx['event'] = request.event

    ctx['html_head'] = "".join(_html_head)
    ctx['footer'] = _footer
    ctx['site_url'] = settings.SITE_URL

    return ctx
<commit_msg>Allow signal receivers of footer_link to add multiple links<commit_after>from django.conf import settings
from django.core.files.storage import default_storage

from pretix.base.i18n import LazyI18nString
from pretix.base.settings import GlobalSettingsObject

from .signals import footer_link, html_head

def contextprocessor(request):
    """
    Adds data to all template contexts
    """
    if request.path.startswith('/control'):
        return {}

    ctx = {
        'css_file': None,
        'DEBUG': settings.DEBUG,
    }
    _html_head = []
    _footer = []

    if hasattr(request, 'event'):
        pretix_settings = request.event.settings
    elif hasattr(request, 'organizer'):
        pretix_settings = request.organizer.settings
    else:
        pretix_settings = GlobalSettingsObject().settings

    text = pretix_settings.get('footer_text', as_type=LazyI18nString)
    link = pretix_settings.get('footer_link', as_type=LazyI18nString)

    if text:
        if link:
            _footer.append({'url': str(link), 'label': str(text)})
        else:
            ctx['footer_text'] = str(text)

    if hasattr(request, 'event'):
        for receiver, response in html_head.send(request.event, request=request):
            _html_head.append(response)
        for receiver, response in footer_link.send(request.event, request=request):
            if isinstance(response, list):
                _footer += response
            else:
                _footer.append(response)

        if request.event.settings.presale_css_file:
            ctx['css_file'] = default_storage.url(request.event.settings.presale_css_file)
        ctx['event_logo'] = request.event.settings.get('logo_image', as_type=str, default='')[7:]
        ctx['event'] = request.event

    ctx['html_head'] = "".join(_html_head)
    ctx['footer'] = _footer
    ctx['site_url'] = settings.SITE_URL

    return ctx
<|endoftext|>


<commit_before>import json
import urllib2

from django.core.files import File
from django.core.files.temp import NamedTemporaryFile

from shared_dataverse_information.dataverse_info.forms import DataverseInfoValidationForm
from shared_dataverse_information.dataverse_info.models import TABULAR_TYPES
from apps.registered_dataverse.registered_dataverse_helper import find_registered_dataverse

from geo_utils.msg_util import *

from apps.gis_shapefiles.models import ShapefileInfo
from apps.worldmap_connect.models import WorldMapImportAttempt, WorldMapLayerInfo

import logging
logger = logging.getLogger(__name__)

FAILED_NOT_A_REGISTERED_DATAVERSE = 'FAILED_NOT_A_REGISTERED_DATAVERSE'

class ErrResultMsg:
    def __init__(self, err_type, err_msg):
        self.err_type = err_type
        self.err_msg = err_msg

def get_shapefile_from_dv_api_info(dv_session_token, dataverse_info_dict):
    """Using Dataverse API information, create a :model:`gis_shapefiles.ShapefileInfo' object.  This function should only receive successful responses.

    return True/False, shp_md5 or ErrResultMsg

    Examples:  True, md5 from ShapefileInfo
               False,  ErrResultMsg
    """
    assert dv_session_token is not None, "dv_session_token cannot be None"
    assert type(dataverse_info_dict) is dict, "dataverse_info_dict must be type 'dict'"

    #------------------------------
    # (1) Validate the data (DataverseInfoValidationForm)
    #------------------------------
    #dataverse_info_dict.update({'datafile_id':None})   # for testing
    validation_form = DataverseInfoValidationForm(dataverse_info_dict)
    if not validation_form.is_valid():
        errs = [ '%s: %s' % (k, v) for k,v in validation_form.errors.items()]
        print (errs)
        form_errs = '\n'.join(errs)
        return False, ErrResultMsg(None, form_errs)

    #-------------------------------------------------
    # (2) Check if this is a Registered Dataverse
    #-------------------------------------------------
    registered_dataverse = find_registered_dataverse(dataverse_info_dict['return_to_dataverse_url'])
    if registered_dataverse is None:
        return False, ErrResultMsg(FAILED_NOT_A_REGISTERED_DATAVERSE\
                        , "This dataverse url was not recognized: %s" % dataverse_info_dict['return_to_dataverse_url']\
                    )

    # quick hack for testing
    #
    #datafile_download_url = dataverse_info_dict.get('datafile_download_url', '--datafile_download_url not avail--')
    #if datafile_download_url.find('dvn-build') > -1 and datafile_download_url.find('https') == -1:
    #    dataverse_info_dict['datafile_download_url'] = datafile_download_url.replace('http', 'https')

    #-------------------------------------------------
    # (3) Look for existing Dataverse files in the database
    #    ShapefileInfo and TabularFileInfo objects are routinely deleted, but if file is already here, use it
    #-------------------------------------------------
    params_for_existing_check = dict(datafile_id=dataverse_info_dict.get('datafile_id', -1)\
                                    , dataverse_installation_name=dataverse_info_dict.get('dataverse_installation_name', -1)\
                                    )

    #-------------------------------------------------
    #   BRANCH OUT HERE TO HANDLE TABULAR VS SHAPEFILES
    #-------------------------------------------------

    existing_sets = ShapefileInfo.objects.filter(**params_for_existing_check\
                                ).values_list('id', flat=True\
                                ).order_by('created')

    existing_shapefile_info_ids = list(existing_sets)
    msgt('existing_shapefile_info_ids: %s' % existing_shapefile_info_ids)

    #-------------------------------------------------
    # add dv_session_token and registered_dataverse to dataverse_info_dict
    #-------------------------------------------------
    dataverse_info_dict['dv_session_token'] = dv_session_token
    dataverse_info_dict['registered_dataverse'] = registered_dataverse

    #------------------------------
    # (4) Existing ShapefileInfo(s) found:
    #  (a) Update the ShapefileInfo object
    #  (b) Delete other groups ShapefileInfo object for this datafile and user
    #  (c) Return the md5
    #------------------------------
    if len(existing_shapefile_info_ids) > 1:

        # pop the last ShapefileInfo id off the list of existing_shapefile_info_ids
        shp_id = existing_shapefile_info_ids.pop()

        # delete the rest
        if len(existing_sets) > 0:
            ShapefileInfo.objects.filter(id__in=existing_shapefile_info_ids).delete()   # delete older ShapefileInfo(s)

    #------------------------------
    # (5) Get or create a new ShapefileInfo object
    #------------------------------
    msgt('(5) Get or create a new ShapefileInfo object')
    try:
        # Existing ShapefileInfo:
        #   (1) Assume file is already saved
        #   (2) update the data
        #
        shapefile_info = ShapefileInfo.objects.get(**params_for_existing_check)

        for key, value in dataverse_info_dict.iteritems():
            setattr(shapefile_info, key, value)

        # Save
        shapefile_info.save()
        msg('shapefile info saved')

        # If the file is still available, return it
        if shapefile_info.is_dv_file_available():
            return True, shapefile_info.md5

        # But the file isn't there!!  Delete ShapefileInfo and make a new one
        shapefile_info.delete()

    except ShapefileInfo.DoesNotExist:
        pass
    except:
        return False, ErrResultMsg(None, 'Failed retrieve an existing ShapefileInfo object')

    msg('new file')

    #------------------------------
    # New shapefile info, create object and attach file
    #------------------------------

    shapefile_info = ShapefileInfo(**dataverse_info_dict)
    shapefile_info.save()

    #------------------------------
    # Download and attach file
    #------------------------------
    datafile_download_url = dataverse_info_dict.get('datafile_download_url', '')

    # Add session token.  Gives permission to download/retrieve the file
    #   - http://localhost:8080/api/access/datafile/FILEID?key=YOURAPIKEY
    #
    datafile_download_url = '%s?key=%s' % (datafile_download_url, dv_session_token)
    msg('datafile_download_url: %s' % datafile_download_url)
    datafile_filename = dataverse_info_dict.get('datafile_label', '')

    img_temp = NamedTemporaryFile(delete=True)

    try:
        img_temp.write(urllib2.urlopen(datafile_download_url).read())
    except urllib2.HTTPError as e:
        shapefile_info.delete() # clear shapefile
        err_msg = 'Failed to download shapefile. HTTPError: %s \n\nurl: %s' % (str(e), datafile_download_url)
        return False, ErrResultMsg(None, err_msg)

        #msg(dir(e))
        #msg('HTTP ERROR!: %s' % e.msg)
        #msg('HTTP ERROR!: %s' % e.msg)
        #msgx('blah')
    img_temp.flush()

    shapefile_info.dv_file.save(datafile_filename, File(img_temp))
    shapefile_info.save()

    return True, shapefile_info.md5

def get_successful_worldmap_attempt_from_shapefile(shapefile_info):
    """
    Given a ShapefileInfo object, check for and return a WorldMapLayerInfo object, if available

    :param shapefile_info: ShapefileInfo object
    :return: WorldMapLayerInfo object or None
    """
    assert type(shapefile_info) is ShapefileInfo, "shapefile_info must be a ShapefileInfo object"

    latest_import_attempt = WorldMapImportAttempt.get_latest_attempt(shapefile_info)
    if latest_import_attempt:
        worldmap_layerinfo = latest_import_attempt.get_success_info()
        if type(worldmap_layerinfo) is WorldMapLayerInfo:
            return worldmap_layerinfo
    return None<commit_msg>debugging<commit_after>import json
import urllib2

from django.core.files import File
from django.core.files.temp import NamedTemporaryFile

from shared_dataverse_information.dataverse_info.forms import DataverseInfoValidationForm
from shared_dataverse_information.dataverse_info.models import TABULAR_TYPES
from apps.registered_dataverse.registered_dataverse_helper import find_registered_dataverse

from geo_utils.msg_util import *

from apps.gis_shapefiles.models import ShapefileInfo
from apps.worldmap_connect.models import WorldMapImportAttempt, WorldMapLayerInfo

import logging
logger = logging.getLogger(__name__)

FAILED_NOT_A_REGISTERED_DATAVERSE = 'FAILED_NOT_A_REGISTERED_DATAVERSE'

class ErrResultMsg:
    def __init__(self, err_type, err_msg):
        self.err_type = err_type
        self.err_msg = err_msg

def get_shapefile_from_dv_api_info(dv_session_token, dataverse_info_dict):
    """Using Dataverse API information, create a :model:`gis_shapefiles.ShapefileInfo' object.  This function should only receive successful responses.

    return True/False, shp_md5 or ErrResultMsg

    Examples:  True, md5 from ShapefileInfo
               False,  ErrResultMsg
    """
    assert dv_session_token is not None, "dv_session_token cannot be None"
    assert type(dataverse_info_dict) is dict, "dataverse_info_dict must be type 'dict'"

    #------------------------------
    # (1) Validate the data (DataverseInfoValidationForm)
    #------------------------------
    #dataverse_info_dict.update({'datafile_id':None})   # for testing
    validation_form = DataverseInfoValidationForm(dataverse_info_dict)
    if not validation_form.is_valid():
        errs = [ '%s: %s' % (k, v) for k,v in validation_form.errors.items()]
        print (errs)
        form_errs = '\n'.join(errs)
        return False, ErrResultMsg(None, form_errs)

    #-------------------------------------------------
    # (2) Check if this is a Registered Dataverse
    #-------------------------------------------------
    registered_dataverse = find_registered_dataverse(dataverse_info_dict['return_to_dataverse_url'])
    if registered_dataverse is None:
        return False, ErrResultMsg(FAILED_NOT_A_REGISTERED_DATAVERSE\
                        , "This dataverse url was not recognized: %s" % dataverse_info_dict['return_to_dataverse_url']\
                    )

    # quick hack for testing
    #
    #datafile_download_url = dataverse_info_dict.get('datafile_download_url', '--datafile_download_url not avail--')
    #if datafile_download_url.find('dvn-build') > -1 and datafile_download_url.find('https') == -1:
    #    dataverse_info_dict['datafile_download_url'] = datafile_download_url.replace('http', 'https')

    #-------------------------------------------------
    # (3) Look for existing Dataverse files in the database
    #    ShapefileInfo and TabularFileInfo objects are routinely deleted, but if file is already here, use it
    #-------------------------------------------------
    params_for_existing_check = dict(datafile_id=dataverse_info_dict.get('datafile_id', -1)\
                                    , dataverse_installation_name=dataverse_info_dict.get('dataverse_installation_name', -1)\
                                    )

    #-------------------------------------------------
    #   BRANCH OUT HERE TO HANDLE TABULAR VS SHAPEFILES
    #-------------------------------------------------

    existing_sets = ShapefileInfo.objects.filter(**params_for_existing_check\
                                ).values_list('id', flat=True\
                                ).order_by('created')

    existing_shapefile_info_ids = list(existing_sets)
    msgt('existing_shapefile_info_ids: %s' % existing_shapefile_info_ids)

    #-------------------------------------------------
    # add dv_session_token and registered_dataverse to dataverse_info_dict
    #-------------------------------------------------
    dataverse_info_dict['dv_session_token'] = dv_session_token
    dataverse_info_dict['registered_dataverse'] = registered_dataverse

    #------------------------------
    # (4) Existing ShapefileInfo(s) found:
    #  (a) Update the ShapefileInfo object
    #  (b) Delete other groups ShapefileInfo object for this datafile and user
    #  (c) Return the md5
    #------------------------------
    if len(existing_shapefile_info_ids) > 1:

        # pop the last ShapefileInfo id off the list of existing_shapefile_info_ids
        shp_id = existing_shapefile_info_ids.pop()

        # delete the rest
        if len(existing_sets) > 0:
            ShapefileInfo.objects.filter(id__in=existing_shapefile_info_ids).delete()   # delete older ShapefileInfo(s)

    #------------------------------
    # (5) Get or create a new ShapefileInfo object
    #------------------------------
    msgt('(5) Get or create a new ShapefileInfo object')
    try:
        # Existing ShapefileInfo:
        #   (1) Assume file is already saved
        #   (2) update the data
        #
        shapefile_info = ShapefileInfo.objects.get(**params_for_existing_check)

        for key, value in dataverse_info_dict.iteritems():
            setattr(shapefile_info, key, value)

        # Save
        shapefile_info.save()
        msg('shapefile info saved')

        # If the file is still available, return it
        if shapefile_info.is_dv_file_available():
            return True, shapefile_info.md5
        else:
            # But the file isn't there!!  Delete ShapefileInfo and make a new one
            shapefile_info.delete()

    except ShapefileInfo.DoesNotExist:
        pass
    #except:
    #    msg('Failed to retrieve an existing ShapefileInfo object -- so create a new one')
    #    #return False, ErrResultMsg(None, 'Failed to retrieve an existing ShapefileInfo object')

    msg('new file')

    #------------------------------
    # New shapefile info, create object and attach file
    #------------------------------

    shapefile_info = ShapefileInfo(**dataverse_info_dict)
    shapefile_info.save()

    #------------------------------
    # Download and attach file
    #------------------------------
    datafile_download_url = dataverse_info_dict.get('datafile_download_url', '')

    # Add session token.  Gives permission to download/retrieve the file
    #   - http://localhost:8080/api/access/datafile/FILEID?key=YOURAPIKEY
    #
    datafile_download_url = '%s?key=%s' % (datafile_download_url, dv_session_token)
    msg('datafile_download_url: %s' % datafile_download_url)
    datafile_filename = dataverse_info_dict.get('datafile_label', '')

    img_temp = NamedTemporaryFile(delete=True)

    try:
        img_temp.write(urllib2.urlopen(datafile_download_url).read())
    except urllib2.HTTPError as e:
        shapefile_info.delete() # clear shapefile
        err_msg = 'Failed to download shapefile. HTTPError: %s \n\nurl: %s' % (str(e), datafile_download_url)
        return False, ErrResultMsg(None, err_msg)

        #msg(dir(e))
        #msg('HTTP ERROR!: %s' % e.msg)
        #msg('HTTP ERROR!: %s' % e.msg)
        #msgx('blah')
    img_temp.flush()

    shapefile_info.dv_file.save(datafile_filename, File(img_temp))
    shapefile_info.save()

    return True, shapefile_info.md5

def get_successful_worldmap_attempt_from_shapefile(shapefile_info):
    """
    Given a ShapefileInfo object, check for and return a WorldMapLayerInfo object, if available

    :param shapefile_info: ShapefileInfo object
    :return: WorldMapLayerInfo object or None
    """
    assert type(shapefile_info) is ShapefileInfo, "shapefile_info must be a ShapefileInfo object"

    latest_import_attempt = WorldMapImportAttempt.get_latest_attempt(shapefile_info)
    if latest_import_attempt:
        worldmap_layerinfo = latest_import_attempt.get_success_info()
        if type(worldmap_layerinfo) is WorldMapLayerInfo:
            return worldmap_layerinfo
    return None<|endoftext|>


<commit_before>import sys
import scipy.optimize
from numpy import *
from numpy.random import random

#
# Step 1: load the concordance counts
#

edges_phrase_to_context = {}
edges_context_to_phrase = {}
types = {}
num_edges = 0

for line in sys.stdin:
    phrase, rest = line.strip().split('\t')
    parts = rest.split('|||')
    for i in range(0, len(parts), 2):
        context, count = parts[i:i+2]

        ctx = tuple(filter(lambda x: x != '<PHRASE>', context.split()))
        cnt = int(count.strip()[2:])
        ccs = edges_phrase_to_context.setdefault(phrase, {})
        ccs[ctx] = cnt
        pcs = edges_context_to_phrase.setdefault(ctx, {})
        pcs[phrase] = cnt

        for token in ctx:
            types.setdefault(token, len(types))
        for token in phrase.split():
            types.setdefault(token, len(types))

        num_edges += 1

print 'Read in', num_edges, 'edges and', len(types), 'word types'

#
# Step 2: initialise the model parameters
#

num_tags = 5
num_types = len(types)
delta = int(sys.argv[1])
gamma = int(sys.argv[2])

def normalise(a):
    return a / sum(a)

# Pr(tag)
tagDist = normalise(random(num_tags)+1)
# Pr(context at pos i = w | tag) indexed by i, tag, word
contextWordDist = [[normalise(random(num_types)+1) for t in range(num_tags)] for i in range(4)]
# PR langrange multipliers
lamba = zeros(num_edges * num_tags)
lamba_index = {}
next = 0
for phrase, ccs in edges_phrase_to_context.items():
    for context in ccs.keys():
        lamba_index[phrase,context] = next
        next += num_tags

#
# Step 3: expectation maximisation
#

for iteration in range(20):
    tagCounts = zeros(num_tags)
    contextWordCounts = [[zeros(num_types) for t in range(num_tags)] for i in range(4)]

    #print 'tagDist', tagDist
    #print 'contextWordCounts[0][0]', contextWordCounts[0][0]

    # Tune lambda
    # dual: min log Z(lamba) s.t. lamba >= 0;
    # sum_c lamba_pct <= delta; sum_p lamba_pct <= gamma
    def dual(ls):
        logz = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for context, count in ccs.items():
                conditionals = zeros(num_tags)
                for t in range(num_tags):
                    prob = tagDist[t]
                    for i in range(4):
                        prob *= contextWordDist[i][t][types[context[i]]]
                    conditionals[t] = prob
                cz = sum(conditionals)
                conditionals /= cz

                local_z = 0
                for t in range(num_tags):
                    local_z += conditionals[t] * exp(-ls[lamba_index[phrase,context] + t])
                logz += log(local_z) * count

        #print 'ls', ls
        #print 'lambda', list(ls)
        #print 'dual', logz
        return logz

    def primal(ls):
        # FIXME: returns negative values for KL (impossible)
        logz = dual(ls)
        kl = -logz

        expectations = zeros(lamba.shape)
        for phrase, ccs in edges_phrase_to_context.items():
            for context, count in ccs.items():
                conditionals = zeros(num_tags)
                for t in range(num_tags):
                    prob = tagDist[t]
                    for i in range(4):
                        prob *= contextWordDist[i][t][types[context[i]]]
                    conditionals[t] = prob
                cz = sum(conditionals)
                conditionals /= cz

                scores = zeros(num_tags)
                for t in range(num_tags):
                    scores[t] = conditionals[t] * exp(-ls[lamba_index[phrase,context] + t])
                local_z = sum(scores)

                for t in range(num_tags):
                    li = lamba_index[phrase,context] + t
                    expectations[li] = scores[t] / local_z * count
                    kl -= expectations[li] * ls[li]

        pt_l1linf = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for t in range(num_tags):
                best = -1e500
                for context, count in ccs.items():
                    li = lamba_index[phrase,context] + t
                    s = expectations[li]
                    if s > best: best = s
                pt_l1linf += best

        ct_l1linf = 0
        for context, pcs in edges_context_to_phrase.items():
            for t in range(num_tags):
                best = -1e500
                for phrase, count in pcs.items():
                    li = lamba_index[phrase,context] + t
                    s = expectations[li]
                    if s > best: best = s
                ct_l1linf += best

        return kl, pt_l1linf, ct_l1linf, kl + delta * pt_l1linf + gamma * ct_l1linf

    def dual_deriv(ls):
        # d/dl log(z) = E_q[phi]
        deriv = zeros(num_edges * num_tags)
        for phrase, ccs in edges_phrase_to_context.items():
            for context, count in ccs.items():
                conditionals = zeros(num_tags)
                for t in range(num_tags):
                    prob = tagDist[t]
                    for i in range(4):
                        prob *= contextWordDist[i][t][types[context[i]]]
                    conditionals[t] = prob
                cz = sum(conditionals)
                conditionals /= cz

                scores = zeros(num_tags)
                for t in range(num_tags):
                    scores[t] = conditionals[t] * exp(-ls[lamba_index[phrase,context] + t])
                local_z = sum(scores)

                for t in range(num_tags):
                    deriv[lamba_index[phrase,context] + t] -= count * scores[t] / local_z

        #print 'ddual', deriv
        return deriv

    def constraints(ls):
        cons = []
        if delta > 0:
            for phrase, ccs in edges_phrase_to_context.items():
                for t in range(num_tags):
                    total = delta
                    for cprime in ccs.keys():
                        total -= ls[lamba_index[phrase, cprime] + t]
                    cons.append(total)

        if gamma > 0:
            for context, pcs in edges_context_to_phrase.items():
                for t in range(num_tags):
                    total = gamma
                    for pprime in pcs.keys():
                        total -= ls[lamba_index[pprime, context] + t]
                    cons.append(total)
        #print 'cons', cons
        return cons

    def constraints_deriv(ls):
        cons = []
        if delta > 0:
            for phrase, ccs in edges_phrase_to_context.items():
                for t in range(num_tags):
                    d = zeros(num_edges * num_tags)
                    for cprime in ccs.keys():
                        d[lamba_index[phrase, cprime] + t] = -1
                    cons.append(d)

        if gamma > 0:
            for context, pcs in edges_context_to_phrase.items():
                for t in range(num_tags):
                    d = zeros(num_edges * num_tags)
                    for pprime in pcs.keys():
                        d[lamba_index[pprime, context] + t] = -1
                    cons.append(d)
        #print 'dcons', cons
        return cons

    print 'Pre lambda optimisation dual', dual(lamba), 'primal', primal(lamba)
    lamba = scipy.optimize.fmin_slsqp(dual, lamba,
                            bounds=[(0, delta)] * (num_edges * num_tags),
                            f_ieqcons=constraints,
                            fprime=dual_deriv,
                            fprime_ieqcons=constraints_deriv,
                            iprint=0)
    print 'Post lambda optimisation dual', dual(lamba), 'primal', primal(lamba)

    # E-step
    llh = z = 0
    for phrase, ccs in edges_phrase_to_context.items():
        for context, count in ccs.items():
            conditionals = zeros(num_tags)
            for t in range(num_tags):
                prob = tagDist[t]
                for i in range(4):
                    prob *= contextWordDist[i][t][types[context[i]]]
                conditionals[t] = prob
            cz = sum(conditionals)
            conditionals /= cz
            llh += log(cz) * count

            scores = zeros(num_tags)
            li = lamba_index[phrase, context]
            for t in range(num_tags):
                scores[t] = conditionals[t] * exp(-lamba[li + t])
            z += count * sum(scores)

            tagCounts += count * scores
            for i in range(4):
                for t in range(num_tags):
                    contextWordCounts[i][t][types[context[i]]] += count * scores[t]

    print 'iteration', iteration, 'llh', llh, 'logz', log(z)

    # M-step
    tagDist = normalise(tagCounts)
    for i in range(4):
        for t in range(num_tags):
            contextWordDist[i][t] = normalise(contextWordCounts[i][t])

for phrase, ccs in edges_phrase_to_context.items():
    for context, count in ccs.items():
        conditionals = zeros(num_tags)
        for t in range(num_tags):
            prob = tagDist[t]
            for i in range(4):
                prob *= contextWordDist[i][t][types[context[i]]]
            conditionals[t] = prob
        cz = sum(conditionals)
        conditionals /= cz

        scores = zeros(num_tags)
        li = lamba_index[phrase, context]
        for t in range(num_tags):
            scores[t] = conditionals[t] * exp(-lamba[li + t])

        print '%s\t%s ||| C=%d ||| %d |||' % (phrase, context, count, argmax(scores)), scores / sum(scores)
<commit_msg>Fixed context-based constraints<commit_after>import sys
import scipy.optimize
from numpy import *
from numpy.random import random

#
# Step 1: load the concordance counts
#

edges_phrase_to_context = {}
edges_context_to_phrase = {}
types = {}
num_edges = 0

for line in sys.stdin:
    phrase, rest = line.strip().split('\t')
    parts = rest.split('|||')
    for i in range(0, len(parts), 2):
        context, count = parts[i:i+2]

        ctx = tuple(filter(lambda x: x != '<PHRASE>', context.split()))
        cnt = int(count.strip()[2:])
        ccs = edges_phrase_to_context.setdefault(phrase, {})
        ccs[ctx] = cnt
        pcs = edges_context_to_phrase.setdefault(ctx, {})
        pcs[phrase] = cnt

        for token in ctx:
            types.setdefault(token, len(types))
        for token in phrase.split():
            types.setdefault(token, len(types))

        num_edges += 1

print 'Read in', num_edges, 'edges and', len(types), 'word types'

#
# Step 2: initialise the model parameters
#

num_tags = 5
num_types = len(types)
num_phrases = len(edges_phrase_to_context)
num_contexts = len(edges_context_to_phrase)
delta = int(sys.argv[1])
gamma = int(sys.argv[2])

def normalise(a):
    return a / sum(a)

# Pr(tag)
tagDist = normalise(random(num_tags)+1)
# Pr(context at pos i = w | tag) indexed by i, tag, word
contextWordDist = [[normalise(random(num_types)+1) for t in range(num_tags)] for i in range(4)]
# PR langrange multipliers
lamba = zeros(2 * num_edges * num_tags)
omega_offset = num_edges * num_tags
lamba_index = {}
next = 0
for phrase, ccs in edges_phrase_to_context.items():
    for context in ccs.keys():
        lamba_index[phrase,context] = next
        next += num_tags

#
# Step 3: expectation maximisation
#

for iteration in range(20):
    tagCounts = zeros(num_tags)
    contextWordCounts = [[zeros(num_types) for t in range(num_tags)] for i in range(4)]

    #print 'tagDist', tagDist
    #print 'contextWordCounts[0][0]', contextWordCounts[0][0]

    # Tune lambda
    # dual: min log Z(lamba) s.t. lamba >= 0;
    # sum_c lamba_pct <= delta; sum_p lamba_pct <= gamma
    def dual(ls):
        logz = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for context, count in ccs.items():
                conditionals = zeros(num_tags)
                for t in range(num_tags):
                    prob = tagDist[t]
                    for i in range(4):
                        prob *= contextWordDist[i][t][types[context[i]]]
                    conditionals[t] = prob
                cz = sum(conditionals)
                conditionals /= cz

                local_z = 0
                for t in range(num_tags):
                    li = lamba_index[phrase,context] + t
                    local_z += conditionals[t] * exp(-ls[li] - ls[omega_offset+li])
                logz += log(local_z) * count

        #print 'ls', ls
        #print 'lambda', list(ls)
        #print 'dual', logz
        return logz

    def loglikelihood():
        llh = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for context, count in ccs.items():
                conditionals = zeros(num_tags)
                for t in range(num_tags):
                    prob = tagDist[t]
                    for i in range(4):
                        prob *= contextWordDist[i][t][types[context[i]]]
                    conditionals[t] = prob
                cz = sum(conditionals)
                llh += log(cz) * count
        return llh

    def primal(ls):
        # FIXME: returns negative values for KL (impossible)
        logz = dual(ls)
        expectations = -dual_deriv(ls)
        kl = -logz - dot(ls, expectations)
        llh = loglikelihood()

        pt_l1linf = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for t in range(num_tags):
                best = -1e500
                for context, count in ccs.items():
                    li = lamba_index[phrase,context] + t
                    s = expectations[li]
                    if s > best: best = s
                pt_l1linf += best

        ct_l1linf = 0
        for context, pcs in edges_context_to_phrase.items():
            for t in range(num_tags):
                best = -1e500
                for phrase, count in pcs.items():
                    li = lamba_index[phrase,context] + t
                    s = expectations[li]
                    if s > best: best = s
                ct_l1linf += best

        return llh, kl, pt_l1linf, ct_l1linf, llh + kl + delta * pt_l1linf + gamma * ct_l1linf

    def dual_deriv(ls):
        # d/dl log(z) = E_q[phi]
        deriv = zeros(2 * num_edges * num_tags)
        for phrase, ccs in edges_phrase_to_context.items():
            for context, count in ccs.items():
                conditionals = zeros(num_tags)
                for t in range(num_tags):
                    prob = tagDist[t]
                    for i in range(4):
                        prob *= contextWordDist[i][t][types[context[i]]]
                    conditionals[t] = prob
                cz = sum(conditionals)
                conditionals /= cz

                scores = zeros(num_tags)
                for t in range(num_tags):
                    li = lamba_index[phrase,context] + t
                    scores[t] = conditionals[t] * exp(-ls[li] - ls[omega_offset + li])
                local_z = sum(scores)

                for t in range(num_tags):
                    if delta > 0:
                        deriv[lamba_index[phrase,context] + t] -= count * scores[t] / local_z
                    if gamma > 0:
                        deriv[omega_offset + lamba_index[phrase,context] + t] -= count * scores[t] / local_z

        #print 'ddual', deriv
        return deriv

    def constraints(ls):
        cons = zeros(num_phrases * num_tags + num_edges * num_tags)

        index = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for t in range(num_tags):
                if delta > 0:
                    total = delta
                    for cprime in ccs.keys():
                        total -= ls[lamba_index[phrase, cprime] + t]
                    cons[index] = total
                index += 1

        for context, pcs in edges_context_to_phrase.items():
            for t in range(num_tags):
                if gamma > 0:
                    total = gamma
                    for pprime in pcs.keys():
                        total -= ls[omega_offset + lamba_index[pprime, context] + t]
                    cons[index] = total
                index += 1

        #print 'cons', cons
        return cons

    def constraints_deriv(ls):
        cons = zeros((num_phrases * num_tags + num_edges * num_tags, 2 * num_edges * num_tags))

        index = 0
        for phrase, ccs in edges_phrase_to_context.items():
            for t in range(num_tags):
                if delta > 0:
                    d = cons[index,:]#zeros(num_edges * num_tags)
                    for cprime in ccs.keys():
                        d[lamba_index[phrase, cprime] + t] = -1
                    #cons[index] = d
                index += 1

        for context, pcs in edges_context_to_phrase.items():
            for t in range(num_tags):
                if gamma > 0:
                    d = cons[index,:]#d = zeros(num_edges * num_tags)
                    for pprime in pcs.keys():
                        d[omega_offset + lamba_index[pprime, context] + t] = -1
                    #cons[index] = d
                index += 1
        #print 'dcons', cons
        return cons

    print 'Pre lambda optimisation dual', dual(lamba), 'primal', primal(lamba)
    #print 'lambda', lamba, lamba.shape
    #print 'bounds', [(0, max(delta, gamma))] * (2 * num_edges * num_tags)

    lamba = scipy.optimize.fmin_slsqp(dual, lamba,
                            bounds=[(0, max(delta, gamma))] * (2 * num_edges * num_tags),
                            f_ieqcons=constraints,
                            fprime=dual_deriv,
                            fprime_ieqcons=constraints_deriv,
                            iprint=0)
    print 'Post lambda optimisation dual', dual(lamba), 'primal', primal(lamba)

    # E-step
    llh = z = 0
    for phrase, ccs in edges_phrase_to_context.items():
        for context, count in ccs.items():
            conditionals = zeros(num_tags)
            for t in range(num_tags):
                prob = tagDist[t]
                for i in range(4):
                    prob *= contextWordDist[i][t][types[context[i]]]
                conditionals[t] = prob
            cz = sum(conditionals)
            conditionals /= cz
            llh += log(cz) * count

            scores = zeros(num_tags)
            li = lamba_index[phrase, context]
            for t in range(num_tags):
                scores[t] = conditionals[t] * exp(-lamba[li + t] - lamba[omega_offset + li + t])
            z += count * sum(scores)

            tagCounts += count * scores
            for i in range(4):
                for t in range(num_tags):
                    contextWordCounts[i][t][types[context[i]]] += count * scores[t]

    print 'iteration', iteration, 'llh', llh, 'logz', log(z)

    # M-step
    tagDist = normalise(tagCounts)
    for i in range(4):
        for t in range(num_tags):
            contextWordDist[i][t] = normalise(contextWordCounts[i][t])

for phrase, ccs in edges_phrase_to_context.items():
    for context, count in ccs.items():
        conditionals = zeros(num_tags)
        for t in range(num_tags):
            prob = tagDist[t]
            for i in range(4):
                prob *= contextWordDist[i][t][types[context[i]]]
            conditionals[t] = prob
        cz = sum(conditionals)
        conditionals /= cz

        #scores = zeros(num_tags)
        #li = lamba_index[phrase, context]
        #for t in range(num_tags):
        #    scores[t] = conditionals[t] * exp(-lamba[li + t])

        #print '%s\t%s ||| C=%d ||| %d |||' % (phrase, context, count, argmax(scores)), scores / sum(scores)
        print '%s\t%s ||| C=%d ||| %d |||' % (phrase, context, count, argmax(conditionals)), conditionals
<|endoftext|>