<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2016 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

# 2015.06.26 13:21:50 CEST
"""
Utilities to build a report writer generating a .rst report for a calculation
"""
from __future__ import print_function, unicode_literals
from openquake.baselib.python3compat import decode
import os
import sys
import mock
import time

from openquake.baselib.general import humansize, AccumDict
from openquake.baselib.python3compat import encode
from openquake.commonlib import readinput, parallel
from openquake.calculators.classical import PSHACalculator
from openquake.calculators import views

def indent(text):
    return '  ' + '\n  '.join(text.splitlines())

def count_eff_ruptures(sources, sitecol, gsims, monitor):
    """
    Count the number of ruptures contained in the given sources and return
    a dictionary src_group_id -> num_ruptures. All sources belong to the
    same tectonic region type.
    """
    acc = AccumDict()
    acc.grp_id = sources[0].src_group_id
    acc.calc_times = []
    acc.eff_ruptures = {acc.grp_id: sum(src.num_ruptures for src in sources)}
    return acc

class ReportWriter(object):
    """
    A particularly smart view over the datastore
    """
    title = dict(
        params='Parameters',
        inputs='Input files',
        csm_info='Composite source model',
        required_params_per_trt='Required parameters per tectonic region type',
        ruptures_per_trt='Number of ruptures per tectonic region type',
        ruptures_events='Specific information for event based',
        rlzs_assoc='Realizations per (TRT, GSIM)',
        job_info='Informational data',
        biggest_ebr_gmf='Maximum memory allocated for the GMFs',
        avglosses_data_transfer='Estimated data transfer for the avglosses',
        exposure_info='Exposure model',
        short_source_info='Slowest sources',
        task_slowest='Slowest task',
        task_info='Information about the tasks',
        times_by_source_class='Computation times by source typology',
        performance='Slowest operations',
    )

    def __init__(self, dstore):
        self.dstore = dstore
        self.oq = oq = dstore['oqparam']
        self.text = (decode(oq.description) + '\n' + '=' * len(oq.description))
        info = dstore['job_info']
        dpath = dstore.hdf5path
        mtime = os.path.getmtime(dpath)
        host = '%s:%s' % (info.hostname, decode(dpath))
        updated = str(time.ctime(mtime))
        versions = sorted(dstore['/'].attrs.items())
        self.text += '\n\n' + views.rst_table([[host, updated]] + versions)
        # NB: in the future, the sitecol could be transferred as
        # an array by leveraging the HDF5 serialization protocol;
        # for the moment however the size of the
        # data to transfer is given by the usual pickle
        sitecol_size = humansize(len(parallel.Pickled(dstore['sitecol'])))
        self.text += '\n\nnum_sites = %d, sitecol = %s' % (
            len(dstore['sitecol']), sitecol_size)

    def add(self, name, obj=None):
        """Add the view named `name` to the report text"""
        title = self.title[name]
        line = '-' * len(title)
        if obj:
            text = '\n::\n\n' + indent(str(obj))
        else:
            text = views.view(name, self.dstore)
        self.text += '\n'.join(['\n\n' + title, line, text])

    def make_report(self):
        """Build the report and return a restructed text string"""
        oq, ds = self.oq, self.dstore
        for name in ('params', 'inputs'):
            self.add(name)
        if 'composite_source_model' in ds:
            self.add('csm_info')
            self.add('required_params_per_trt')
        self.add('rlzs_assoc', ds['csm_info'].get_rlzs_assoc())
        if 'source_info' in ds:
            self.add('ruptures_per_trt')
        if 'scenario' not in oq.calculation_mode:
            self.add('job_info')
        if 'ruptures' in ds:
            self.add('ruptures_events')
        if oq.calculation_mode in ('event_based_risk',):
            if 'ruptures' in ds:
                self.add('biggest_ebr_gmf')
            self.add('avglosses_data_transfer')
        if 'exposure' in oq.inputs:
            self.add('exposure_info')
        if 'source_info' in ds:
            self.add('short_source_info')
            self.add('times_by_source_class')
        if 'task_info' in ds:
            self.add('task_info')
            if 'classical' in ds['task_info']:
                self.add('task_slowest')
        if 'performance_data' in ds:
            self.add('performance')
        return self.text

    def save(self, fname):
        """Save the report"""
        with open(fname, 'wb') as f:
            f.write(encode(self.text))

def build_report(job_ini, output_dir=None):
    """
    Write a `report.csv` file with information about the calculation
    without running it

    :param job_ini:
        full pathname of the job.ini file
    :param output_dir:
        the directory where the report is written (default the input directory)
    """
    oq = readinput.get_oqparam(job_ini)
    output_dir = output_dir or os.path.dirname(job_ini)
    from openquake.calculators import base  # ugly
    calc = base.calculators(oq)
    calc.save_params()  # needed to save oqparam

    # some taken is care so that the real calculation is not run:
    # the goal is to extract information about the source management only
    with mock.patch.object(PSHACalculator, 'core_task', count_eff_ruptures):
        calc.pre_execute()
    if hasattr(calc, '_composite_source_model'):
        calc.datastore['csm_info'] = calc.csm.info
    rw = ReportWriter(calc.datastore)
    rw.make_report()
    report = (os.path.join(output_dir, 'report.rst') if output_dir
              else calc.datastore.export_path('report.rst'))
    try:
        rw.save(report)
    except IOError as exc:  # permission error
        sys.stderr.write(str(exc) + '\n')
    return report
<commit_msg>Fixed the reports<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2016 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

# 2015.06.26 13:21:50 CEST
"""
Utilities to build a report writer generating a .rst report for a calculation
"""
from __future__ import print_function, unicode_literals
from openquake.baselib.python3compat import decode
import os
import sys
import mock
import time

from openquake.baselib.general import humansize, AccumDict
from openquake.baselib.python3compat import encode
from openquake.commonlib import readinput, parallel
from openquake.calculators.classical import PSHACalculator
from openquake.calculators import views

def indent(text):
    return '  ' + '\n  '.join(text.splitlines())

def count_eff_ruptures(sources, sitecol, gsims, monitor):
    """
    Count the number of ruptures contained in the given sources and return
    a dictionary src_group_id -> num_ruptures. All sources belong to the
    same tectonic region type.
    """
    acc = AccumDict()
    acc.grp_id = sources[0].src_group_id
    acc.calc_times = []
    acc.eff_ruptures = {acc.grp_id: sum(src.num_ruptures for src in sources)}
    return acc

class ReportWriter(object):
    """
    A particularly smart view over the datastore
    """
    title = dict(
        params='Parameters',
        inputs='Input files',
        csm_info='Composite source model',
        required_params_per_trt='Required parameters per tectonic region type',
        ruptures_per_trt='Number of ruptures per tectonic region type',
        ruptures_events='Specific information for event based',
        rlzs_assoc='Realizations per (TRT, GSIM)',
        job_info='Informational data',
        biggest_ebr_gmf='Maximum memory allocated for the GMFs',
        avglosses_data_transfer='Estimated data transfer for the avglosses',
        exposure_info='Exposure model',
        short_source_info='Slowest sources',
        task_slowest='Slowest task',
        task_info='Information about the tasks',
        times_by_source_class='Computation times by source typology',
        performance='Slowest operations',
    )

    def __init__(self, dstore):
        self.dstore = dstore
        self.oq = oq = dstore['oqparam']
        self.text = (decode(oq.description) + '\n' + '=' * len(oq.description))
        info = dstore['job_info']
        dpath = dstore.hdf5path
        mtime = os.path.getmtime(dpath)
        host = '%s:%s' % (info.hostname, decode(dpath))
        updated = str(time.ctime(mtime))
        versions = sorted(dstore['/'].attrs.items())
        self.text += '\n\n' + views.rst_table([[host, updated]] + versions)
        # NB: in the future, the sitecol could be transferred as
        # an array by leveraging the HDF5 serialization protocol;
        # for the moment however the size of the
        # data to transfer is given by the usual pickle
        sitecol_size = humansize(len(parallel.Pickled(dstore['sitecol'])))
        self.text += '\n\nnum_sites = %d, sitecol = %s' % (
            len(dstore['sitecol']), sitecol_size)

    def add(self, name, obj=None):
        """Add the view named `name` to the report text"""
        title = self.title[name]
        line = '-' * len(title)
        if obj:
            text = '\n::\n\n' + indent(str(obj))
        else:
            text = views.view(name, self.dstore)
        self.text += '\n'.join(['\n\n' + title, line, text])

    def make_report(self):
        """Build the report and return a restructed text string"""
        oq, ds = self.oq, self.dstore
        for name in ('params', 'inputs'):
            self.add(name)
        if 'composite_source_model' in ds:
            self.add('csm_info')
            self.add('required_params_per_trt')
        self.add('rlzs_assoc', ds['csm_info'].get_rlzs_assoc())
        if 'source_info' in ds:
            self.add('ruptures_per_trt')
        if 'scenario' not in oq.calculation_mode:
            self.add('job_info')
        if 'events' in ds:
            self.add('ruptures_events')
        if oq.calculation_mode in ('event_based_risk',):
            if 'events' in ds:
                self.add('biggest_ebr_gmf')
            self.add('avglosses_data_transfer')
        if 'exposure' in oq.inputs:
            self.add('exposure_info')
        if 'source_info' in ds:
            self.add('short_source_info')
            self.add('times_by_source_class')
        if 'task_info' in ds:
            self.add('task_info')
            if 'classical' in ds['task_info']:
                self.add('task_slowest')
        if 'performance_data' in ds:
            self.add('performance')
        return self.text

    def save(self, fname):
        """Save the report"""
        with open(fname, 'wb') as f:
            f.write(encode(self.text))

def build_report(job_ini, output_dir=None):
    """
    Write a `report.csv` file with information about the calculation
    without running it

    :param job_ini:
        full pathname of the job.ini file
    :param output_dir:
        the directory where the report is written (default the input directory)
    """
    oq = readinput.get_oqparam(job_ini)
    output_dir = output_dir or os.path.dirname(job_ini)
    from openquake.calculators import base  # ugly
    calc = base.calculators(oq)
    calc.save_params()  # needed to save oqparam

    # some taken is care so that the real calculation is not run:
    # the goal is to extract information about the source management only
    with mock.patch.object(PSHACalculator, 'core_task', count_eff_ruptures):
        calc.pre_execute()
    if hasattr(calc, '_composite_source_model'):
        calc.datastore['csm_info'] = calc.csm.info
    rw = ReportWriter(calc.datastore)
    rw.make_report()
    report = (os.path.join(output_dir, 'report.rst') if output_dir
              else calc.datastore.export_path('report.rst'))
    try:
        rw.save(report)
    except IOError as exc:  # permission error
        sys.stderr.write(str(exc) + '\n')
    return report
<|endoftext|>


<commit_before>import unittest
from openquake.commonlib import sap

def f(a, b, c, d=1):
    return [a, b, c, d]

class SapTestCase(unittest.TestCase):
    def test_ok(self):
        p = sap.Parser(f)
        p.arg('a', 'first argument')
        p.opt('b', 'second argument')
        p.flg('c', 'third argument')
        p.opt('d', 'fourth argument')
        self.assertEqual(
            ['1', '2', False, '3'], p.callfunc('1 -b=2 -d=3'.split()))

    def test_group(self):
        p = sap.Parser(f, version=None)
        p.arg('a', 'first argument')
        p.opt('b', 'second argument')
        p.group('other arguments')
        p.flg('c', 'third argument')
        p.opt('d', 'fourth argument')
        self.assertEqual(p.help(), '''\
usage: nosetests [-h] [-b B] [-c] [-d 1] a

positional arguments:
  a            first argument

optional arguments:
  -h, --help   show this help message and exit
  -b B, --b B  second argument

other arguments:
  -c, --c      third argument
  -d 1, --d 1  fourth argument
''')

    def test_NameError(self):
        p = sap.Parser(f)
        p.arg('a', 'first argument')
        with self.assertRaises(NameError):
            p.flg('c', 'third argument')

    def test_help(self):
        p = sap.Parser(f, version=None)
        p.arg('a', 'first argument')
        p.opt('b', 'second argument')
        self.assertEqual(p.help(), '''\
usage: nosetests [-h] [-b B] a

positional arguments:
  a            first argument

optional arguments:
  -h, --help   show this help message and exit
  -b B, --b B  second argument
''')
        # missing argparse description for 'c' and 'd'
        with self.assertRaises(NameError):
            self.assertEqual(
                ['1', '2', False, '3'], p.callfunc('1 -b=2 -d=3'.split()))

    def test_long_argument(self):
        # test the replacement '_' -> '-' in variable names
        p = sap.Parser(lambda a_long_argument: None)
        p.opt('a_long_argument', 'a long argument')
        self.assertEqual(p.help(), '''\
usage: nosetests [-h] [-v] [-a A_LONG_ARGUMENT]

optional arguments:
  -h, --help            show this help message and exit
  -v, --version         show program's version number and exit
  -a A_LONG_ARGUMENT, --a-long-argument A_LONG_ARGUMENT
                        a long argument
''')
<commit_msg>Tested disabling of the help facility<commit_after>import unittest
from openquake.commonlib import sap

def f(a, b, c, d=1):
    return [a, b, c, d]

class SapTestCase(unittest.TestCase):
    def test_ok(self):
        p = sap.Parser(f)
        p.arg('a', 'first argument')
        p.opt('b', 'second argument')
        p.flg('c', 'third argument')
        p.opt('d', 'fourth argument')
        self.assertEqual(
            ['1', '2', False, '3'], p.callfunc('1 -b=2 -d=3'.split()))

    def test_group(self):
        p = sap.Parser(f, version=None)
        p.arg('a', 'first argument')
        p.opt('b', 'second argument')
        p.group('other arguments')
        p.flg('c', 'third argument')
        p.opt('d', 'fourth argument')
        self.assertEqual(p.help(), '''\
usage: nosetests [-h] [-b B] [-c] [-d 1] a

positional arguments:
  a            first argument

optional arguments:
  -h, --help   show this help message and exit
  -b B, --b B  second argument

other arguments:
  -c, --c      third argument
  -d 1, --d 1  fourth argument
''')

    def test_NameError(self):
        p = sap.Parser(f)
        p.arg('a', 'first argument')
        with self.assertRaises(NameError):
            p.flg('c', 'third argument')

    def test_no_help(self):
        p = sap.Parser(f, help=None, version=None)
        p.arg('a', 'first argument')
        p.opt('b', 'second argument')
        self.assertEqual(p.help(), '''\
usage: nosetests [-b B] a

positional arguments:
  a            first argument

optional arguments:
  -b B, --b B  second argument
''')
        # missing argparse specification for 'c'
        with self.assertRaises(NameError):
            p.check_arguments()

    def test_long_argument(self):
        # test the replacement '_' -> '-' in variable names
        p = sap.Parser(lambda a_long_argument: None)
        p.opt('a_long_argument', 'a long argument')
        self.assertEqual(p.help(), '''\
usage: nosetests [-h] [-v] [-a A_LONG_ARGUMENT]

optional arguments:
  -h, --help            show this help message and exit
  -v, --version         show program's version number and exit
  -a A_LONG_ARGUMENT, --a-long-argument A_LONG_ARGUMENT
                        a long argument
''')
<|endoftext|>


<commit_before>from pru_speak import pru_speak
import socket
#Server on local host @ port 6060
TCP_IP = '127.0.0.1'
TCP_PORT = 6060

#The max size upto which is recieved at one go
BUFFER_SIZE = 1024 * 2

#create a passive socket and bind it to 127.0.0.1:6060
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind((TCP_IP, TCP_PORT))

#Take max of one connection
s.listen(1)
print "PRU Speak server at ", TCP_IP + ":" + str(TCP_PORT)

#connection is a one time affair - handles only one client
conn, addr = s.accept()
print 'Connection address:', addr

while True:
        try:
                #blocking wait for data from client
                bs_code = conn.recv(BUFFER_SIZE)
                print "recieved : \n", bs_code, "\n"

                if (not bs_code) or (bs_code == 'EXIT'):
			break #send empty data to shutdown server

		try:
                	ret = pru_speak.execute_instruction(bs_code)
			print "Return value : ", ret
			if ret == []:
				#no return value
				conn.send('\n')
			else:
				#the sends the list of ret values as a string sperated by " "
	        	        conn.send(" ".join(map(str, ret)))

	        except Exception as e:
                	print e

        except Exception as e:
                print "error : ", e
                break

conn.close()
<commit_msg>Add newline to socket listener<commit_after>from pru_speak import pru_speak
import socket
#Server on local host @ port 6060
TCP_IP = '127.0.0.1'
TCP_PORT = 6060

#The max size upto which is recieved at one go
BUFFER_SIZE = 1024 * 2

#create a passive socket and bind it to 127.0.0.1:6060
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind((TCP_IP, TCP_PORT))

#Take max of one connection
s.listen(1)
print "PRU Speak server at ", TCP_IP + ":" + str(TCP_PORT)

#connection is a one time affair - handles only one client
conn, addr = s.accept()
print 'Connection address:', addr

while True:
        try:
                #blocking wait for data from client
                bs_code = conn.recv(BUFFER_SIZE)
                print "recieved : \n", bs_code, "\n"

                if (not bs_code) or (bs_code == 'EXIT'):
			break #send empty data to shutdown server

		try:
                	ret = pru_speak.execute_instruction(bs_code)
			print "Return value : ", ret
			if ret == []:
				#no return value
				conn.send('\n')
			else:
				#the sends the list of ret values as a string sperated by " "
	        	        conn.send(" ".join(map(str, ret)) + "\n")

	        except Exception as e:
                	print e

        except Exception as e:
                print "error : ", e
                break

conn.close()
<|endoftext|>


<commit_before># -*- cpy-indent-level: 4; indent-tabs-mode: nil -*-
# ex: set expandtab softtabstop=4 shiftwidth=4:
#
# Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016  Contributor
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Export data to systems outside of the broker.
"""

import logging
from collections import defaultdict
from aquilon.exceptions_ import InternalError, RollbackException

LOGGER = logging.getLogger(__name__)

"""
Handle the creation of export events.

Implementers should inherit from this class in order to hook into lifecycle
events within the broker.  To acheive this the class needs to be registered
with the Exporter using either the register_exporter decorator or the
Exporter.register class method (see the latter for the calling convention).

Each of the methods of this class should return either an ExporterTransaction
or a ExporterNotification, depending if external state is being modified, or
only notifications are required.  None may be returned is neithing is required.
"""
class ExportHandler(object):
    """Called when an object is created."""
    def create(self, obj, **kwargs):
        pass

    """Called when an object is updated."""
    def update(self, obj, **kwargs):
        pass

    """Called when an object is deleted."""
    def delete(self, obj, **kwargs):
        pass

    """Called by the Exporter to apply a transaction"""
    def commit(self, action, **kwargs):
        raise NotImplementedError

    """Called by the Exporter to undo a transaction"""
    def rollback(self, action, **kwargs):
        raise NotImplementedError

    """Called by the Exporter to emmit the notifications"""
    def publish(self, notifications, **kwargs):
        raise NotImplementedError

"""
Exporter Transaction

Implementers should inherit from this class in order to manage transactions
with external systems.  The exporter will call the handler commit method
for all of the transactions it has stored.  Any errors will cause the
corrisponding rollback method to be called.  On success the exporter will
the process all of the notifications has stored.
"""
class ExporterTransaction(object):
    pass

"""
Exporter Notification

Like the transaction this should be subclassed by implementers.  The publish
method of the handler will be called with a list of notifications.
"""
class ExporterNotification(object):
    pass

"""
Exporter

The exporter manages transactions and notifications for ExportHandler classes
(see above).  ExportHandler are registed as follow:

    @register_exporter('SomeClass')
    class SomeClassExporter(ExportHandler):
        ...

The constructor of the Exporter will save all of the kwargs and pass them
to all of the various methods of the above classes.
"""
class Exporter(object):
    # Dictoionary of handlers indexed by the class name.  Each entry
    # contains a list of the handlers for the given class type.
    _handlers = defaultdict(list)

    """
    Resgister an ExportHandler for the given class types.  This method should
    be called durning initalisation to register the handler.
    """
    @classmethod
    def register(self, handler, *cls_names):
        if not isinstance(handler, ExportHandler):
            raise InternalError('Handler must be of type ExportHandler')
        for cls_name in cls_names:
            self._handlers[cls_name].append(handler)

    """
    Create a new exporter instance.  The exporter class should be used by
    BrokerCommand's to facilite the export of data to other systems.  The
    contents of kwargs will be passed to the methods of registered
    ExportHandler classes.  If kwargs contrains 'logger' then this will
    be used as the logger object for this class.
    """
    def __init__(self, **kwargs):
        self._kwargs = kwargs
        # Indexed by handler, contains transactions and notifications
        self._actions = defaultdict(lambda: defaultdict(list))
        self._stash = defaultdict(dict)
        self._logger = kwargs.get('logger', LOGGER)

    # Queue a up an action created by an ExportHandler.  Each even should be
    # None or a supported class depending on the type it will be queued up internally.
    def _queue_action(self, handler, action):
        if action is None:
            return
        if isinstance(action, ExporterNotification):
            self._actions[handler]['notifications'].append(action)
        elif isinstance(action, ExporterTransaction):
            self._actions[handler]['transactions'].append(action)
        else:
            raise InternalError('ExportHandler returned an unknown action')

    # Handle the create and delete methods; they are essentially the same
    # so method implements the common functionality.
    def _do_create_or_delete(self, task, obj):
        name = obj.__class__.__name__
        if name in self._handlers:
            for handler in self._handlers[name]:
                action = getattr(handler, task)(obj, **self._kwargs)
                self._queue_action(handler, action)

    """
    Indicate that the passed object has been created
    """
    def create(self, obj):
        self._do_create_or_delete('create', obj)

    """
    Indicate that the passed object has been updated.
    """
    def update(self, obj):
        oname = obj.__class__.__name__
        oid = id(obj)
        if oname in self._handlers:
            for handler in self._handlers[oname]:
                hname = handler.__class__.__name__
                action = handler.update(obj, **self._kwargs)
                self._queue_action(handler, action)

    """
    Indicate that the object has been deleted.
    """
    def delete(self, obj):
        self._do_create_or_delete('delete', obj)

    """
    Commit the changes
    """
    def commit(self):
        for handler, actions in self._actions.items():
            transactions = actions['transactions']
            rollback = actions['rollback']
            for action in transactions:
                # Exceptions in the following are propergated
                handler.commit(action, **self._kwargs)
                rollback.append(action)

    """
    Publish any queued notifications
    """
    def publish(self):
        for handler, actions in self._actions.items():
            notificatons = actions['notifications']
            try:
                handler.publish(notificatons, **self._kwargs)
            except Exception as e:
                self._logger.info("Caught notification error: %s", e)

    """
    Perform any rollback opperations
    """
    def rollback(self):
        undo_failed = False
        for handler, actions in self._actions.items():
            rollback = actions['rollback']
            for action in rollback.reverse():
                try:
                    handler.rollback(action, **self._kwargs)
                except Exception as e:
                    self._logger.warning('Caught exporter rollback error: %s', e)
                    undo_failed = True

        if undo_failed:
            raise RollbackException()

    # The following methods implment a context hander.  This can be used
    # to wrap the commit and rollback methods

    def __enter__(self):
        try:
            self.commit()
        except Exception as e:
            self._logger.info("Exporter commit failed: %s", e)
            self.rollback()

    def __exit__(self, exc_type, exc_val, exc_tb):
        # exc_type is !None if an exception occured
        if exc_type:
            self.rollback()
        else:
            self.publish()

    """
    The following method is called by the ORM when session.flush() is called.
    Internally this calls create, update and delete on the modified objects
    as required.
    """
    def event_after_flush(self, session, flush_context):
        for obj in session.new:
            self.create(obj)
        for obj in session.dirty:
            self.update(obj)
        for obj in session.deleted:
            self.delete(obj)

"""
Register an ExportHandler for the supplied class names.

This decorator takes a list of classes (as arguments) and calls
Exporter.register.
"""
def register_exporter(*class_names):
    return (lambda exporter: Exporter.register(exporter(), *class_names))

# Import all of the plugins
from aquilon.worker.exporters import *

<commit_msg>Drop transactions from exporter framework<commit_after># -*- cpy-indent-level: 4; indent-tabs-mode: nil -*-
# ex: set expandtab softtabstop=4 shiftwidth=4:
#
# Copyright (C) 2008,2009,2010,2011,2012,2013,2014,2015,2016  Contributor
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Export data to systems outside of the broker.
"""

import logging
from collections import defaultdict
from aquilon.exceptions_ import InternalError, RollbackException

LOGGER = logging.getLogger(__name__)

"""
Handle the creation of export events.

Implementers should inherit from this class in order to hook into lifecycle
events within the broker.  To acheive this the class needs to be registered
with the Exporter using either the register_exporter decorator or the
Exporter.register class method (see the latter for the calling convention).
"""
class ExportHandler(object):
    """Called when an object is created."""
    def create(self, obj, **kwargs):
        pass

    """Called when an object is updated."""
    def update(self, obj, **kwargs):
        pass

    """Called when an object is deleted."""
    def delete(self, obj, **kwargs):
        pass

    """Called by the Exporter to emmit the notifications"""
    def publish(self, notifications, **kwargs):
        raise NotImplementedError

"""
Exporter Notification

Like the transaction this should be subclassed by implementers.  The publish
method of the handler will be called with a list of notifications.
"""
class ExporterNotification(object):
    pass

"""
Exporter

The exporter manages notifications for ExportHandler classes
(see above).  ExportHandler are registed as follow:

    @register_exporter('SomeClass')
    class SomeClassExporter(ExportHandler):
        ...

The constructor of the Exporter will save all of the kwargs and pass them
to all of the various methods of the above classes.
"""
class Exporter(object):
    # Dictoionary of handlers indexed by the class name.  Each entry
    # contains a list of the handlers for the given class type.
    _handlers = defaultdict(list)

    """
    Resgister an ExportHandler for the given class types.  This method should
    be called durning initalisation to register the handler.
    """
    @classmethod
    def register(self, handler, *cls_names):
        if not isinstance(handler, ExportHandler):
            raise InternalError('Handler must be of type ExportHandler')
        for cls_name in cls_names:
            self._handlers[cls_name].append(handler)

    """
    Create a new exporter instance.  The exporter class should be used by
    BrokerCommand's to facilite the export of data to other systems.  The
    contents of kwargs will be passed to the methods of registered
    ExportHandler classes.  If kwargs contrains 'logger' then this will
    be used as the logger object for this class.
    """
    def __init__(self, **kwargs):
        self._kwargs = kwargs
        # Indexed by handler, contains notifications
        self._actions = defaultdict(lambda: defaultdict(list))
        self._stash = defaultdict(dict)
        self._logger = kwargs.get('logger', LOGGER)

    # Queue a up an action created by an ExportHandler.  Each even should be
    # None or a supported class depending on the type it will be queued up internally.
    def _queue_action(self, handler, action):
        if action is None:
            return
        if isinstance(action, ExporterNotification):
            self._actions[handler]['notifications'].append(action)
        else:
            raise InternalError('ExportHandler returned an unknown action')

    # Handle the create and delete methods; they are essentially the same
    # so method implements the common functionality.
    def _do_create_or_delete(self, task, obj):
        name = obj.__class__.__name__
        if name in self._handlers:
            for handler in self._handlers[name]:
                action = getattr(handler, task)(obj, **self._kwargs)
                self._queue_action(handler, action)

    """
    Indicate that the passed object has been created
    """
    def create(self, obj):
        self._do_create_or_delete('create', obj)

    """
    Indicate that the passed object has been updated.
    """
    def update(self, obj):
        oname = obj.__class__.__name__
        oid = id(obj)
        if oname in self._handlers:
            for handler in self._handlers[oname]:
                hname = handler.__class__.__name__
                action = handler.update(obj, **self._kwargs)
                self._queue_action(handler, action)

    """
    Indicate that the object has been deleted.
    """
    def delete(self, obj):
        self._do_create_or_delete('delete', obj)

    """
    Publish any queued notifications
    """
    def publish(self):
        for handler, actions in self._actions.items():
            notificatons = actions['notifications']
            try:
                handler.publish(notificatons, **self._kwargs)
            except Exception as e:
                self._logger.info("Caught notification error: %s", e)

    # The following methods implment a context hander.  This can be used
    # to wrap the commit and rollback methods

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        # exc_type is !None if an exception occured
        if not exc_type:
            self.publish()

    """
    The following method is called by the ORM when session.flush() is called.
    Internally this calls create, update and delete on the modified objects
    as required.
    """
    def event_after_flush(self, session, flush_context):
        for obj in session.new:
            self.create(obj)
        for obj in session.dirty:
            self.update(obj)
        for obj in session.deleted:
            self.delete(obj)

"""
Register an ExportHandler for the supplied class names.

This decorator takes a list of classes (as arguments) and calls
Exporter.register.
"""
def register_exporter(*class_names):
    return (lambda exporter: Exporter.register(exporter(), *class_names))

# Import all of the plugins
from aquilon.worker.exporters import *

<|endoftext|>


<commit_before>#!/usr/bin/env Python
"""
Copyright (c) 2013 The Regents of the University of California, AMERICAN INSTITUTES FOR RESEARCH
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
"""
@author Gabe Fierro gt.fierro@berkeley.edu github.com/gtfierro
"""
"""
Performs a basic assignee disambiguation
"""
from collections import defaultdict, deque
import uuid
from string import lowercase as alphabet
import re
import md5
import cPickle as pickle
import alchemy
from collections import Counter
from Levenshtein import jaro_winkler
from alchemy import get_config, match
from alchemy.schema import *
from alchemy.match import commit_inserts, commit_updates
from handlers.xml_util import normalize_utf8
from datetime import datetime
from sqlalchemy.sql import or_
from sqlalchemy.sql.expression import bindparam
from unidecode import unidecode
from tasks import bulk_commit_inserts, bulk_commit_updates
import multiprocessing
import itertools
import sys
import json

config = get_config()

THRESHOLD = config.get("assignee").get("threshold")

uuid_to_object = {}
uuid_to_cleanid = {}
letter_to_cleanid = {}
uuids_by_cleanidletter = defaultdict(list)

grant_uuids = set()
app_uuids = set()
grantsessiongen = alchemy.session_generator(dbtype='grant')
appsessiongen = alchemy.session_generator(dbtype='application')

nodigits = re.compile(r'[a-z ]')
stoplist = ['the','of','and','a','an','at']
substitutions = json.load(open('nber_substitutions.json'))

def isgrant(obj):
    """
    returns True of obj is from Grant table, False if from App table
    """
    return hasattr(obj, 'patent')

def get_cleanid(obj):
    """
    Returns a cleaned string version of the object representation:

    if obj has an organization, uses that. If obj has first/last name,
    uses "firstname|lastname"

    Changes severything to lowercase and removes everything that isn't [a-z ]
    """
    cleanid = ''
    if obj.organization:
        cleanid = obj.organization
    else:
        try:
            cleanid = obj.name_first + obj.name_last
        except:
            cleanid = ''
    cleanid = cleanid.lower()
    cleanid = ' '.join(filter(lambda x:
                        x not in stoplist,
                        cleanid.split()))
    cleanid = ''.join(nodigits.findall(cleanid)).strip()
    for pair in substitutions:
      cleanid = cleanid.replace(pair[0], pair[1])
    return cleanid

def get_similarity(uuid1, uuid2):
    clean1 = uuid_to_cleanid[uuid1]
    clean2 = uuid_to_cleanid[uuid2]
    if clean1 == clean2:
      return 1.0
    return jaro_winkler(clean1,clean2,0.0)

def disambiguate_letter(letter):
    groups = defaultdict(list)
    bucket = uuids_by_cleanidletter[letter]
    uuidsremaining = bucket[:]
    groupkeys = []
    print len(bucket),'raw assignees for letter:', letter
    i = 1
    while True:
        if not uuidsremaining:
            break
        i += 1
        uuid = uuidsremaining.pop()
        if i%10000 == 0:
            print i, datetime.now()
        matcheduuid = False
        for groupkey in groupkeys:
            if get_similarity(uuid, groupkey) >= THRESHOLD:
                groups[groupkey].append(uuid)
                matcheduuid = True
                break
        if matcheduuid: continue
        groups[uuid].append(uuid)
        groupkeys.append(uuid)
    return groups

def create_disambiguated_record_for_block(block):
    grant_assignee_inserts = []
    app_assignee_inserts = []
    patentassignee_inserts = []
    applicationassignee_inserts = []
    grant_rawassignee_updates = []
    app_rawassignee_updates = []
    ra_objs = [uuid_to_object[uuid] for uuid in block]
    # vote on the disambiguated assignee parameters
    freq = defaultdict(Counter)
    param = {}

    for ra in ra_objs:
        for k,v in ra.summarize.items():
            if not v:
                v = ''
            param[k] = v

    if not param.has_key('organization'):
        param['organization'] = ''
    if not param.has_key('type'):
        param['type'] = ''
    if not param.has_key('name_last'):
        param['name_last'] = ''
    if not param.has_key('name_first'):
        param['name_first'] = ''
    if not param.has_key('residence'):
        param['residence'] = ''
    if not param.has_key('nationality'):
        param['nationality'] = ''
    if param.has_key('type'):
        if not param['type'].isdigit():
            param['type'] = ''

    # create persistent identifier
    if param["organization"]:
        param["id"] = md5.md5(unidecode(param["organization"])).hexdigest()
    elif param["name_last"]:
        param["id"] = md5.md5(unidecode(param["name_last"]+param["name_first"])).hexdigest()
    else:
        param["id"] = md5.md5('').hexdigest()

    grant_assignee_inserts.append(param)
    app_assignee_inserts.append(param)

    # inserts for patent_assignee and appliation_assignee tables
    patents = filter(lambda x: x, map(lambda x: getattr(x,'patent_id',None), ra_objs))
    patentassignee_inserts.extend({'patent_id': x, 'assignee_id': param['id']} for x in patents)
    applications = filter(lambda x: x, map(lambda x: getattr(x,'application_id',None), ra_objs))
    applicationassignee_inserts.extend([{'application_id': x, 'assignee_id': param['id']} for x in applications])

    # update statements for rawassignee tables
    for ra in ra_objs:
        if isgrant(ra):
            grant_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
        else:
            app_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
    return grant_assignee_inserts, app_assignee_inserts, patentassignee_inserts, applicationassignee_inserts, grant_rawassignee_updates, app_rawassignee_updates

def run_disambiguation():
    """
    Runs disambiguation algorithm on grant and application assignees from
    the database indicated by lib/alchemy/config
    """
    # retrieve database connections and pull in all assignees from
    # both grant and application databases
    grtsesh = grantsessiongen()
    appsesh = appsessiongen()
    print 'fetching raw assignees',datetime.now()
    rawassignees = list(grtsesh.query(RawAssignee))
    rawassignees.extend(list(appsesh.query(App_RawAssignee)))
    # clear the destination tables
    if alchemy.is_mysql():
        grtsesh.execute('truncate assignee; truncate patent_assignee;')
        appsesh.execute('truncate assignee; truncate application_assignee;')
    else:
        grtsesh.execute('delete from assignee; delete from patent_assignee;')
        appsesh.execute('delete from assignee; delete from patent_assignee;')
    print 'cleaning ids', datetime.now()
    # uses the get_cleanid method to remove undesirable characters and
    # normalize to case and group by first letter
    for ra in rawassignees:
        uuid_to_object[ra.uuid] = ra
        cleanid = get_cleanid(ra)
        uuid_to_cleanid[ra.uuid] = cleanid
        if not cleanid:
          continue
        firstletter = cleanid[0]
        uuids_by_cleanidletter[firstletter].append(ra.uuid)

    print 'disambiguating blocks', datetime.now()
    # disambiguates each of the letter blocks using
    # the list of assignees as a stack and only performing
    # jaro-winkler comparisons on the first item of each block
    allrecords = []
    for letter in alphabet:
        print 'disambiguating','({0})'.format(letter),datetime.now()
        lettergroup = disambiguate_letter(letter)
        print 'got',len(lettergroup),'records'
        print 'creating disambiguated records','({0})'.format(letter),datetime.now()
        allrecords.extend(lettergroup.values())
    # create the attributes for the disambiguated assignee record from the
    # raw records placed into a block in the disambiguation phase
    res = map(create_disambiguated_record_for_block, allrecords)
    mid = itertools.izip(*res)
    grant_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    app_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    patentassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    applicationassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    grant_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))
    app_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))

    # write out the insert counts for each table into a text file
    with open('mid.txt','wb') as f:
      f.write(str(len(grant_assignee_inserts))+'\n')
      f.write(str(len(app_assignee_inserts))+'\n')
      f.write(str(len(patentassignee_inserts))+'\n')
      f.write(str(len(applicationassignee_inserts))+'\n')
      f.write(str(len(grant_rawassignee_updates))+'\n')
      f.write(str(len(app_rawassignee_updates))+'\n')
    # insert disambiguated assignee records
    bulk_commit_inserts(grant_assignee_inserts, Assignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(app_assignee_inserts, App_Assignee.__table__, alchemy.is_mysql(), 20000, 'application')
    # insert patent/assignee link records
    bulk_commit_inserts(patentassignee_inserts, patentassignee, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(applicationassignee_inserts, applicationassignee, alchemy.is_mysql(), 20000, 'application')
    # update rawassignees with their disambiguated record
    bulk_commit_updates('assignee_id', grant_rawassignee_updates, RawAssignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_updates('assignee_id', app_rawassignee_updates, App_RawAssignee.__table__, alchemy.is_mysql(), 20000, 'application')

if __name__=='__main__':
  run_disambiguation()
<commit_msg>Remove residence and nationality from assignee disambiguation<commit_after>#!/usr/bin/env Python
"""
Copyright (c) 2013 The Regents of the University of California, AMERICAN INSTITUTES FOR RESEARCH
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
"""
@author Gabe Fierro gt.fierro@berkeley.edu github.com/gtfierro
"""
"""
Performs a basic assignee disambiguation
"""
from collections import defaultdict, deque
import uuid
from string import lowercase as alphabet
import re
import md5
import cPickle as pickle
import alchemy
from collections import Counter
from Levenshtein import jaro_winkler
from alchemy import get_config, match
from alchemy.schema import *
from alchemy.match import commit_inserts, commit_updates
from handlers.xml_util import normalize_utf8
from datetime import datetime
from sqlalchemy.sql import or_
from sqlalchemy.sql.expression import bindparam
from unidecode import unidecode
from tasks import bulk_commit_inserts, bulk_commit_updates
import multiprocessing
import itertools
import sys
import json

config = get_config()

THRESHOLD = config.get("assignee").get("threshold")

uuid_to_object = {}
uuid_to_cleanid = {}
letter_to_cleanid = {}
uuids_by_cleanidletter = defaultdict(list)

grant_uuids = set()
app_uuids = set()
grantsessiongen = alchemy.session_generator(dbtype='grant')
appsessiongen = alchemy.session_generator(dbtype='application')

nodigits = re.compile(r'[a-z ]')
stoplist = ['the','of','and','a','an','at']
substitutions = json.load(open('nber_substitutions.json'))

def isgrant(obj):
    """
    returns True of obj is from Grant table, False if from App table
    """
    return hasattr(obj, 'patent')

def get_cleanid(obj):
    """
    Returns a cleaned string version of the object representation:

    if obj has an organization, uses that. If obj has first/last name,
    uses "firstname|lastname"

    Changes severything to lowercase and removes everything that isn't [a-z ]
    """
    cleanid = ''
    if obj.organization:
        cleanid = obj.organization
    else:
        try:
            cleanid = obj.name_first + obj.name_last
        except:
            cleanid = ''
    cleanid = cleanid.lower()
    cleanid = ' '.join(filter(lambda x:
                        x not in stoplist,
                        cleanid.split()))
    cleanid = ''.join(nodigits.findall(cleanid)).strip()
    for pair in substitutions:
      cleanid = cleanid.replace(pair[0], pair[1])
    return cleanid

def get_similarity(uuid1, uuid2):
    clean1 = uuid_to_cleanid[uuid1]
    clean2 = uuid_to_cleanid[uuid2]
    if clean1 == clean2:
      return 1.0
    return jaro_winkler(clean1,clean2,0.0)

def disambiguate_letter(letter):
    groups = defaultdict(list)
    bucket = uuids_by_cleanidletter[letter]
    uuidsremaining = bucket[:]
    groupkeys = []
    print len(bucket),'raw assignees for letter:', letter
    i = 1
    while True:
        if not uuidsremaining:
            break
        i += 1
        uuid = uuidsremaining.pop()
        if i%10000 == 0:
            print i, datetime.now()
        matcheduuid = False
        for groupkey in groupkeys:
            if get_similarity(uuid, groupkey) >= THRESHOLD:
                groups[groupkey].append(uuid)
                matcheduuid = True
                break
        if matcheduuid: continue
        groups[uuid].append(uuid)
        groupkeys.append(uuid)
    return groups

def create_disambiguated_record_for_block(block):
    grant_assignee_inserts = []
    app_assignee_inserts = []
    patentassignee_inserts = []
    applicationassignee_inserts = []
    grant_rawassignee_updates = []
    app_rawassignee_updates = []
    ra_objs = [uuid_to_object[uuid] for uuid in block]
    # vote on the disambiguated assignee parameters
    freq = defaultdict(Counter)
    param = {}

    for ra in ra_objs:
        for k,v in ra.summarize.items():
            if not v:
                v = ''
            param[k] = v

    if not param.has_key('organization'):
        param['organization'] = ''
    if not param.has_key('type'):
        param['type'] = ''
    if not param.has_key('name_last'):
        param['name_last'] = ''
    if not param.has_key('name_first'):
        param['name_first'] = ''
    if param.has_key('type'):
        if not param['type'].isdigit():
            param['type'] = ''

    # create persistent identifier
    if param["organization"]:
        param["id"] = md5.md5(unidecode(param["organization"])).hexdigest()
    elif param["name_last"]:
        param["id"] = md5.md5(unidecode(param["name_last"]+param["name_first"])).hexdigest()
    else:
        param["id"] = md5.md5('').hexdigest()

    grant_assignee_inserts.append(param)
    app_assignee_inserts.append(param)

    # inserts for patent_assignee and appliation_assignee tables
    patents = filter(lambda x: x, map(lambda x: getattr(x,'patent_id',None), ra_objs))
    patentassignee_inserts.extend({'patent_id': x, 'assignee_id': param['id']} for x in patents)
    applications = filter(lambda x: x, map(lambda x: getattr(x,'application_id',None), ra_objs))
    applicationassignee_inserts.extend([{'application_id': x, 'assignee_id': param['id']} for x in applications])

    # update statements for rawassignee tables
    for ra in ra_objs:
        if isgrant(ra):
            grant_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
        else:
            app_rawassignee_updates.append({'pk': ra.uuid, 'update': param['id']})
    return grant_assignee_inserts, app_assignee_inserts, patentassignee_inserts, applicationassignee_inserts, grant_rawassignee_updates, app_rawassignee_updates

def run_disambiguation():
    """
    Runs disambiguation algorithm on grant and application assignees from
    the database indicated by lib/alchemy/config
    """
    # retrieve database connections and pull in all assignees from
    # both grant and application databases
    grtsesh = grantsessiongen()
    appsesh = appsessiongen()
    print 'fetching raw assignees',datetime.now()
    rawassignees = list(grtsesh.query(RawAssignee))
    rawassignees.extend(list(appsesh.query(App_RawAssignee)))
    # clear the destination tables
    if alchemy.is_mysql():
        grtsesh.execute('truncate assignee; truncate patent_assignee;')
        appsesh.execute('truncate assignee; truncate application_assignee;')
    else:
        grtsesh.execute('delete from assignee; delete from patent_assignee;')
        appsesh.execute('delete from assignee; delete from patent_assignee;')
    print 'cleaning ids', datetime.now()
    # uses the get_cleanid method to remove undesirable characters and
    # normalize to case and group by first letter
    for ra in rawassignees:
        uuid_to_object[ra.uuid] = ra
        cleanid = get_cleanid(ra)
        uuid_to_cleanid[ra.uuid] = cleanid
        if not cleanid:
          continue
        firstletter = cleanid[0]
        uuids_by_cleanidletter[firstletter].append(ra.uuid)

    print 'disambiguating blocks', datetime.now()
    # disambiguates each of the letter blocks using
    # the list of assignees as a stack and only performing
    # jaro-winkler comparisons on the first item of each block
    allrecords = []
    for letter in alphabet:
        print 'disambiguating','({0})'.format(letter),datetime.now()
        lettergroup = disambiguate_letter(letter)
        print 'got',len(lettergroup),'records'
        print 'creating disambiguated records','({0})'.format(letter),datetime.now()
        allrecords.extend(lettergroup.values())
    # create the attributes for the disambiguated assignee record from the
    # raw records placed into a block in the disambiguation phase
    res = map(create_disambiguated_record_for_block, allrecords)
    mid = itertools.izip(*res)
    grant_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    app_assignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    patentassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    applicationassignee_inserts = list(itertools.chain.from_iterable(mid.next()))
    grant_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))
    app_rawassignee_updates = list(itertools.chain.from_iterable(mid.next()))

    # write out the insert counts for each table into a text file
    with open('mid.txt','wb') as f:
      f.write(str(len(grant_assignee_inserts))+'\n')
      f.write(str(len(app_assignee_inserts))+'\n')
      f.write(str(len(patentassignee_inserts))+'\n')
      f.write(str(len(applicationassignee_inserts))+'\n')
      f.write(str(len(grant_rawassignee_updates))+'\n')
      f.write(str(len(app_rawassignee_updates))+'\n')
    # insert disambiguated assignee records
    bulk_commit_inserts(grant_assignee_inserts, Assignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(app_assignee_inserts, App_Assignee.__table__, alchemy.is_mysql(), 20000, 'application')
    # insert patent/assignee link records
    bulk_commit_inserts(patentassignee_inserts, patentassignee, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_inserts(applicationassignee_inserts, applicationassignee, alchemy.is_mysql(), 20000, 'application')
    # update rawassignees with their disambiguated record
    bulk_commit_updates('assignee_id', grant_rawassignee_updates, RawAssignee.__table__, alchemy.is_mysql(), 20000, 'grant')
    bulk_commit_updates('assignee_id', app_rawassignee_updates, App_RawAssignee.__table__, alchemy.is_mysql(), 20000, 'application')

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


<commit_before>"""
This module implements corfunc
"""
import warnings
import numpy as np
from scipy.optimize import curve_fit
from scipy.interpolate import interp1d
from scipy.fftpack import dct
from scipy.signal import argrelextrema
from numpy.linalg import lstsq
from sas.sascalc.dataloader.data_info import Data1D
from sas.sascalc.corfunc.transform_thread import FourierThread
from sas.sascalc.corfunc.transform_thread import HilbertThread

class CorfuncCalculator(object):

    class _Interpolator(object):
        """
        Interpolates between curve f and curve g over the range start:stop and
        caches the result of the function when it's called

        :param f: The first curve to interpolate
        :param g: The second curve to interpolate
        :param start: The value at which to start the interpolation
        :param stop: The value at which to stop the interpolation
        """
        def __init__(self, f, g, start, stop):
            self.f = f
            self.g = g
            self.start = start
            self.stop = stop
            self._lastx = []
            self._lasty = []

        def __call__(self, x):
            # If input is a single number, evaluate the function at that number
            # and return a single number
            if type(x) == float or type(x) == int:
                return self._smoothed_function(np.array([x]))[0]
            # If input is a list, and is different to the last input, evaluate
            # the function at each point. If the input is the same as last time
            # the function was called, return the result that was calculated
            # last time instead of explicity evaluating the function again.
            elif self._lastx == [] or x.tolist() != self._lastx.tolist():
                self._lasty = self._smoothed_function(x)
                self._lastx = x
            return self._lasty

        def _smoothed_function(self,x):
            ys = np.zeros(x.shape)
            ys[x <= self.start] = self.f(x[x <= self.start])
            ys[x >= self.stop] = self.g(x[x >= self.stop])
            with warnings.catch_warnings():
                # Ignore divide by zero error
                warnings.simplefilter('ignore')
                h = 1/(1+(x-self.stop)**2/(self.start-x)**2)
            mask = np.logical_and(x > self.start, x < self.stop)
            ys[mask] = h[mask]*self.g(x[mask])+(1-h[mask])*self.f(x[mask])
            return ys

    def __init__(self, data=None, lowerq=None, upperq=None, scale=1):
        """
        Initialize the class.

        :param data: Data of the type DataLoader.Data1D
        :param lowerq: The Q value to use as the boundary for
            Guinier extrapolation
        :param upperq: A tuple of the form (lower, upper).
            Values between lower and upper will be used for Porod extrapolation
        :param scale: Scaling factor for I(q)
        """
        self._data = None
        self.set_data(data, scale)
        self.lowerq = lowerq
        self.upperq = upperq
        self.background = self.compute_background()
        self._transform_thread = None

    def set_data(self, data, scale=1):
        """
        Prepares the data for analysis

        :return: new_data = data * scale - background
        """
        if data is None:
            return
        # Only process data of the class Data1D
        if not issubclass(data.__class__, Data1D):
            raise ValueError, "Data must be of the type DataLoader.Data1D"

        # Prepare the data
        new_data = Data1D(x=data.x, y=data.y)
        new_data *= scale

        # Ensure the errors are set correctly
        if new_data.dy is None or len(new_data.x) != len(new_data.dy) or \
            (min(new_data.dy) == 0 and max(new_data.dy) == 0):
            new_data.dy = np.ones(len(new_data.x))

        self._data = new_data

    def compute_background(self, upperq=None):
        """
        Compute the background level from the Porod region of the data
        """
        if self._data is None: return 0
        elif upperq is None and self.upperq is not None: upperq = self.upperq
        elif upperq is None and self.upperq is None: return 0
        q = self._data.x
        mask = np.logical_and(q > upperq[0], q < upperq[1])
        _, _, bg = self._fit_porod(q[mask], self._data.y[mask])

        return bg

    def compute_extrapolation(self):
        """
        Extrapolate and interpolate scattering data

        :return: The extrapolated data
        """
        q = self._data.x
        iq = self._data.y

        params, s2 = self._fit_data(q, iq)
        qs = np.arange(0, q[-1]*100, (q[1]-q[0]))
        iqs = s2(qs)

        extrapolation = Data1D(qs, iqs)

        return params, extrapolation, s2

    def compute_transform(self, extrapolation, trans_type, background=None,
        completefn=None, updatefn=None):
        """
        Transform an extrapolated scattering curve into a correlation function.

        :param extrapolation: The extrapolated data
        :param background: The background value (if not provided, previously
            calculated value will be used)
        :param extrap_fn: A callable function representing the extraoplated data
        :param completefn: The function to call when the transform calculation
            is complete
        :param updatefn: The function to call to update the GUI with the status
            of the transform calculation
        :return: The transformed data
        """
        if self._transform_thread is not None:
            if self._transform_thread.isrunning(): return

        if background is None: background = self.background

        if trans_type == 'fourier':
            self._transform_thread = FourierThread(self._data, extrapolation,
            background, completefn=completefn,
            updatefn=updatefn)
        elif trans_type == 'hilbert':
            self._transform_thread = HilbertThread(self._data, extrapolation,
            background, completefn=completefn, updatefn=updatefn)
        else:
            err = ("Incorrect transform type supplied, must be 'fourier'",
                " or 'hilbert'")
            raise ValueError, err

        self._transform_thread.queue()

    def transform_isrunning(self):
        if self._transform_thread is None: return False
        return self._transform_thread.isrunning()

    def stop_transform(self):
        if self._transform_thread.isrunning():
            self._transform_thread.stop()

    def extract_parameters(self, transformed_data):
        """
        Extract the interesting measurements from a correlation function
        :param transformed_data: Fourier transformation of the
            extrapolated data
        """
        # Calculate indexes of maxima and minima
        x = transformed_data.x
        y = transformed_data.y
        maxs = argrelextrema(y, np.greater)[0]
        mins = argrelextrema(y, np.less)[0]

        # If there are no maxima, return None
        if len(maxs) == 0:
            return None

        GammaMin = y[mins[0]]  # The value at the first minimum

        ddy = (y[:-2]+y[2:]-2*y[1:-1])/(x[2:]-x[:-2])**2  # 2nd derivative of y
        dy = (y[2:]-y[:-2])/(x[2:]-x[:-2])  # 1st derivative of y
        # Find where the second derivative goes to zero
        zeros = argrelextrema(np.abs(ddy), np.less)[0]
        # locate the first inflection point
        linear_point = zeros[0]

        # Try to calculate slope around linear_point using 80 data points
        lower = linear_point - 40
        upper = linear_point + 40

        # If too few data points to the left, use linear_point*2 data points
        if lower < 0:
            lower = 0
            upper = linear_point * 2
        # If too few to right, use 2*(dy.size - linear_point) data points
        elif upper > len(dy):
            upper = len(dy)
            width = len(dy) - linear_point
            lower = 2*linear_point - dy.size

        m = np.mean(dy[lower:upper])  # Linear slope
        b = y[1:-1][linear_point]-m*x[1:-1][linear_point]  # Linear intercept

        Lc = (GammaMin-b)/m  # Hard block thickness

        # Find the data points where the graph is linear to within 1%
        mask = np.where(np.abs((y-(m*x+b))/y) < 0.01)[0]
        if len(mask) == 0:  # Return garbage for bad fits
            return { 'max': self._round_sig_figs(x[maxs[0]], 6) }
        dtr = x[mask[0]]  # Beginning of Linear Section
        d0 = x[mask[-1]]  # End of Linear Section
        GammaMax = y[mask[-1]]
        A = np.abs(GammaMin/GammaMax)  # Normalized depth of minimum

        params = {
            'max': x[maxs[0]],
            'dtr': dtr,
            'Lc': Lc,
            'd0': d0,
            'A': A,
            'fill': Lc/x[maxs[0]]
        }

        return params

    def _porod(self, q, K, sigma, bg):
        """Equation for the Porod region of the data"""
        return bg + (K*q**(-4))*np.exp(-q**2*sigma**2)

    def _fit_guinier(self, q, iq):
        """Fit the Guinier region of the curve"""
        A = np.vstack([q**2, np.ones(q.shape)]).T
        return lstsq(A, np.log(iq))

    def _fit_porod(self, q, iq):
        """Fit the Porod region of the curve"""
        fitp = curve_fit(lambda q, k, sig, bg: self._porod(q, k, sig, bg)*q**2,
                         q, iq*q**2, bounds=([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]))[0]
        k, sigma, bg = fitp
        return k, sigma, bg

    def _fit_data(self, q, iq):
        """
        Given a data set, extrapolate out to large q with Porod and
        to q=0 with Guinier
        """
        mask = np.logical_and(q > self.upperq[0], q < self.upperq[1])

        # Returns an array where the 1st and 2nd elements are the values of k
        # and sigma for the best-fit Porod function
        k, sigma, _ = self._fit_porod(q[mask], iq[mask])
        bg = self.background

        # Smooths between the best-fit porod function and the data to produce a
        # better fitting curve
        data = interp1d(q, iq)
        s1 = self._Interpolator(data,
            lambda x: self._porod(x, k, sigma, bg), self.upperq[0], q[-1])

        mask = np.logical_and(q < self.lowerq, 0 < q)

        # Returns parameters for the best-fit Guinier function
        g = self._fit_guinier(q[mask], iq[mask])[0]

        # Smooths between the best-fit Guinier function and the Porod curve
        s2 = self._Interpolator((lambda x: (np.exp(g[1]+g[0]*x**2))), s1, q[0],
            self.lowerq)

        params = {'A': g[1], 'B': g[0], 'K': k, 'sigma': sigma}

        return params, s2
<commit_msg>Added a comment to the source code<commit_after>"""
This module implements corfunc
"""
import warnings
import numpy as np
from scipy.optimize import curve_fit
from scipy.interpolate import interp1d
from scipy.fftpack import dct
from scipy.signal import argrelextrema
from numpy.linalg import lstsq
from sas.sascalc.dataloader.data_info import Data1D
from sas.sascalc.corfunc.transform_thread import FourierThread
from sas.sascalc.corfunc.transform_thread import HilbertThread

class CorfuncCalculator(object):

    class _Interpolator(object):
        """
        Interpolates between curve f and curve g over the range start:stop and
        caches the result of the function when it's called

        :param f: The first curve to interpolate
        :param g: The second curve to interpolate
        :param start: The value at which to start the interpolation
        :param stop: The value at which to stop the interpolation
        """
        def __init__(self, f, g, start, stop):
            self.f = f
            self.g = g
            self.start = start
            self.stop = stop
            self._lastx = []
            self._lasty = []

        def __call__(self, x):
            # If input is a single number, evaluate the function at that number
            # and return a single number
            if type(x) == float or type(x) == int:
                return self._smoothed_function(np.array([x]))[0]
            # If input is a list, and is different to the last input, evaluate
            # the function at each point. If the input is the same as last time
            # the function was called, return the result that was calculated
            # last time instead of explicity evaluating the function again.
            elif self._lastx == [] or x.tolist() != self._lastx.tolist():
                self._lasty = self._smoothed_function(x)
                self._lastx = x
            return self._lasty

        def _smoothed_function(self,x):
            ys = np.zeros(x.shape)
            ys[x <= self.start] = self.f(x[x <= self.start])
            ys[x >= self.stop] = self.g(x[x >= self.stop])
            with warnings.catch_warnings():
                # Ignore divide by zero error
                warnings.simplefilter('ignore')
                h = 1/(1+(x-self.stop)**2/(self.start-x)**2)
            mask = np.logical_and(x > self.start, x < self.stop)
            ys[mask] = h[mask]*self.g(x[mask])+(1-h[mask])*self.f(x[mask])
            return ys

    def __init__(self, data=None, lowerq=None, upperq=None, scale=1):
        """
        Initialize the class.

        :param data: Data of the type DataLoader.Data1D
        :param lowerq: The Q value to use as the boundary for
            Guinier extrapolation
        :param upperq: A tuple of the form (lower, upper).
            Values between lower and upper will be used for Porod extrapolation
        :param scale: Scaling factor for I(q)
        """
        self._data = None
        self.set_data(data, scale)
        self.lowerq = lowerq
        self.upperq = upperq
        self.background = self.compute_background()
        self._transform_thread = None

    def set_data(self, data, scale=1):
        """
        Prepares the data for analysis

        :return: new_data = data * scale - background
        """
        if data is None:
            return
        # Only process data of the class Data1D
        if not issubclass(data.__class__, Data1D):
            raise ValueError, "Data must be of the type DataLoader.Data1D"

        # Prepare the data
        new_data = Data1D(x=data.x, y=data.y)
        new_data *= scale

        # Ensure the errors are set correctly
        if new_data.dy is None or len(new_data.x) != len(new_data.dy) or \
            (min(new_data.dy) == 0 and max(new_data.dy) == 0):
            new_data.dy = np.ones(len(new_data.x))

        self._data = new_data

    def compute_background(self, upperq=None):
        """
        Compute the background level from the Porod region of the data
        """
        if self._data is None: return 0
        elif upperq is None and self.upperq is not None: upperq = self.upperq
        elif upperq is None and self.upperq is None: return 0
        q = self._data.x
        mask = np.logical_and(q > upperq[0], q < upperq[1])
        _, _, bg = self._fit_porod(q[mask], self._data.y[mask])

        return bg

    def compute_extrapolation(self):
        """
        Extrapolate and interpolate scattering data

        :return: The extrapolated data
        """
        q = self._data.x
        iq = self._data.y

        params, s2 = self._fit_data(q, iq)
        # Extrapolate to 100*Qmax in experimental data
        qs = np.arange(0, q[-1]*100, (q[1]-q[0]))
        iqs = s2(qs)

        extrapolation = Data1D(qs, iqs)

        return params, extrapolation, s2

    def compute_transform(self, extrapolation, trans_type, background=None,
        completefn=None, updatefn=None):
        """
        Transform an extrapolated scattering curve into a correlation function.

        :param extrapolation: The extrapolated data
        :param background: The background value (if not provided, previously
            calculated value will be used)
        :param extrap_fn: A callable function representing the extraoplated data
        :param completefn: The function to call when the transform calculation
            is complete
        :param updatefn: The function to call to update the GUI with the status
            of the transform calculation
        :return: The transformed data
        """
        if self._transform_thread is not None:
            if self._transform_thread.isrunning(): return

        if background is None: background = self.background

        if trans_type == 'fourier':
            self._transform_thread = FourierThread(self._data, extrapolation,
            background, completefn=completefn,
            updatefn=updatefn)
        elif trans_type == 'hilbert':
            self._transform_thread = HilbertThread(self._data, extrapolation,
            background, completefn=completefn, updatefn=updatefn)
        else:
            err = ("Incorrect transform type supplied, must be 'fourier'",
                " or 'hilbert'")
            raise ValueError, err

        self._transform_thread.queue()

    def transform_isrunning(self):
        if self._transform_thread is None: return False
        return self._transform_thread.isrunning()

    def stop_transform(self):
        if self._transform_thread.isrunning():
            self._transform_thread.stop()

    def extract_parameters(self, transformed_data):
        """
        Extract the interesting measurements from a correlation function
        :param transformed_data: Fourier transformation of the
            extrapolated data
        """
        # Calculate indexes of maxima and minima
        x = transformed_data.x
        y = transformed_data.y
        maxs = argrelextrema(y, np.greater)[0]
        mins = argrelextrema(y, np.less)[0]

        # If there are no maxima, return None
        if len(maxs) == 0:
            return None

        GammaMin = y[mins[0]]  # The value at the first minimum

        ddy = (y[:-2]+y[2:]-2*y[1:-1])/(x[2:]-x[:-2])**2  # 2nd derivative of y
        dy = (y[2:]-y[:-2])/(x[2:]-x[:-2])  # 1st derivative of y
        # Find where the second derivative goes to zero
        zeros = argrelextrema(np.abs(ddy), np.less)[0]
        # locate the first inflection point
        linear_point = zeros[0]

        # Try to calculate slope around linear_point using 80 data points
        lower = linear_point - 40
        upper = linear_point + 40

        # If too few data points to the left, use linear_point*2 data points
        if lower < 0:
            lower = 0
            upper = linear_point * 2
        # If too few to right, use 2*(dy.size - linear_point) data points
        elif upper > len(dy):
            upper = len(dy)
            width = len(dy) - linear_point
            lower = 2*linear_point - dy.size

        m = np.mean(dy[lower:upper])  # Linear slope
        b = y[1:-1][linear_point]-m*x[1:-1][linear_point]  # Linear intercept

        Lc = (GammaMin-b)/m  # Hard block thickness

        # Find the data points where the graph is linear to within 1%
        mask = np.where(np.abs((y-(m*x+b))/y) < 0.01)[0]
        if len(mask) == 0:  # Return garbage for bad fits
            return { 'max': self._round_sig_figs(x[maxs[0]], 6) }
        dtr = x[mask[0]]  # Beginning of Linear Section
        d0 = x[mask[-1]]  # End of Linear Section
        GammaMax = y[mask[-1]]
        A = np.abs(GammaMin/GammaMax)  # Normalized depth of minimum

        params = {
            'max': x[maxs[0]],
            'dtr': dtr,
            'Lc': Lc,
            'd0': d0,
            'A': A,
            'fill': Lc/x[maxs[0]]
        }

        return params

    def _porod(self, q, K, sigma, bg):
        """Equation for the Porod region of the data"""
        return bg + (K*q**(-4))*np.exp(-q**2*sigma**2)

    def _fit_guinier(self, q, iq):
        """Fit the Guinier region of the curve"""
        A = np.vstack([q**2, np.ones(q.shape)]).T
        return lstsq(A, np.log(iq))

    def _fit_porod(self, q, iq):
        """Fit the Porod region of the curve"""
        fitp = curve_fit(lambda q, k, sig, bg: self._porod(q, k, sig, bg)*q**2,
                         q, iq*q**2, bounds=([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]))[0]
        k, sigma, bg = fitp
        return k, sigma, bg

    def _fit_data(self, q, iq):
        """
        Given a data set, extrapolate out to large q with Porod and
        to q=0 with Guinier
        """
        mask = np.logical_and(q > self.upperq[0], q < self.upperq[1])

        # Returns an array where the 1st and 2nd elements are the values of k
        # and sigma for the best-fit Porod function
        k, sigma, _ = self._fit_porod(q[mask], iq[mask])
        bg = self.background

        # Smooths between the best-fit porod function and the data to produce a
        # better fitting curve
        data = interp1d(q, iq)
        s1 = self._Interpolator(data,
            lambda x: self._porod(x, k, sigma, bg), self.upperq[0], q[-1])

        mask = np.logical_and(q < self.lowerq, 0 < q)

        # Returns parameters for the best-fit Guinier function
        g = self._fit_guinier(q[mask], iq[mask])[0]

        # Smooths between the best-fit Guinier function and the Porod curve
        s2 = self._Interpolator((lambda x: (np.exp(g[1]+g[0]*x**2))), s1, q[0],
            self.lowerq)

        params = {'A': g[1], 'B': g[0], 'K': k, 'sigma': sigma}

        return params, s2
<|endoftext|>


<commit_before>from twisted.python import log
from twisted.python.log import logging
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.internet import reactor

from vumi.service import Worker, Consumer, Publisher
from vumi.workers.smpp.client import EsmeTransceiverFactory, EsmeTransceiver

import json

import os
os.environ['DJANGO_SETTINGS_MODULE'] = 'vumi.webapp.settings'
from vumi.webapp.api import models
from vumi.webapp.api import forms

class SmppConsumer(Consumer):
    """
    This consumer creates the generic outbound SMPP transport.
    Anything published to the `vumi.smpp` exchange with
    routing key smpp.* (* == single word match, # == zero or more words)
    """
    exchange_name = "vumi"
    exchange_type = "direct"
    durable = True
    auto_delete = False
    queue_name = "sms_send"
    routing_key = "vumi.*"

    def __init__(self, send_callback):
        self.send = send_callback

    def consume_json(self, dictionary):
        log.msg("Consumed JSON %s" % dictionary)
        payload = []
        kwargs = dictionary.get('kwargs')
        if kwargs:
            payload = kwargs.get('payload', [])
        for mess in payload:
            sequence_number = self.send(**mess)
            formdict = {
                    "sent_sms":mess.get("id"),
                    "sequence_number": sequence_number,
                    }
            log.msg("SMPPLinkForm <- %s" % formdict)
            form = forms.SMPPLinkForm(formdict)
            form.save()
        return True

    def consume(self, message):
        if self.consume_json(json.loads(message.content.body)):
            self.ack(message)

class SmppPublisher(Publisher):
    """
    This publisher publishes all incoming SMPP messages to the
    `vumi.smpp` exchange, its default routing key is `smpp.fallback`
    """
    exchange_name = "vumi.smpp"
    exchange_type = "topic"             # -> route based on pattern matching
    routing_key = 'smpp.fallback'       # -> overriden in publish method
    durable = False                     # -> not created at boot
    auto_delete = False                 # -> auto delete if no consumers bound
    delivery_mode = 2                   # -> do not save to disk

    def publish_json(self, dictionary, **kwargs):
        log.msg("Publishing JSON %s with extra args: %s" % (dictionary, kwargs))
        super(SmppPublisher, self).publish_json(dictionary, **kwargs)

class SmppTransport(Worker):
    """
    The SmppTransport
    """

    def startWorker(self):
        log.msg("Starting the SmppTransport")
        # start the Smpp transport
        factory = EsmeTransceiverFactory(
                int(self.config['smpp_increment']),
                int(self.config['smpp_offset']))
        factory.loadDefaults(self.config)
        factory.setLatestSequenceNumber(self.getLatestSequenceNumber())
        factory.setConnectCallback(self.esme_connected)
        factory.setDisconnectCallback(self.esme_disconnected)
        factory.setSubmitSMRespCallback(self.submit_sm_resp)
        factory.setDeliverSMCallback(self.deliver_sm)
        print factory.defaults
        reactor.connectTCP(
                factory.defaults['host'],
                factory.defaults['port'],
                factory)

    def getLatestSequenceNumber(self):
        sequence_number = 0
        try:
            sequence_number = models.SMPPLink.objects.latest().sequence_number
        except Exception, e:
            log.msg("No SMPPLink entries yet")
        return sequence_number

    @inlineCallbacks
    def esme_connected(self, client):
        log.msg("ESME Connected, adding handlers")
        self.esme_client = client
        self.esme_client.set_handler(self)

        # Start the publisher
        self.publisher = yield self.start_publisher(SmppPublisher)
        # Start the consumer, pass along the send_smpp callback for sending
        # back consumed AMQP messages over SMPP.
        self.consumer = yield self.start_consumer(SmppConsumer, self.send_smpp)

    @inlineCallbacks
    def esme_disconnected(self):
        log.msg("ESME Disconnected, stopping consumer")
        stop = yield self.consumer.stop()

    @inlineCallbacks
    def submit_sm_resp(self, *args, **kwargs):
        smpplink = models.SMPPLink.objects \
                .filter(sequence_number=kwargs['sequence_number']) \
                .order_by('-created_at')[:1].get()
        kwargs.update({'sent_sms':smpplink.sent_sms_id})
        log.msg("SMPPRespForm <- %s" % kwargs)
        form = forms.SMPPRespForm(kwargs)
        form.save()
        yield log.msg("SUBMIT SM RESP %s" % (kwargs))

    @inlineCallbacks
    def deliver_sm(self, *args, **kwargs):
        print "#####DELIVERED#####", kwargs

    @inlineCallbacks
    def deliver_sm__(self, *args, **kwargs):
        groupdict = {'title':'reply', 'user':1}
        groupform = forms.SendGroupForm(groupdict)
        if not groupform.is_valid():
            raise FormValidationError(groupform)
        send_group = groupform.save()
        sentdict = {
                'transport_name': 'smpp',
                'from_msisdn': u'27123456789',
                'send_group': send_group.id,
                'user': 1,
                'to_msisdn': kwargs['source_addr'],
                'message': 'You said: "'+kwargs['short_message']+'"'
                }
        sentform = forms.SentSMSForm(sentdict)
        if not sentform.is_valid():
            raise FormValidationError(sentform)
        sent_sms = sentform.save()
        sequence_number = self.send_smpp(
                sent_sms.id,
                sent_sms.to_msisdn,
                sent_sms.message
                )
        linkdict = {
                "sent_sms":sent_sms.id,
                "sequence_number": sequence_number,
                }
        log.msg("SMPPLinkForm <- %s" % linkdict)
        linkform = forms.SMPPLinkForm(linkdict)
        linkform.save()
        yield log.msg("DELIVER SM %s" % (sentdict))

    def send_smpp(self, id, to_msisdn, message, *args, **kwargs):
        print "Sending SMPP, to: %s, message: %s" % (to_msisdn, message)
        sequence_number = self.esme_client.submit_sm(
                short_message = str(message),
                destination_addr = str(to_msisdn),
                )
        return sequence_number

    def sms_callback(self, *args, **kwargs):
        print "Got SMS:", args, kwargs

    def errback(self, *args, **kwargs):
        print "Got Error: ", args, kwargs

    def stopWorker(self):
        log.msg("Stopping the SMPPTransport")

<commit_msg>msisdn validation forwarder<commit_after>from twisted.python import log
from twisted.python.log import logging
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.internet import reactor

from vumi.service import Worker, Consumer, Publisher
from vumi.workers.smpp.client import EsmeTransceiverFactory, EsmeTransceiver

import json

import os
os.environ['DJANGO_SETTINGS_MODULE'] = 'vumi.webapp.settings'
from vumi.webapp.api import models
from vumi.webapp.api import forms

import urllib
import urllib2

class SmppConsumer(Consumer):
    """
    This consumer creates the generic outbound SMPP transport.
    Anything published to the `vumi.smpp` exchange with
    routing key smpp.* (* == single word match, # == zero or more words)
    """
    exchange_name = "vumi"
    exchange_type = "direct"
    durable = True
    auto_delete = False
    queue_name = "sms_send"
    routing_key = "vumi.*"

    def __init__(self, send_callback):
        self.send = send_callback

    def consume_json(self, dictionary):
        log.msg("Consumed JSON %s" % dictionary)
        payload = []
        kwargs = dictionary.get('kwargs')
        if kwargs:
            payload = kwargs.get('payload', [])
        for mess in payload:
            sequence_number = self.send(**mess)
            formdict = {
                    "sent_sms":mess.get("id"),
                    "sequence_number": sequence_number,
                    }
            log.msg("SMPPLinkForm <- %s" % formdict)
            form = forms.SMPPLinkForm(formdict)
            form.save()
        return True

    def consume(self, message):
        if self.consume_json(json.loads(message.content.body)):
            self.ack(message)

class SmppPublisher(Publisher):
    """
    This publisher publishes all incoming SMPP messages to the
    `vumi.smpp` exchange, its default routing key is `smpp.fallback`
    """
    exchange_name = "vumi.smpp"
    exchange_type = "topic"             # -> route based on pattern matching
    routing_key = 'smpp.fallback'       # -> overriden in publish method
    durable = False                     # -> not created at boot
    auto_delete = False                 # -> auto delete if no consumers bound
    delivery_mode = 2                   # -> do not save to disk

    def publish_json(self, dictionary, **kwargs):
        log.msg("Publishing JSON %s with extra args: %s" % (dictionary, kwargs))
        super(SmppPublisher, self).publish_json(dictionary, **kwargs)

class SmppTransport(Worker):
    """
    The SmppTransport
    """

    def startWorker(self):
        log.msg("Starting the SmppTransport")
        # start the Smpp transport
        factory = EsmeTransceiverFactory(
                int(self.config['smpp_increment']),
                int(self.config['smpp_offset']))
        factory.loadDefaults(self.config)
        factory.setLatestSequenceNumber(self.getLatestSequenceNumber())
        factory.setConnectCallback(self.esme_connected)
        factory.setDisconnectCallback(self.esme_disconnected)
        factory.setSubmitSMRespCallback(self.submit_sm_resp)
        factory.setDeliverSMCallback(self.deliver_sm)
        print factory.defaults
        reactor.connectTCP(
                factory.defaults['host'],
                factory.defaults['port'],
                factory)

    def getLatestSequenceNumber(self):
        sequence_number = 0
        try:
            sequence_number = models.SMPPLink.objects.latest().sequence_number
        except Exception, e:
            log.msg("No SMPPLink entries yet")
        return sequence_number

    @inlineCallbacks
    def esme_connected(self, client):
        log.msg("ESME Connected, adding handlers")
        self.esme_client = client
        self.esme_client.set_handler(self)

        # Start the publisher
        self.publisher = yield self.start_publisher(SmppPublisher)
        # Start the consumer, pass along the send_smpp callback for sending
        # back consumed AMQP messages over SMPP.
        self.consumer = yield self.start_consumer(SmppConsumer, self.send_smpp)

    @inlineCallbacks
    def esme_disconnected(self):
        log.msg("ESME Disconnected, stopping consumer")
        stop = yield self.consumer.stop()

    @inlineCallbacks
    def submit_sm_resp(self, *args, **kwargs):
        smpplink = models.SMPPLink.objects \
                .filter(sequence_number=kwargs['sequence_number']) \
                .order_by('-created_at')[:1].get()
        kwargs.update({'sent_sms':smpplink.sent_sms_id})
        log.msg("SMPPRespForm <- %s" % kwargs)
        form = forms.SMPPRespForm(kwargs)
        form.save()
        yield log.msg("SUBMIT SM RESP %s" % (kwargs))

    @inlineCallbacks
    def deliver_sm(self, *args, **kwargs):
        url = "http://localhost:8080/"
        params = {'json' : '{"route":"%s", "msisdn":"%s", "message":"%s"}' % (
            kwargs.get('source_addr'),
            kwargs.get('source_addr'),
            kwargs.get('short_message')
            )}
        data = urllib.urlencode(params)
        req = urllib2.Request(url, data)
        resp = urllib2.urlopen(req)
        yield log.msg("DELIVER SM %s" % (json.dumps(kwargs)))

    @inlineCallbacks
    def deliver_sm__(self, *args, **kwargs):
        groupdict = {'title':'reply', 'user':1}
        groupform = forms.SendGroupForm(groupdict)
        if not groupform.is_valid():
            raise FormValidationError(groupform)
        send_group = groupform.save()
        sentdict = {
                'transport_name': 'smpp',
                'from_msisdn': u'27123456789',
                'send_group': send_group.id,
                'user': 1,
                'to_msisdn': kwargs['source_addr'],
                'message': 'You said: "'+kwargs['short_message']+'"'
                }
        sentform = forms.SentSMSForm(sentdict)
        if not sentform.is_valid():
            raise FormValidationError(sentform)
        sent_sms = sentform.save()
        sequence_number = self.send_smpp(
                sent_sms.id,
                sent_sms.to_msisdn,
                sent_sms.message
                )
        linkdict = {
                "sent_sms":sent_sms.id,
                "sequence_number": sequence_number,
                }
        log.msg("SMPPLinkForm <- %s" % linkdict)
        linkform = forms.SMPPLinkForm(linkdict)
        linkform.save()
        yield log.msg("DELIVER SM %s" % (sentdict))

    def send_smpp(self, id, to_msisdn, message, *args, **kwargs):
        print "Sending SMPP, to: %s, message: %s" % (to_msisdn, message)
        sequence_number = self.esme_client.submit_sm(
                short_message = str(message),
                destination_addr = str(to_msisdn),
                )
        return sequence_number

    def sms_callback(self, *args, **kwargs):
        print "Got SMS:", args, kwargs

    def errback(self, *args, **kwargs):
        print "Got Error: ", args, kwargs

    def stopWorker(self):
        log.msg("Stopping the SMPPTransport")

<|endoftext|>


<commit_before>#!/usr/bin/env python

from ..common import *
from ..extractor import VideoExtractor

class Bilibili(VideoExtractor):
    name = "Bilibili"

    # Bilibili media encoding options, in descending quality order.
    stream_types = [
        {'id': 'flv_p60', 'quality': 116, 'audio_quality': 30280,
         'container': 'MP4', 'video_resolution': '1080p', 'desc': '高清 1080P60'},
        # 'id': 'hdflv2', 'quality': 112?
        {'id': 'flv', 'quality': 80, 'audio_quality': 30280,
         'container': 'MP4', 'video_resolution': '1080p', 'desc': '高清 1080P'},
        {'id': 'flv720_p60', 'quality': 74, 'audio_quality': 30280,
         'container': 'MP4', 'video_resolution': '720p', 'desc': '高清 720P60'},
        {'id': 'flv720', 'quality': 64, 'audio_quality': 30280,
         'container': 'MP4', 'video_resolution': '720p', 'desc': '高清 720P'},
        {'id': 'flv480', 'quality': 32, 'audio_quality': 30280,
         'container': 'MP4', 'video_resolution': '480p', 'desc': '清晰 480P'},
        {'id': 'flv360', 'quality': 16, 'audio_quality': 30216,
         'container': 'MP4', 'video_resolution': '360p', 'desc': '流畅 360P'},

        {'id': 'default', 'quality': 0}
    ]

    @staticmethod
    def bilibili_headers(referer=None, cookie=None):
        # a reasonable UA
        ua = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.84 Safari/537.36'
        headers = {'User-Agent': ua}
        if referer is not None:
            headers.update({'Referer': referer})
        if cookie is not None:
            headers.update({'Cookie': cookie})
        return headers

    @staticmethod
    def bilibili_bangumi_api(avid, cid, ep_id):
        return 'https://api.bilibili.com/pgc/player/web/playurl?avid=%s&cid=%s&qn=0&type=&otype=json&ep_id=%s&fnver=0&fnval=16' % (avid, cid, ep_id)

    def prepare(self, **kwargs):
        self.stream_qualities = {s['quality']: s for s in self.stream_types}

        html_content = get_content(self.url, headers=self.bilibili_headers())
        #self.title = match1(html_content,
        #                    r'<h1 title="([^"]+)"')

        # regular av
        if re.match(r'https?://(www)?\.bilibili\.com/video/av(\d+)', self.url):
            playinfo_text = match1(html_content, r'__playinfo__=(.*?)</script>