<commit_before>
from django.db import connection, transaction, models
from django.db.backends.util import truncate_name
from django.dispatch import dispatcher
from django.conf import settings

class DatabaseOperations(object):

    """
    Generic SQL implementation of the DatabaseOperations.
    Some of this code comes from Django Evolution.
    """

    def __init__(self):
        self.debug = False
        self.deferred_sql = []

    def execute(self, sql, params=[]):
        """
        Executes the given SQL statement, with optional parameters.
        If the instance's debug attribute is True, prints out what it executes.
        """
        cursor = connection.cursor()
        if self.debug:
            print "   = %s" % sql, params
        cursor.execute(sql, params)
        try:
            return cursor.fetchall()
        except:
            return []

    def add_deferred_sql(self, sql):
        """
        Add a SQL statement to the deferred list, that won't be executed until
        this instance's execute_deferred_sql method is run.
        """
        self.deferred_sql.append(sql)

    def execute_deferred_sql(self):
        """
        Executes all deferred SQL, resetting the deferred_sql list
        """
        for sql in self.deferred_sql:
            self.execute(sql)

        self.deferred_sql = []

    def create_table(self, table_name, fields):
        """
        Creates the table 'table_name'. 'fields' is a tuple of fields,
        each repsented by a 2-part tuple of field name and a
        django.db.models.fields.Field object
        """
        qn = connection.ops.quote_name
        columns = [
            self.column_sql(table_name, field_name, field)
            for field_name, field in fields
        ]

        self.execute('CREATE TABLE %s (%s);' % (table_name, ', '.join([col for col in columns])))

    add_table = create_table # Alias for consistency's sake

    def rename_table(self, old_table_name, table_name):
        """
        Renames the table 'old_table_name' to 'table_name'.
        """
        if old_table_name == table_name:
            # No Operation
            return
        qn = connection.ops.quote_name
        params = (qn(old_table_name), qn(table_name))
        self.execute('ALTER TABLE %s RENAME TO %s;' % params)

    def delete_table(self, table_name):
        """
        Deletes the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), )
        self.execute('DROP TABLE %s;' % params)

    drop_table = delete_table

    def add_column(self, table_name, name, field):
        """
        Adds the column 'name' to the table 'table_name'.
        Uses the 'field' paramater, a django.db.models.fields.Field instance,
        to generate the necessary sql

        @param table_name: The name of the table to add the column to
        @param name: The name of the column to add
        @param field: The field to use
        """
        qn = connection.ops.quote_name
        sql = self.column_sql(table_name, name, field)
        params = (
            qn(table_name),
            sql,
        )
        sql = 'ALTER TABLE %s ADD COLUMN %s;' % params
        self.execute(sql)

    def column_sql(self, table_name, field_name, field, tablespace=''):
        """
        Creates the SQL snippet for a column. Used by add_column and add_table.
        """
        qn = connection.ops.quote_name
        field.set_attributes_from_name(field_name)

        # hook for the field to do any resolution prior to it's attributes being queried
        if hasattr(field, 'south_init'):
            field.south_init()

        sql = field.db_type()
        if not sql:
            return None

        field_output = [field.column, sql]
        field_output.append('%sNULL' % (not field.null and 'NOT ' or ''))
        if field.primary_key:
            field_output.append('PRIMARY KEY')
        elif field.unique:
            field_output.append('UNIQUE')

        tablespace = field.db_tablespace or tablespace
        if tablespace and connection.features.supports_tablespaces and field.unique:
            # We must specify the index tablespace inline, because we
            # won't be generating a CREATE INDEX statement for this field.
            field_output.append(connection.ops.tablespace_sql(tablespace, inline=True))

        sql = ' '.join(field_output)
        sqlparams = ()
        # if the field is "NOT NULL" and a default value is provided, create the column with it
        # this allows the addition of a NOT NULL field to a table with existing rows
        if not field.null and field.has_default():
            sql += " DEFAULT %s"
            sqlparams = (field.get_default())

        if field.rel:
            self.add_deferred_sql(
                self.foreign_key_sql(
                    table_name,
                    field.column,
                    field.rel.to._meta.db_table,
                    field.rel.to._meta.get_field(field.rel.field_name).column
                )
            )

        return sql % sqlparams

    def foreign_key_sql(self, from_table_name, from_column_name, to_table_name, to_column_name):
        """
        Generates a full SQL statement to add a foreign key constraint
        """
        constraint_name = '%s_refs_%s_%x' % (from_column_name, to_column_name, abs(hash((from_table_name, to_table_name))))
        return 'ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % (
            from_table_name,
            truncate_name(constraint_name, connection.ops.max_name_length()),
            from_column_name,
            to_table_name,
            to_column_name,
            connection.ops.deferrable_sql() # Django knows this
        )

    def delete_column(self, table_name, name):
        """
        Deletes the column 'column_name' from the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), qn(name))
        self.execute('ALTER TABLE %s DROP COLUMN %s CASCADE;' % params, [])

    def rename_column(self, table_name, old, new):
        """
        Renames the column 'old' from the table 'table_name' to 'new'.
        """
        raise NotImplementedError("rename_column has no generic SQL syntax")

    def start_transaction(self):
        """
        Makes sure the following commands are inside a transaction.
        Must be followed by a (commit|rollback)_transaction call.
        """
        transaction.commit_unless_managed()
        transaction.enter_transaction_management()
        transaction.managed(True)

    def commit_transaction(self):
        """
        Commits the current transaction.
        Must be preceded by a start_transaction call.
        """
        transaction.commit()
        transaction.leave_transaction_management()

    def rollback_transaction(self):
        """
        Rolls back the current transaction.
        Must be preceded by a start_transaction call.
        """
        transaction.rollback()
        transaction.leave_transaction_management()

    def send_create_signal(self, app_label, model_names):
        """
        Sends a post_syncdb signal for the model specified.

        If the model is not found (perhaps it's been deleted?),
        no signal is sent.

        TODO: The behavior of django.contrib.* apps seems flawed in that
        they don't respect created_models.  Rather, they blindly execute
        over all models within the app sending the signal.  This is a
        patch we should push Django to make  For now, this should work.
        """
        app = models.get_app(app_label)
        if not app:
            return

        created_models = []
        for model_name in model_names:
            model = models.get_model(app_label, model_name)
            if model:
                created_models.append(model)

        if created_models:
            # syncdb defaults -- perhaps take these as options?
            verbosity = 1
            interactive = True
            dispatcher.send(signal=models.signals.post_syncdb, sender=app,
                app=app, created_models=created_models,
                verbosity=verbosity, interactive=interactive)

    def mock_model(self, model_name, db_table, db_tablespace='',
                    pk_field_name='id', pk_field_type=models.AutoField,
                    pk_field_kwargs={}):
        """
        Generates a MockModel class that provides enough information
        to be used by a foreign key/many-to-many relationship.

        Migrations should prefer to use these rather than actual models
        as models could get deleted over time, but these can remain in
        migration files forever.
        """
        class MockOptions(object):
            def __init__(self):
                self.db_table = db_table
                self.db_tablespace = db_tablespace or settings.DEFAULT_TABLESPACE
                self.object_name = model_name
                self.module_name = model_name.lower()

                if pk_field_type == models.AutoField:
                    pk_field_kwargs['primary_key'] = True

                self.pk = pk_field_type(**pk_field_kwargs)
                self.pk.set_attributes_from_name(pk_field_name)

            def get_field_by_name(self, field_name):
                # we only care about the pk field
                return (self.pk, self.model, True, False)

            def get_field(self, name):
                # we only care about the pk field
                return self.pk

        class MockModel(object):
            _meta = None

        # We need to return an actual class object here, not an instance
        MockModel._meta = MockOptions()
        MockModel._meta.model = MockModel
        return MockModel
<commit_msg>added support for looking at the db_index property of a field and generating a CREATE INDEX statement automatically for it<commit_after>
from django.db import connection, transaction, models
from django.db.backends.util import truncate_name
from django.dispatch import dispatcher
from django.conf import settings

class DatabaseOperations(object):

    """
    Generic SQL implementation of the DatabaseOperations.
    Some of this code comes from Django Evolution.
    """

    def __init__(self):
        self.debug = False
        self.deferred_sql = []

    def execute(self, sql, params=[]):
        """
        Executes the given SQL statement, with optional parameters.
        If the instance's debug attribute is True, prints out what it executes.
        """
        cursor = connection.cursor()
        if self.debug:
            print "   = %s" % sql, params
        cursor.execute(sql, params)
        try:
            return cursor.fetchall()
        except:
            return []

    def add_deferred_sql(self, sql):
        """
        Add a SQL statement to the deferred list, that won't be executed until
        this instance's execute_deferred_sql method is run.
        """
        self.deferred_sql.append(sql)

    def execute_deferred_sql(self):
        """
        Executes all deferred SQL, resetting the deferred_sql list
        """
        for sql in self.deferred_sql:
            self.execute(sql)

        self.deferred_sql = []

    def create_table(self, table_name, fields):
        """
        Creates the table 'table_name'. 'fields' is a tuple of fields,
        each repsented by a 2-part tuple of field name and a
        django.db.models.fields.Field object
        """
        qn = connection.ops.quote_name
        columns = [
            self.column_sql(table_name, field_name, field)
            for field_name, field in fields
        ]

        self.execute('CREATE TABLE %s (%s);' % (table_name, ', '.join([col for col in columns])))

    add_table = create_table # Alias for consistency's sake

    def rename_table(self, old_table_name, table_name):
        """
        Renames the table 'old_table_name' to 'table_name'.
        """
        if old_table_name == table_name:
            # No Operation
            return
        qn = connection.ops.quote_name
        params = (qn(old_table_name), qn(table_name))
        self.execute('ALTER TABLE %s RENAME TO %s;' % params)

    def delete_table(self, table_name):
        """
        Deletes the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), )
        self.execute('DROP TABLE %s;' % params)

    drop_table = delete_table

    def add_column(self, table_name, name, field):
        """
        Adds the column 'name' to the table 'table_name'.
        Uses the 'field' paramater, a django.db.models.fields.Field instance,
        to generate the necessary sql

        @param table_name: The name of the table to add the column to
        @param name: The name of the column to add
        @param field: The field to use
        """
        qn = connection.ops.quote_name
        sql = self.column_sql(table_name, name, field)
        params = (
            qn(table_name),
            sql,
        )
        sql = 'ALTER TABLE %s ADD COLUMN %s;' % params
        self.execute(sql)

    def column_sql(self, table_name, field_name, field, tablespace=''):
        """
        Creates the SQL snippet for a column. Used by add_column and add_table.
        """
        qn = connection.ops.quote_name
        field.set_attributes_from_name(field_name)

        # hook for the field to do any resolution prior to it's attributes being queried
        if hasattr(field, 'south_init'):
            field.south_init()

        sql = field.db_type()
        if not sql:
            return None

        field_output = [field.column, sql]
        field_output.append('%sNULL' % (not field.null and 'NOT ' or ''))
        if field.primary_key:
            field_output.append('PRIMARY KEY')
        elif field.unique:
            field_output.append('UNIQUE')

        tablespace = field.db_tablespace or tablespace
        if tablespace and connection.features.supports_tablespaces and field.unique:
            # We must specify the index tablespace inline, because we
            # won't be generating a CREATE INDEX statement for this field.
            field_output.append(connection.ops.tablespace_sql(tablespace, inline=True))

        sql = ' '.join(field_output)
        sqlparams = ()
        # if the field is "NOT NULL" and a default value is provided, create the column with it
        # this allows the addition of a NOT NULL field to a table with existing rows
        if not field.null and field.has_default():
            sql += " DEFAULT %s"
            sqlparams = (field.get_default())

        if field.rel:
            self.add_deferred_sql(
                self.foreign_key_sql(
                    table_name,
                    field.column,
                    field.rel.to._meta.db_table,
                    field.rel.to._meta.get_field(field.rel.field_name).column
                )
            )

        if field.db_index:
            self.add_deferred_sql(self.create_index_sql(table_name, [field.column]))

        return sql % sqlparams

    def foreign_key_sql(self, from_table_name, from_column_name, to_table_name, to_column_name):
        """
        Generates a full SQL statement to add a foreign key constraint
        """
        constraint_name = '%s_refs_%s_%x' % (from_column_name, to_column_name, abs(hash((from_table_name, to_table_name))))
        return 'ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % (
            from_table_name,
            truncate_name(constraint_name, connection.ops.max_name_length()),
            from_column_name,
            to_table_name,
            to_column_name,
            connection.ops.deferrable_sql() # Django knows this
        )

    def create_index_sql(self, table_name, field_names, db_tablespace=''):
        """
        Generates a create index statement on 'table_name' for a list of 'fields'
        """
        if not field_names:
            print "No fields found to create an index for"
            return ''

        if db_tablespace and connection.features.supports_tablespaces:
            tablespace_sql = ' ' + connection.ops.tablespace_sql(db_tablespace)
        else:
            tablespace_sql = ''

        index_unique_name = ''
        if len(field_names) > 1:
            index_unique_name = '_%x' % abs(hash((table_name, ','.join(field_names))))

        index_name = '%s_%s%s' % (table_name, field_names[0], index_unique_name)
        qn = connection.ops.quote_name
        return 'CREATE INDEX %s ON %s (%s)%s;' % (
            index_name,
            table_name,
            ','.join([qn(field) for field in field_names]),
            tablespace_sql
            )

    def delete_column(self, table_name, name):
        """
        Deletes the column 'column_name' from the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), qn(name))
        self.execute('ALTER TABLE %s DROP COLUMN %s CASCADE;' % params, [])

    def rename_column(self, table_name, old, new):
        """
        Renames the column 'old' from the table 'table_name' to 'new'.
        """
        raise NotImplementedError("rename_column has no generic SQL syntax")

    def start_transaction(self):
        """
        Makes sure the following commands are inside a transaction.
        Must be followed by a (commit|rollback)_transaction call.
        """
        transaction.commit_unless_managed()
        transaction.enter_transaction_management()
        transaction.managed(True)

    def commit_transaction(self):
        """
        Commits the current transaction.
        Must be preceded by a start_transaction call.
        """
        transaction.commit()
        transaction.leave_transaction_management()

    def rollback_transaction(self):
        """
        Rolls back the current transaction.
        Must be preceded by a start_transaction call.
        """
        transaction.rollback()
        transaction.leave_transaction_management()

    def send_create_signal(self, app_label, model_names):
        """
        Sends a post_syncdb signal for the model specified.

        If the model is not found (perhaps it's been deleted?),
        no signal is sent.

        TODO: The behavior of django.contrib.* apps seems flawed in that
        they don't respect created_models.  Rather, they blindly execute
        over all models within the app sending the signal.  This is a
        patch we should push Django to make  For now, this should work.
        """
        app = models.get_app(app_label)
        if not app:
            return

        created_models = []
        for model_name in model_names:
            model = models.get_model(app_label, model_name)
            if model:
                created_models.append(model)

        if created_models:
            # syncdb defaults -- perhaps take these as options?
            verbosity = 1
            interactive = True
            dispatcher.send(signal=models.signals.post_syncdb, sender=app,
                app=app, created_models=created_models,
                verbosity=verbosity, interactive=interactive)

    def mock_model(self, model_name, db_table, db_tablespace='',
                    pk_field_name='id', pk_field_type=models.AutoField,
                    pk_field_kwargs={}):
        """
        Generates a MockModel class that provides enough information
        to be used by a foreign key/many-to-many relationship.

        Migrations should prefer to use these rather than actual models
        as models could get deleted over time, but these can remain in
        migration files forever.
        """
        class MockOptions(object):
            def __init__(self):
                self.db_table = db_table
                self.db_tablespace = db_tablespace or settings.DEFAULT_TABLESPACE
                self.object_name = model_name
                self.module_name = model_name.lower()

                if pk_field_type == models.AutoField:
                    pk_field_kwargs['primary_key'] = True

                self.pk = pk_field_type(**pk_field_kwargs)
                self.pk.set_attributes_from_name(pk_field_name)

            def get_field_by_name(self, field_name):
                # we only care about the pk field
                return (self.pk, self.model, True, False)

            def get_field(self, name):
                # we only care about the pk field
                return self.pk

        class MockModel(object):
            _meta = None

        # We need to return an actual class object here, not an instance
        MockModel._meta = MockOptions()
        MockModel._meta.model = MockModel
        return MockModel
<|endoftext|>


<commit_before># ~/ceam/ceam/modules/disease_models.py

from datetime import timedelta

from ceam import config
from ceam.state_machine import Transition, State, TransitionSet
from ceam.modules.disease import DiseaseModule, DiseaseState, ExcessMortalityState, IncidenceRateTransition, ProportionTransition

def heart_disease_factory():
    module = DiseaseModule('ihd')

    healthy = State('healthy')

    # Calculate an adjusted disability weight for the acute heart attack phase that
    # accounts for the fact that our timestep is longer than the phase length
    # TODO: This doesn't account for the fact that our timestep is longer than 28 days
    timestep = config.getfloat('simulation_parameters', 'time_step')
    weight = 0.43*(2/timestep) + 0.07*(28/timestep)
    heart_attack = ExcessMortalityState('heart_attack', disability_weight=weight, dwell_time=timedelta(days=28), modelable_entity_id=1814, prevalence_meid=1814)

    #
    mild_heart_failure = ExcessMortalityState('mild_heart_failure', disability_weight=0.04, modelable_entity_id=2412, prevalence_meid=1821)
    moderate_heart_failure = ExcessMortalityState('moderate_heart_failure', disability_weight=0.07, modelable_entity_id=2412, prevalence_meid=1822)
    severe_heart_failure = ExcessMortalityState('severe_heart_failure', disability_weight=0.18, modelable_entity_id=2412, prevalence_meid=1823)

    asymptomatic_angina = ExcessMortalityState('asymptomatic_angina', disability_weight=0.0, modelable_entity_id=1817, prevalence_meid=3102)
    mild_angina = ExcessMortalityState('mild_angina', disability_weight=0.03, modelable_entity_id=1817, prevalence_meid=1818)
    moderate_angina = ExcessMortalityState('moderate_angina', disability_weight=0.08, modelable_entity_id=1817, prevalence_meid=1819)
    severe_angina = ExcessMortalityState('severe_angina', disability_weight=0.17, modelable_entity_id=1817, prevalence_meid=1820)

    asymptomatic_ihd = ExcessMortalityState('asymptomatic_ihd', disability_weight=0.08, modelable_entity_id=3233, prevalence_meid=3233)

    heart_attack_transition = IncidenceRateTransition(heart_attack, 'heart_attack', modelable_entity_id=1814)
    healthy.transition_set.append(heart_attack_transition)

    heart_failure_buckets = TransitionSet(allow_null_transition=False)
    heart_failure_buckets.extend([
        ProportionTransition(mild_heart_failure, proportion=0.182074),
        ProportionTransition(moderate_heart_failure, proportion=0.149771),
        ProportionTransition(severe_heart_failure, proportion=0.402838),
        ])

    angina_buckets = TransitionSet(allow_null_transition=False)
    angina_buckets.extend([
        ProportionTransition(asymptomatic_angina, proportion=0.304553),
        ProportionTransition(mild_angina, proportion=0.239594),
        ProportionTransition(moderate_angina, proportion=0.126273),
        ProportionTransition(severe_angina, proportion=0.32958),
        ])
    healthy.transition_set.append(IncidenceRateTransition(angina_buckets, 'non_mi_angina', modelable_entity_id=1817))

    heart_attack.transition_set.allow_null_transition=False
    heart_attack.transition_set.append(ProportionTransition(heart_failure_buckets, proportion=0.01))
    heart_attack.transition_set.append(ProportionTransition(angina_buckets, proportion=0.15))
    heart_attack.transition_set.append(ProportionTransition(asymptomatic_ihd, proportion=0.84))

    mild_heart_failure.transition_set.append(heart_attack_transition)
    moderate_heart_failure.transition_set.append(heart_attack_transition)
    severe_heart_failure.transition_set.append(heart_attack_transition)
    asymptomatic_angina.transition_set.append(heart_attack_transition)
    mild_angina.transition_set.append(heart_attack_transition)
    moderate_angina.transition_set.append(heart_attack_transition)
    severe_angina.transition_set.append(heart_attack_transition)
    asymptomatic_ihd.transition_set.append(heart_attack_transition)

    module.states.extend([healthy, heart_attack, mild_heart_failure, moderate_heart_failure, severe_heart_failure, asymptomatic_angina, mild_angina, moderate_angina, severe_angina, asymptomatic_ihd])
    return module

def stroke_factory():
    module = DiseaseModule('hemorrhagic_stroke')

    healthy = State('healthy')
    # TODO: disability weight for stroke
    hemorrhagic_stroke = ExcessMortalityState('hemorrhagic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9311)
    ischemic_stroke = ExcessMortalityState('ischemic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9310)
    chronic_stroke = ExcessMortalityState('chronic_stroke', disability_weight=0.32, modelable_entity_id=9312)

    hemorrhagic_transition = IncidenceRateTransition(hemorrhagic_stroke, 'hemorrhagic_stroke', modelable_entity_id=9311)
    ischemic_transition = IncidenceRateTransition(ischemic_stroke, 'ischemic_stroke', modelable_entity_id=9310)
    healthy.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    hemorrhagic_stroke.transition_set.append(Transition(chronic_stroke))
    ischemic_stroke.transition_set.append(Transition(chronic_stroke))

    chronic_stroke.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    module.states.extend([healthy, hemorrhagic_stroke, ischemic_stroke, chronic_stroke])

    return module

# End.
<commit_msg>Cosmetic: line spacing.<commit_after># ~/ceam/ceam/modules/disease_models.py

from datetime import timedelta

from ceam import config
from ceam.state_machine import Transition, State, TransitionSet
from ceam.modules.disease import DiseaseModule, DiseaseState, ExcessMortalityState, IncidenceRateTransition, ProportionTransition

def heart_disease_factory():
    module = DiseaseModule('ihd')

    healthy = State('healthy')

    # Calculate an adjusted disability weight for the acute heart attack phase that
    # accounts for the fact that our timestep is longer than the phase length
    # TODO: This doesn't account for the fact that our timestep is longer than 28 days
    timestep = config.getfloat('simulation_parameters', 'time_step')
    weight = 0.43*(2/timestep) + 0.07*(28/timestep)
    heart_attack = ExcessMortalityState('heart_attack', disability_weight=weight, dwell_time=timedelta(days=28), modelable_entity_id=1814, prevalence_meid=1814)

    #
    mild_heart_failure = ExcessMortalityState('mild_heart_failure', disability_weight=0.04, modelable_entity_id=2412, prevalence_meid=1821)
    moderate_heart_failure = ExcessMortalityState('moderate_heart_failure', disability_weight=0.07, modelable_entity_id=2412, prevalence_meid=1822)
    severe_heart_failure = ExcessMortalityState('severe_heart_failure', disability_weight=0.18, modelable_entity_id=2412, prevalence_meid=1823)

    asymptomatic_angina = ExcessMortalityState('asymptomatic_angina', disability_weight=0.0, modelable_entity_id=1817, prevalence_meid=3102)
    mild_angina = ExcessMortalityState('mild_angina', disability_weight=0.03, modelable_entity_id=1817, prevalence_meid=1818)
    moderate_angina = ExcessMortalityState('moderate_angina', disability_weight=0.08, modelable_entity_id=1817, prevalence_meid=1819)
    severe_angina = ExcessMortalityState('severe_angina', disability_weight=0.17, modelable_entity_id=1817, prevalence_meid=1820)

    asymptomatic_ihd = ExcessMortalityState('asymptomatic_ihd', disability_weight=0.08, modelable_entity_id=3233, prevalence_meid=3233)

    heart_attack_transition = IncidenceRateTransition(heart_attack, 'heart_attack', modelable_entity_id=1814)
    healthy.transition_set.append(heart_attack_transition)

    heart_failure_buckets = TransitionSet(allow_null_transition=False)
    heart_failure_buckets.extend([
        ProportionTransition(mild_heart_failure, proportion=0.182074),
        ProportionTransition(moderate_heart_failure, proportion=0.149771),
        ProportionTransition(severe_heart_failure, proportion=0.402838),
        ])

    angina_buckets = TransitionSet(allow_null_transition=False)
    angina_buckets.extend([
        ProportionTransition(asymptomatic_angina, proportion=0.304553),
        ProportionTransition(mild_angina, proportion=0.239594),
        ProportionTransition(moderate_angina, proportion=0.126273),
        ProportionTransition(severe_angina, proportion=0.32958),
        ])
    healthy.transition_set.append(IncidenceRateTransition(angina_buckets, 'non_mi_angina', modelable_entity_id=1817))

    heart_attack.transition_set.allow_null_transition=False
    heart_attack.transition_set.append(ProportionTransition(heart_failure_buckets, proportion=0.01))
    heart_attack.transition_set.append(ProportionTransition(angina_buckets, proportion=0.15))
    heart_attack.transition_set.append(ProportionTransition(asymptomatic_ihd, proportion=0.84))

    mild_heart_failure.transition_set.append(heart_attack_transition)
    moderate_heart_failure.transition_set.append(heart_attack_transition)
    severe_heart_failure.transition_set.append(heart_attack_transition)
    asymptomatic_angina.transition_set.append(heart_attack_transition)
    mild_angina.transition_set.append(heart_attack_transition)
    moderate_angina.transition_set.append(heart_attack_transition)
    severe_angina.transition_set.append(heart_attack_transition)
    asymptomatic_ihd.transition_set.append(heart_attack_transition)

    module.states.extend([healthy, heart_attack, mild_heart_failure, moderate_heart_failure, severe_heart_failure, asymptomatic_angina, mild_angina, moderate_angina, severe_angina, asymptomatic_ihd])
    return module

def stroke_factory():
    module = DiseaseModule('hemorrhagic_stroke')

    healthy = State('healthy')
    # TODO: disability weight for stroke
    hemorrhagic_stroke = ExcessMortalityState('hemorrhagic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9311)
    ischemic_stroke = ExcessMortalityState('ischemic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9310)
    chronic_stroke = ExcessMortalityState('chronic_stroke', disability_weight=0.32, modelable_entity_id=9312)

    hemorrhagic_transition = IncidenceRateTransition(hemorrhagic_stroke, 'hemorrhagic_stroke', modelable_entity_id=9311)
    ischemic_transition = IncidenceRateTransition(ischemic_stroke, 'ischemic_stroke', modelable_entity_id=9310)
    healthy.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    hemorrhagic_stroke.transition_set.append(Transition(chronic_stroke))
    ischemic_stroke.transition_set.append(Transition(chronic_stroke))

    chronic_stroke.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    module.states.extend([healthy, hemorrhagic_stroke, ischemic_stroke, chronic_stroke])

    return module

# End.
<|endoftext|>


<commit_before>import pytest
from pupa.scrape import Jurisdiction as JurisdictionBase
from pupa.importers import JurisdictionImporter
from opencivicdata.models import Jurisdiction

class FakeJurisdiction(JurisdictionBase):
    division_id = 'division-id'
    name = 'test'
    url = 'http://example.com'
    classification = 'government'

@pytest.mark.django_db
def test_jurisdiction_import():
    tj = FakeJurisdiction()
    juris_dict = tj.as_dict()
    JurisdictionImporter('jurisdiction-id').import_data([juris_dict])

    dbj = Jurisdiction.objects.get()
    assert dbj.id == tj.jurisdiction_id
    assert dbj.division_id == tj.division_id
    assert dbj.name == tj.name
    assert dbj.url == tj.url

@pytest.mark.django_db
def test_jurisdiction_no_duplicates():
    tj = FakeJurisdiction()
    ji = JurisdictionImporter('jurisdiction-id')
    ji.import_data([tj.as_dict()])
    ji.import_data([tj.as_dict()])
    assert Jurisdiction.objects.count() == 1
<commit_msg>test update and noop on jurisdiction<commit_after>import pytest
from pupa.scrape import Jurisdiction as JurisdictionBase
from pupa.importers import JurisdictionImporter
from opencivicdata.models import Jurisdiction

class FakeJurisdiction(JurisdictionBase):
    division_id = 'division-id'
    name = 'test'
    url = 'http://example.com'
    classification = 'government'

@pytest.mark.django_db
def test_jurisdiction_import():
    tj = FakeJurisdiction()
    juris_dict = tj.as_dict()
    JurisdictionImporter('jurisdiction-id').import_data([juris_dict])

    dbj = Jurisdiction.objects.get()
    assert dbj.id == tj.jurisdiction_id
    assert dbj.division_id == tj.division_id
    assert dbj.name == tj.name
    assert dbj.url == tj.url

@pytest.mark.django_db
def test_jurisdiction_update():
    tj = FakeJurisdiction()
    ji = JurisdictionImporter('jurisdiction-id')
    _, what = ji.import_item(tj.as_dict())
    assert what == 'insert'

    _, what = ji.import_item(tj.as_dict())
    assert what == 'noop'
    assert Jurisdiction.objects.count() == 1

    tj.name = 'different name'
    obj, what = ji.import_item(tj.as_dict())
    assert what == 'update'
    assert Jurisdiction.objects.count() == 1
    assert obj.name == 'different name'
<|endoftext|>


<commit_before># Ben Jones bjones99@gatech.edu
# Georgia Tech Fall 2014
#
# tcpdump.py: interface to tcpdump to stop and start captures and do
# second passes over existing pcaps

from base64 import b64encode
import logging
import os
import tempfile

# local imports
import centinel
from centinel import command

class Tcpdump():
    """Class to interface between tcpdump and Python"""

    def __init__(self, filename=None, pcap_args=None):
        if filename is None:
            temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
            temp_file.close()
            filename = temp_file.name
        self.filename = filename
        # don't change this to default value because it is a mutable
        # type and whatever you change it to will become the new
        # default value until the interpreter is restarted
        if pcap_args is None:
            # use the centinel configured tcpdump options if available
            # (if not specified by the user, this will be -i any, so
            # the same as below
            if 'tcpdump_params' in centinel.conf['experiments']:
                pcap_args = centinel.conf['experiments']['tcpdump_params']
            # for backwards compatability, ensure that we give some
            # pcap args for what to capture
            else:
                pcap_args = ["-i", "any"]
                logging.warning("Global config not available, so falling "
                                "back on -i any pcap args")
        self.pcap_args = pcap_args

    def start(self):
        cmd = ['sudo', 'tcpdump', '-w', self.filename]
        cmd.extend(self.pcap_args)
        self.caller = command.Command(cmd, _tcpdump_callback)
        self.caller.start()

    def stop(self):
        if self.caller is not None:
            self.caller.stop()

    def post_processing(self, out_filter, out_file=None):
        if out_file is None:
            temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
            temp_file.close()
            out_file = temp_file.name
        cmd = ['tcpdump', '-r', self.filename, '-w', out_file]
        caller = command.Command(cmd, _tcpdump_callback)
        caller.start()

    def b64_output(self):
        with open(self.filename, 'r') as file_p:
            return b64encode(file_p.read())

    def pcap(self):
        with open(self.filename, 'r') as file_p:
            return file_p.read()

    def delete(self):
        os.remove(self.filename)

def _tcpdump_callback(self, line, kill_switch):
    """Callback function to handle tcpdump"""

    line = line.lower()
    if ("listening" in line) or ("reading" in line):
        self.started = True
    if ("no suitable device" in line):
        self.error = True
        self.kill_switch()
    if "by kernel" in line:
        self.stopped = True
<commit_msg>added tcpdump file address return<commit_after># Ben Jones bjones99@gatech.edu
# Georgia Tech Fall 2014
#
# tcpdump.py: interface to tcpdump to stop and start captures and do
# second passes over existing pcaps

from base64 import b64encode
import logging
import os
import tempfile

# local imports
import centinel
from centinel import command

class Tcpdump():
    """Class to interface between tcpdump and Python"""

    def __init__(self, filename=None, pcap_args=None):
        if filename is None:
            temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
            temp_file.close()
            filename = temp_file.name
        self.filename = filename
        # don't change this to default value because it is a mutable
        # type and whatever you change it to will become the new
        # default value until the interpreter is restarted
        if pcap_args is None:
            # use the centinel configured tcpdump options if available
            # (if not specified by the user, this will be -i any, so
            # the same as below
            if 'tcpdump_params' in centinel.conf['experiments']:
                pcap_args = centinel.conf['experiments']['tcpdump_params']
            # for backwards compatability, ensure that we give some
            # pcap args for what to capture
            else:
                pcap_args = ["-i", "any"]
                logging.warning("Global config not available, so falling "
                                "back on -i any pcap args")
        self.pcap_args = pcap_args

    def start(self):
        cmd = ['sudo', 'tcpdump', '-w', self.filename]
        cmd.extend(self.pcap_args)
        self.caller = command.Command(cmd, _tcpdump_callback)
        self.caller.start()

    def stop(self):
        if self.caller is not None:
            self.caller.stop()

    def post_processing(self, out_filter, out_file=None):
        if out_file is None:
            temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
            temp_file.close()
            out_file = temp_file.name
        cmd = ['tcpdump', '-r', self.filename, '-w', out_file]
        caller = command.Command(cmd, _tcpdump_callback)
        caller.start()

    def b64_output(self):
        with open(self.filename, 'r') as file_p:
            return b64encode(file_p.read())

    def pcap(self):
        with open(self.filename, 'r') as file_p:
            return file_p.read()

    def pcap_filename(self):
        return self.filename

    def delete(self):
        os.remove(self.filename)

def _tcpdump_callback(self, line, kill_switch):
    """Callback function to handle tcpdump"""

    line = line.lower()
    if ("listening" in line) or ("reading" in line):
        self.started = True
    if ("no suitable device" in line):
        self.error = True
        self.kill_switch()
    if "by kernel" in line:
        self.stopped = True
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.test import TestCase
from django.utils import timezone

from dailyreport.forms import DailyForm, CommentForm, TaskForm, SearchForm, DateForm
from dailyreport.models import Daily, Task, Comment

# 10 tests
class DailyFormTest(TestCase):
    def test_valid(self):
        params = dict(title='test_title', create_date=timezone.now().date())
        daily = Daily()
        form = DailyForm(params, instance=daily)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        params = dict()
        daily = Daily()
        form = DailyForm(params, instance=daily)
        self.assertFalse(form.is_valid())

    def test_invalid_spaces(self):
        params = dict(title=' ', create_date=timezone.now().date())
        daily = Daily()
        form = DailyForm(params, instance=daily)
        self.assertFalse(form.is_valid())

class TaskFormTest(TestCase):
    def test_valid(self):
        params = dict(name='test_name', time_plan=1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertTrue(form.is_valid())

    def test_valid_include_complete_task(self):
        params = dict(name='test_name', time_plan=1, time_real=1,
                      implement_date=timezone.now().date(), complete_task=True)
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        params = dict()
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_minus_plan(self):
        params = dict(name='test_name', time_plan=-1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_minus_real(self):
        params = dict(name='test_name', time_plan=1, time_real=-1, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_zero_time(self):
        params = dict(name='test_name', time_plan=0, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_complete_task(self):
        params = dict(name='test_name', time_plan=0, time_real=0,
                      implement_date=timezone.now().date(), complete_task=True)
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_name1(self):
        params = dict(name='', time_plan=1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_name2(self):
        params = dict(name=' ', time_plan=1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

class CommentFormTest(TestCase):
    def test_valid(self):
        params = dict(comment='test_comment')
        comment = Comment()
        form = CommentForm(params, instance=comment)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        params = dict()
        comment = Comment()
        form = CommentForm(params, instance=comment)
        self.assertFalse(form.is_valid())

    def test_invalid_spaces(self):
        params = dict(comment=' ')
        comment = Comment()
        form = CommentForm(params, instance=comment)
        self.assertFalse(form.is_valid())

class SearchFormTest(TestCase):
    def test_valid(self):
        param = dict(keyword='test_keyword')
        form = SearchForm(param)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        u"""
            現行では、空文字での検索を許容しています
            追記：最新ではスペースのみなどの検索キーワードを認めません
        :return:
        """
        param = dict()
        form = SearchForm(param)
        self.assertFalse(form.is_valid())

    def test_invalid_spaces(self):
        param = dict(keyword=' ')
        form = SearchForm(param)
        self.assertFalse(form.is_valid())

class DateFormTest(TestCase):
    def test_valid(self):
        param = dict(date=timezone.now().date())
        form = DateForm(param)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        param = dict()
        form = DateForm(param)
        self.assertFalse(form.is_valid())
<commit_msg>[rename]spaces to space<commit_after># -*- coding: utf-8 -*-
from django.test import TestCase
from django.utils import timezone

from dailyreport.forms import DailyForm, CommentForm, TaskForm, SearchForm, DateForm
from dailyreport.models import Daily, Task, Comment

# 10 tests
class DailyFormTest(TestCase):
    def test_valid(self):
        params = dict(title='test_title', create_date=timezone.now().date())
        daily = Daily()
        form = DailyForm(params, instance=daily)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        params = dict()
        daily = Daily()
        form = DailyForm(params, instance=daily)
        self.assertFalse(form.is_valid())

    def test_invalid_space(self):
        params = dict(title=' ', create_date=timezone.now().date())
        daily = Daily()
        form = DailyForm(params, instance=daily)
        self.assertFalse(form.is_valid())

class TaskFormTest(TestCase):
    def test_valid(self):
        params = dict(name='test_name', time_plan=1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertTrue(form.is_valid())

    def test_valid_include_complete_task(self):
        params = dict(name='test_name', time_plan=1, time_real=1,
                      implement_date=timezone.now().date(), complete_task=True)
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        params = dict()
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_minus_plan(self):
        params = dict(name='test_name', time_plan=-1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_minus_real(self):
        params = dict(name='test_name', time_plan=1, time_real=-1, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_zero_time(self):
        params = dict(name='test_name', time_plan=0, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_complete_task(self):
        params = dict(name='test_name', time_plan=0, time_real=0,
                      implement_date=timezone.now().date(), complete_task=True)
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_name1(self):
        params = dict(name='', time_plan=1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

    def test_invalid_name2(self):
        params = dict(name=' ', time_plan=1, time_real=0, implement_date=timezone.now().date())
        task = Task()
        form = TaskForm(params, instance=task)
        self.assertFalse(form.is_valid())

class CommentFormTest(TestCase):
    def test_valid(self):
        params = dict(comment='test_comment')
        comment = Comment()
        form = CommentForm(params, instance=comment)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        params = dict()
        comment = Comment()
        form = CommentForm(params, instance=comment)
        self.assertFalse(form.is_valid())

    def test_invalid_space(self):
        params = dict(comment=' ')
        comment = Comment()
        form = CommentForm(params, instance=comment)
        self.assertFalse(form.is_valid())

class SearchFormTest(TestCase):
    def test_valid(self):
        param = dict(keyword='test_keyword')
        form = SearchForm(param)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        u"""
            現行では、空文字での検索を許容しています
            追記：最新ではスペースのみなどの検索キーワードを認めません
        :return:
        """
        param = dict()
        form = SearchForm(param)
        self.assertFalse(form.is_valid())

    def test_invalid_space(self):
        param = dict(keyword=' ')
        form = SearchForm(param)
        self.assertFalse(form.is_valid())

class DateFormTest(TestCase):
    def test_valid(self):
        param = dict(date=timezone.now().date())
        form = DateForm(param)
        self.assertTrue(form.is_valid())

    def test_invalid(self):
        param = dict()
        form = DateForm(param)
        self.assertFalse(form.is_valid())
<|endoftext|>


<commit_before>"""
Mercurial version control
"""
from rez.release_vcs import ReleaseVCS
from rez.exceptions import ReleaseVCSError
import os.path

class HgReleaseVCSError(ReleaseVCSError):
    pass

class HgReleaseVCS(ReleaseVCS):
    @classmethod
    def name(cls):
        return 'hg'

    def __init__(self, path):
        super(HgReleaseVCS, self).__init__(path)
        self.executable = self.find_executable('hg')

        hgdir = os.path.join(self.path, '.hg')
        if not os.path.isdir(hgdir):
            raise HgReleaseVCSError(
                "'%s' is not a mercurial working copy" % self.path)
        try:
            assert self.hg('root')[0] == self.path
        except AssertionError:
            raise HgReleaseVCSError(
                "'%s' is not the root of a mercurial working copy" % self.path)
        except Exception, err:
            raise HgReleaseVCSError("failed to call hg binary: " + str(err))

        self.patch_path = os.path.join(hgdir, 'patches')
        if not os.path.isdir(self.patch_path):
            self.patch_path = None

    @classmethod
    def is_valid_root(cls, path):
        return os.path.isdir(os.path.join(path, '.hg'))

    def hg(self, *nargs):
        return self._cmd(self.executable, *nargs)

    def _create_tag_impl(self, tag_name, message=None):
        if self.patch_path:
            # patch queue
            self.hg('tag', '-f', tag_name, '--message', message, '--mq')
            # use a bookmark on the main repo since we can't change it
            self.hg('bookmark', '-f', tag_name)
        else:
            self.hg('tag', '-f', tag_name)

    def validate_repostate(self):
        def _check(modified, path):
            if modified:
                modified = [line.split()[-1] for line in modified]
                raise ReleaseVCSError(
                    "%s is not in a state to release - please commit outstanding "
                    "changes: %s" % (path, ', '.join(modified)))

        _check(self.hg('status', '-m', '-a'), self.path)
        if self.patch_path:
            _check(self.hg('status', '-m', '-a', '--mq'), self.patch_path)

    def get_changelog(self, previous_revision=None):
        return "TODO"

def register_plugin():
    return HgReleaseVCS
<commit_msg>implemented some missing methods for HgReleaseVCS<commit_after>"""
Mercurial version control
"""
from rez.release_vcs import ReleaseVCS
from rez.exceptions import ReleaseVCSError
from rez.utils.logging_ import print_debug, print_error
import os.path

class HgReleaseVCSError(ReleaseVCSError):
    pass

class HgReleaseVCS(ReleaseVCS):
    @classmethod
    def name(cls):
        return 'hg'

    def __init__(self, path):
        super(HgReleaseVCS, self).__init__(path)
        self.executable = self.find_executable('hg')

        hgdir = os.path.join(self.path, '.hg')
        if not os.path.isdir(hgdir):
            raise HgReleaseVCSError(
                "'%s' is not a mercurial working copy" % self.path)
        try:
            assert self.hg('root')[0] == self.path
        except AssertionError:
            raise HgReleaseVCSError(
                "'%s' is not the root of a mercurial working copy" % self.path)
        except Exception, err:
            raise HgReleaseVCSError("failed to call hg binary: " + str(err))

        self.patch_path = os.path.join(hgdir, 'patches')
        if not os.path.isdir(self.patch_path):
            self.patch_path = None

    @classmethod
    def is_valid_root(cls, path):
        return os.path.isdir(os.path.join(path, '.hg'))

    def hg(self, *nargs, **kwargs):
        if kwargs.pop('patch', False):
            nargs += ('--mq',)
        if kwargs:
            raise HgReleaseVCSError("Unrecognized keyword args to hg command:"
                                    " %s" % ", ".join(kwargs))
        return self._cmd(self.executable, *nargs)

    def _create_tag_impl(self, tag_name, message=None):
        tag_args = ['tag', '-f', tag_name]
        if message:
            tag_args += ['--message', message]
        if self.patch_path:
            # patch queue
            self.hg(patch=True, *tag_args)
            # use a bookmark on the main repo since we can't change it
            self.hg('bookmark', '-f', tag_name)
        else:
            self.hg(*tag_args)

    def get_paths(self, patch=False):
        paths = self.hg("paths", patch=patch)
        return dict(line.split(' = ', 1) for line in paths if line)

    def get_default_url(self, patch=False):
        return self.get_paths(patch=patch).get('default')

    def validate_repostate(self):
        def _check(modified, path):
            if modified:
                modified = [line.split()[-1] for line in modified]
                raise ReleaseVCSError(
                    "%s is not in a state to release - please commit outstanding "
                    "changes: %s" % (path, ', '.join(modified)))

        _check(self.hg('status', '-m', '-a'), self.path)
        if self.patch_path:
            _check(self.hg('status', '-m', '-a', '--mq'), self.patch_path)

    def get_current_revision(self):
        doc = {
            'commit': self.hg("log", "--template", "{node}", "-r", ".")[0],
            'branch': self.hg("branch")[0],
        }

        def _get(key, fn):
            try:
                value = fn()
                doc[key] = value
                return (value is not None)
            except Exception as e:
                print_error("Error retrieving %s: %s" % (key, str(e)))
                return False

        _get("push_url", self.get_default_url)

        return doc

    def get_changelog(self, previous_revision=None):
        prev_commit = None
        if previous_revision is not None:
            try:
                prev_commit = previous_revision["commit"]
            except:
                if self.package.config.debug("package_release"):
                    print_debug("couldn't determine previous commit from: %r"
                                % previous_revision)

        if prev_commit:
            # git behavior is to simply print the log from the last common
            # ancsestor... which is apparently desired. so we'll mimic that
            commit_range = "ancestor(%s, .)::." % prev_commit
            stdout = self.hg("log", "-r", commit_range)
        else:
            stdout = self.hg("log")
        return '\n'.join(stdout)

    def create_release_tag(self, tag_name, message=None):
        # create tag
        print "Creating tag '%s'..." % tag_name
        self._create_tag_impl(tag_name, message=message)

        # push tag
        main_url = self.get_default_url()
        if self.patch_path:
            # push the tag on the patch repo...
            patch_url = self.get_default_url(patch=True)
            if patch_url:
                self.hg('push', patch_url, patch=True)

            # ...and push the bookmark on the main
            if main_url:
                self.hg('push', '--bookmark', tag_name, main_url)
        elif main_url:
            self.hg('push', main_url)

def register_plugin():
    return HgReleaseVCS
<|endoftext|>


<commit_before># Copyright 2015 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

import os
import sys

_CATAPULT_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..'))

# Directories in catapult/third_party required by dashboard.
THIRD_PARTY_LIBRARIES = [
    'apiclient',
    'beautifulsoup4',
    'cachetools',
    'certifi',
    'chardet',
    'click',
    'cloudstorage',
    'depot_tools',
    'flask',
    'flot',
    'gae_ts_mon',
    'google-auth',
    'graphy',
    'html5lib-python',
    'idna',
    'ijson',
    'itsdangerous',
    'jinja2',
    'jquery',
    'mapreduce',
    'markupsafe',
    'mock',
    'oauth2client',
    'pipeline',
    'polymer',
    'polymer-svg-template',
    'polymer2/bower_components',
    'polymer2/bower_components/chopsui',
    'pyasn1',
    'pyasn1_modules',
    'pyparsing',
    'redux/redux.min.js',
    'requests',
    'requests_toolbelt',
    'rsa',
    'six',
    'uritemplate',
    'urllib3',
    'webapp2',
    'webtest',
    'werkzeug',
]

THIRD_PARTY_LIBRARIES_PY2 = THIRD_PARTY_LIBRARIES + [
    'httplib2/python2/httplib2'
]

THIRD_PARTY_LIBRARIES_PY3 = THIRD_PARTY_LIBRARIES + [
    'httplib2/python3/httplib2'
]

# Files and directories in catapult/dashboard.
DASHBOARD_FILES = [
    'api.yaml',
    'app.yaml',
    'appengine_config.py',
    'cron.yaml',
    'dashboard',
    'dispatch.yaml',
    'index.yaml',
    'pinpoint.yaml',
    'queue.yaml',
    'scripts.yaml',
    'upload-processing.yaml',
    'upload.yaml',
]

TRACING_PATHS = [
    'tracing/tracing',
    'tracing/tracing_build',
    'tracing/third_party/gl-matrix/dist/gl-matrix-min.js',
    'tracing/third_party/mannwhitneyu',
]

def PathsForDeployment():
  """Returns a list of paths to things required for deployment.

  This includes both Python libraries that are required, and also
  other files, such as config files.

  This list is used when building a temporary deployment directory;
  each of the items in this list will have a corresponding file or
  directory with the same basename in the deployment directory.
  """
  paths = []
  paths.extend(_CatapultThirdPartyLibraryPaths())
  for name in DASHBOARD_FILES:
    paths.append(os.path.join(_CATAPULT_PATH, 'dashboard', name))
  paths.append(os.path.join(_CATAPULT_PATH, 'tracing', 'tracing_project.py'))
  paths.append(os.path.join(_CATAPULT_PATH, 'common', 'py_utils', 'py_utils'))
  # Required by py_utils
  paths.append(os.path.join(_CATAPULT_PATH, 'devil', 'devil'))
  paths.extend(_TracingPaths())
  return paths

def PathsForTesting():
  """Returns a list of Python library paths required for dashboard tests."""
  return _AllSdkThirdPartyLibraryPaths() + _CatapultThirdPartyLibraryPaths() + [
      os.path.join(_CATAPULT_PATH, 'dashboard'),
      os.path.join(_CATAPULT_PATH, 'tracing'),
      os.path.join(_CATAPULT_PATH, 'common', 'py_utils', 'py_utils'),

      # Required by py_utils
      os.path.join(_CATAPULT_PATH, 'devil', 'devil'),

      # Isolate the sheriff_config package, since it's deployed independently.
      os.path.join(_CATAPULT_PATH, 'dashboard', 'dashboard', 'sheriff_config'),
  ]

def _AllSdkThirdPartyLibraryPaths():
  """Returns a list of all third party library paths from the SDK.

  The AppEngine documentation directs us to add App Engine libraries from the
  SDK to our Python path for local unit tests.
    https://cloud.google.com/appengine/docs/python/tools/localunittesting
  """
  paths = []
  for sdk_bin_path in os.environ['PATH'].split(os.pathsep):
    if 'google-cloud-sdk' not in sdk_bin_path:
      continue

    if not os.path.isdir(sdk_bin_path):
      sdk_bin_path = os.path.dirname(sdk_bin_path)

    appengine_path = os.path.join(sdk_bin_path, 'platform', 'google_appengine')
    paths.append(appengine_path)
    sys.path.insert(0, appengine_path)

  try:
    import dev_appserver
  except ImportError:
    # TODO: Put the Cloud SDK in the path with the binary dependency manager.
    # https://github.com/catapult-project/catapult/issues/2135
    print('This script requires the Google Cloud SDK to be in PATH.')
    print('Install at https://cloud.google.com/sdk and then run')
    print('`gcloud components install app-engine-python`')
    sys.exit(1)

  paths.extend(dev_appserver.EXTRA_PATHS)
  return paths

def _CatapultThirdPartyLibraryPaths():
  """Returns a list of required third-party libraries in catapult."""
  paths = []
  paths.append(
      os.path.join(_CATAPULT_PATH, 'common', 'node_runner', 'node_runner',
                   'node_modules', '@chopsui', 'tsmon-client',
                   'tsmon-client.js'))
  third_party_libraries = (
      THIRD_PARTY_LIBRARIES_PY3 if sys.version_info.major == 3
      else THIRD_PARTY_LIBRARIES_PY2)
  for library in third_party_libraries:
    paths.append(os.path.join(_CATAPULT_PATH, 'third_party', library))
  return paths

def _TracingPaths():
  """Returns a list of paths that may be imported from tracing."""
  # TODO(sullivan): This should either pull from tracing_project or be generated
  # via gypi. See https://github.com/catapult-project/catapult/issues/3048.
  paths = []
  for path in TRACING_PATHS:
    paths.append(os.path.join(_CATAPULT_PATH, os.path.normpath(path)))
  return paths
<commit_msg>[pinpoint] fix import issue on httplib2<commit_after># Copyright 2015 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

import os
import sys

_CATAPULT_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '..', '..'))

# Directories in catapult/third_party required by dashboard.
THIRD_PARTY_LIBRARIES = [
    'apiclient',
    'beautifulsoup4',
    'cachetools',
    'certifi',
    'chardet',
    'click',
    'cloudstorage',
    'depot_tools',
    'flask',
    'flot',
    'gae_ts_mon',
    'google-auth',
    'graphy',
    'html5lib-python',
    'idna',
    'ijson',
    'itsdangerous',
    'jinja2',
    'jquery',
    'mapreduce',
    'markupsafe',
    'mock',
    'oauth2client',
    'pipeline',
    'polymer',
    'polymer-svg-template',
    'polymer2/bower_components',
    'polymer2/bower_components/chopsui',
    'pyasn1',
    'pyasn1_modules',
    'pyparsing',
    'redux/redux.min.js',
    'requests',
    'requests_toolbelt',
    'rsa',
    'six',
    'uritemplate',
    'urllib3',
    'webapp2',
    'webtest',
    'werkzeug',
]

THIRD_PARTY_LIBRARIES_PY2 = THIRD_PARTY_LIBRARIES + ['httplib2/python2']

THIRD_PARTY_LIBRARIES_PY3 = THIRD_PARTY_LIBRARIES + ['httplib2/python3']

# Files and directories in catapult/dashboard.
DASHBOARD_FILES = [
    'api.yaml',
    'app.yaml',
    'appengine_config.py',
    'cron.yaml',
    'dashboard',
    'dispatch.yaml',
    'index.yaml',
    'pinpoint.yaml',
    'queue.yaml',
    'scripts.yaml',
    'upload-processing.yaml',
    'upload.yaml',
]

TRACING_PATHS = [
    'tracing/tracing',
    'tracing/tracing_build',
    'tracing/third_party/gl-matrix/dist/gl-matrix-min.js',
    'tracing/third_party/mannwhitneyu',
]

def PathsForDeployment():
  """Returns a list of paths to things required for deployment.

  This includes both Python libraries that are required, and also
  other files, such as config files.

  This list is used when building a temporary deployment directory;
  each of the items in this list will have a corresponding file or
  directory with the same basename in the deployment directory.
  """
  paths = []
  paths.extend(_CatapultThirdPartyLibraryPaths())
  for name in DASHBOARD_FILES:
    paths.append(os.path.join(_CATAPULT_PATH, 'dashboard', name))
  paths.append(os.path.join(_CATAPULT_PATH, 'tracing', 'tracing_project.py'))
  paths.append(os.path.join(_CATAPULT_PATH, 'common', 'py_utils', 'py_utils'))
  # Required by py_utils
  paths.append(os.path.join(_CATAPULT_PATH, 'devil', 'devil'))
  paths.extend(_TracingPaths())
  return paths

def PathsForTesting():
  """Returns a list of Python library paths required for dashboard tests."""
  return _AllSdkThirdPartyLibraryPaths() + _CatapultThirdPartyLibraryPaths() + [
      os.path.join(_CATAPULT_PATH, 'dashboard'),
      os.path.join(_CATAPULT_PATH, 'tracing'),
      os.path.join(_CATAPULT_PATH, 'common', 'py_utils', 'py_utils'),

      # Required by py_utils
      os.path.join(_CATAPULT_PATH, 'devil', 'devil'),

      # Isolate the sheriff_config package, since it's deployed independently.
      os.path.join(_CATAPULT_PATH, 'dashboard', 'dashboard', 'sheriff_config'),
  ]

def _AllSdkThirdPartyLibraryPaths():
  """Returns a list of all third party library paths from the SDK.

  The AppEngine documentation directs us to add App Engine libraries from the
  SDK to our Python path for local unit tests.
    https://cloud.google.com/appengine/docs/python/tools/localunittesting
  """
  paths = []
  for sdk_bin_path in os.environ['PATH'].split(os.pathsep):
    if 'google-cloud-sdk' not in sdk_bin_path:
      continue

    if not os.path.isdir(sdk_bin_path):
      sdk_bin_path = os.path.dirname(sdk_bin_path)

    appengine_path = os.path.join(sdk_bin_path, 'platform', 'google_appengine')
    paths.append(appengine_path)
    sys.path.insert(0, appengine_path)

  try:
    import dev_appserver
  except ImportError:
    # TODO: Put the Cloud SDK in the path with the binary dependency manager.
    # https://github.com/catapult-project/catapult/issues/2135
    print('This script requires the Google Cloud SDK to be in PATH.')
    print('Install at https://cloud.google.com/sdk and then run')
    print('`gcloud components install app-engine-python`')
    sys.exit(1)

  paths.extend(dev_appserver.EXTRA_PATHS)
  return paths

def _CatapultThirdPartyLibraryPaths():
  """Returns a list of required third-party libraries in catapult."""
  paths = []
  paths.append(
      os.path.join(_CATAPULT_PATH, 'common', 'node_runner', 'node_runner',
                   'node_modules', '@chopsui', 'tsmon-client',
                   'tsmon-client.js'))
  third_party_libraries = (
      THIRD_PARTY_LIBRARIES_PY3 if sys.version_info.major == 3
      else THIRD_PARTY_LIBRARIES_PY2)
  for library in third_party_libraries:
    paths.append(os.path.join(_CATAPULT_PATH, 'third_party', library))
  return paths

def _TracingPaths():
  """Returns a list of paths that may be imported from tracing."""
  # TODO(sullivan): This should either pull from tracing_project or be generated
  # via gypi. See https://github.com/catapult-project/catapult/issues/3048.
  paths = []
  for path in TRACING_PATHS:
    paths.append(os.path.join(_CATAPULT_PATH, os.path.normpath(path)))
  return paths
<|endoftext|>


<commit_before>
import os, json, unittest, time, shutil, sys
sys.path.extend(['.','..','py'])

import h2o, h2o_cmd, h2o_hosts, h2o_glm
import h2o_browse as h2b
import h2o_import as h2i
import time, random

def infoFromSummary(summary):
    columnsList = summary['columns']
    for columns in columnsList:
        N = columns['N']
        # self.assertEqual(N, rowCount)
        name = columns['name']
        stype = columns['type']
        histogram = columns['histogram']
        bin_size = histogram['bin_size']
        bin_names = histogram['bin_names']
        for b in bin_names:
            print "bin_name:", b

        bins = histogram['bins']
        nbins = histogram['bins']
        print "\n\n************************"
        print "N:", N
        print "name:", name
        print "type:", stype
        print "bin_size:", bin_size
        print "len(bin_names):", len(bin_names), bin_names
        print "len(bins):", len(bins), bins
        print "len(nbins):", len(nbins), nbins

        # not done if enum
        if stype != "enum":
            smax = columns['max']
            smin = columns['min']
            percentiles = columns['percentiles']
            thresholds = percentiles['thresholds']
            values = percentiles['values']
            mean = columns['mean']
            sigma = columns['sigma']

            print "len(max):", len(smax), smax
            print "len(min):", len(smin), smin
            print "len(thresholds):", len(thresholds), thresholds
            print "len(values):", len(values), values
            print "mean:", mean
            print "sigma:", sigma

            for v in values:
                self.assertTrue(v >= smin,
                    "Percentile value %s should all be >= the min dataset value %s" % (v, expectedMin))
                self.assertTrue(v <= smax,
                    "Percentile value %s should all be <= the max dataset value %s" % (v, expectedMax))

class Basic(unittest.TestCase):
    def tearDown(self):
        h2o.check_sandbox_for_errors()

    @classmethod
    def setUpClass(cls):
        # assume we're at 0xdata with it's hdfs namenode
        global localhost
        localhost = h2o.decide_if_localhost()
        if (localhost):
            h2o.build_cloud(1)
        else:
            # all hdfs info is done thru the hdfs_config michal's ec2 config sets up?
            h2o_hosts.build_cloud_with_hosts(1,
                # this is for our amazon ec hdfs
                # see https://github.com/0xdata/h2o/wiki/H2O-and-s3n
                hdfs_name_node='10.78.14.235:9000',
                hdfs_version='0.20.2')

    @classmethod
    def tearDownClass(cls):
        h2o.tear_down_cloud()

    def test_parse_summary_airline_s3n(self):
        URI = "s3n://h2o-airlines-unpacked/"
        csvFilelist = [
            ("allyears2k.csv",   300), #130.1MB
            ("allyears.csv",     900), #12GB
            ("allyears_10.csv", 1800), #119.98GB
        ]
        # IMPORT**********************************************
        # since H2O deletes the source key, we should re-import every iteration if we re-use the src in the list
        importHDFSResult = h2o.nodes[0].import_hdfs(URI)
        s3nFullList = importHDFSResult['files']
        ### print "s3nFullList:", h2o.dump_json(s3nFullList)

        self.assertGreater(len(s3nFullList),8,"Should see more than 8 files in s3n?")
        storeView = h2o.nodes[0].store_view()
        for s in storeView['keys']:
            print "\nkey:", s['key']
            if 'rows' in s:
                print "rows:", s['rows'], "value_size_bytes:", s['value_size_bytes']

        trial = 0
        for (csvFilename, timeoutSecs) in csvFilelist:
            trialStart = time.time()
            csvPathname = csvFilename
            s3nKey = URI + csvPathname

            # PARSE****************************************
            key2 = csvFilename + "_" + str(trial) + ".hex"
            print "Loading s3n key: ", s3nKey, 'thru HDFS'
            timeoutSecs = 500
            start = time.time()
            parseKey = h2o.nodes[0].parse(s3nKey, key2,
                timeoutSecs=timeoutSecs, retryDelaySecs=10, pollTimeoutSecs=60, noise=('JStack', none))
            elapsed = time.time() - start
            print "parse end on ", s3nKey, 'took', elapsed, 'seconds',\
                "%d pct. of timeout" % ((elapsed*100)/timeoutSecs)
            print "parse result:", parseKey['destination_key']

            # INSPECT******************************************
            # We should be able to see the parse result?
            start = time.time()
            inspect = h2o_cmd.runInspect(None, parseKey['destination_key'], timeoutSecs=500)
            print "Inspect:", parseKey['destination_key'], "took", time.time() - start, "seconds"
            h2o_cmd.infoFromInspect(inspect, csvPathname)
            num_rows = inspect['num_rows']
            num_cols = inspect['num_cols']

            (missingValuesDict, constantValuesDict, enumSizeDict, colTypeDict, colNameDict) = \
                h2o_cmd.get_column_info_from_inspect(parseKey, timeoutSecs=300)

            if missingValuesDict:
                print len(missingValuesDict), "columns with missing values"
                m = [str(k) + ":" + str(v) for k,v in missingValuesDict.iteritems()]
                print "Maybe columns got flipped to NAs: " + ", ".join(m)
                ### raise Exception("Looks like columns got flipped to NAs: " + ", ".join(m))

            if constantValuesDict:
                print len(constantValuesDict), "columns with constant values"

            print "\n" + csvPathname, \
                "    num_rows:", "{:,}".format(num_rows), \
                "    num_cols:", "{:,}".format(num_cols)

            # SUMMARY****************************************
            summaryResult = h2o.nodes[0].summary_page(key2)
            summary = summaryResult['summary']
            # print h2o.dump_json(summary)
            infoFromSummary(summary)

            # STOREVIEW***************************************
            storeView = h2o.nodes[0].store_view()
            for s in storeView['keys']:
                print "\nStoreView: key:", s['key']
                if 'rows' in s:
                    print "StoreView: rows:", s['rows'], "value_size_bytes:", s['value_size_bytes']

            print "Trial #", trial, "completed in", time.time() - trialStart, "seconds."
            trial += 0

if __name__ == '__main__':
    h2o.unit_main()
<commit_msg>updated for airline ec2 testing<commit_after>
import os, json, unittest, time, shutil, sys
sys.path.extend(['.','..','py'])

import h2o, h2o_cmd, h2o_hosts, h2o_glm
import h2o_browse as h2b
import h2o_import as h2i
import time, random

def infoFromSummary(self,summary):
    columnsList = summary['columns']
    for columns in columnsList:
        N = columns['N']
        # self.assertEqual(N, rowCount)
        name = columns['name']
        stype = columns['type']
        histogram = columns['histogram']
        bin_size = histogram['bin_size']
        bin_names = histogram['bin_names']
        for b in bin_names:
            print "bin_name:", b

        bins = histogram['bins']
        nbins = histogram['bins']
        print "\n\n************************"
        print "N:", N
        print "name:", name
        print "type:", stype
        print "bin_size:", bin_size
        print "len(bin_names):", len(bin_names), bin_names
        print "len(bins):", len(bins), bins
        print "len(nbins):", len(nbins), nbins

        # not done if enum
        if stype != "enum":
            smax = columns['max']
            smin = columns['min']
            mean = columns['mean']
            sigma = columns['sigma']
            print "smax:", smax
            print "smin:", smin
            print "mean:", mean
            print "sigma:", sigma

            # sometimes we don't get percentiles?
            if len(bins)!=0:
                percentiles = columns['percentiles']
                thresholds = percentiles['thresholds']
                values = percentiles['values']

                # h2o shows 5 of them, ordered
                print "len(max):", len(smax), smax
                print "len(min):", len(smin), smin
                print "len(thresholds):", len(thresholds), thresholds
                print "len(values):", len(values), values

                for v in values:
                    # 0 is the most max or most min
                    self.assertTrue(v >= smin[0],
                        "Percentile value %s should all be >= the min dataset value %s" % (v, smin[0]))
                    self.assertTrue(v <= smax[0],
                        "Percentile value %s should all be <= the max dataset value %s" % (v, smax[0]))

class Basic(unittest.TestCase):
    def tearDown(self):
        h2o.check_sandbox_for_errors()

    @classmethod
    def setUpClass(cls):
        # assume we're at 0xdata with it's hdfs namenode
        global localhost
        localhost = h2o.decide_if_localhost()
        if (localhost):
            h2o.build_cloud(1)
        else:
            # all hdfs info is done thru the hdfs_config michal's ec2 config sets up?
            h2o_hosts.build_cloud_with_hosts(1,
                # this is for our amazon ec hdfs
                # see https://github.com/0xdata/h2o/wiki/H2O-and-s3n
                hdfs_name_node='10.78.14.235:9000',
                hdfs_version='0.20.2')

    @classmethod
    def tearDownClass(cls):
        h2o.tear_down_cloud()

    def test_parse_summary_airline_s3n(self):
        URI = "s3n://h2o-airlines-unpacked/"
        csvFilelist = [
            ("allyears2k.csv",   300), #130.1MB
            ("allyears.csv",     900), #12GB
            ("allyears_10.csv", 1800), #119.98GB
        ]
        # IMPORT**********************************************
        # since H2O deletes the source key, we should re-import every iteration if we re-use the src in the list
        importHDFSResult = h2o.nodes[0].import_hdfs(URI)
        ### print "importHDFSResult:", h2o.dump_json(importHDFSResult)
        s3nFullList = importHDFSResult['succeeded']
        ### print "s3nFullList:", h2o.dump_json(s3nFullList)

        self.assertGreater(len(s3nFullList),8,"Should see more than 8 files in s3n?")
        # why does this hang?
        if 1==0:
            storeView = h2o.nodes[0].store_view(timeoutSecs=60)
            for s in storeView['keys']:
                print "\nkey:", s['key']
                if 'rows' in s:
                    print "rows:", s['rows'], "value_size_bytes:", s['value_size_bytes']

        trial = 0
        for (csvFilename, timeoutSecs) in csvFilelist:
            trialStart = time.time()
            csvPathname = csvFilename
            s3nKey = URI + csvPathname

            # PARSE****************************************
            key2 = csvFilename + "_" + str(trial) + ".hex"
            print "Loading s3n key: ", s3nKey, 'thru HDFS'
            start = time.time()
            parseKey = h2o.nodes[0].parse(s3nKey, key2,
                timeoutSecs=timeoutSecs, retryDelaySecs=10, pollTimeoutSecs=120)
            elapsed = time.time() - start
            print "parse end on ", s3nKey, 'took', elapsed, 'seconds',\
                "%d pct. of timeout" % ((elapsed*100)/timeoutSecs)
            print "parse result:", parseKey['destination_key']

            # INSPECT******************************************
            # We should be able to see the parse result?
            start = time.time()
            inspect = h2o_cmd.runInspect(None, parseKey['destination_key'], timeoutSecs=500)
            print "Inspect:", parseKey['destination_key'], "took", time.time() - start, "seconds"
            h2o_cmd.infoFromInspect(inspect, csvPathname)
            num_rows = inspect['num_rows']
            num_cols = inspect['num_cols']

            (missingValuesDict, constantValuesDict, enumSizeDict, colTypeDict, colNameDict) = \
                h2o_cmd.get_column_info_from_inspect(parseKey, timeoutSecs=300)

            if missingValuesDict:
                print len(missingValuesDict), "columns with missing values"
                m = [str(k) + ":" + str(v) for k,v in missingValuesDict.iteritems()]
                print "Maybe columns got flipped to NAs: " + ", ".join(m)
                ### raise Exception("Looks like columns got flipped to NAs: " + ", ".join(m))

            if constantValuesDict:
                print len(constantValuesDict), "columns with constant values"
                m = [str(k) + ":" + str(v) for k,v in constantValuesDict.iteritems()]
                print "constant columns: " + ", ".join(m)

            print "\n" + csvPathname, \
                "    num_rows:", "{:,}".format(num_rows), \
                "    num_cols:", "{:,}".format(num_cols)

            # SUMMARY****************************************
            summaryResult = h2o.nodes[0].summary_page(key2)
            summary = summaryResult['summary']
            # print h2o.dump_json(summary)
            infoFromSummary(self, summary)

            # STOREVIEW***************************************
            storeView = h2o.nodes[0].store_view()
            for s in storeView['keys']:
                print "\nStoreView: key:", s['key']
                if 'rows' in s:
                    print "StoreView: rows:", s['rows'], "value_size_bytes:", s['value_size_bytes']

            print "Trial #", trial, "completed in", time.time() - trialStart, "seconds."
            trial += 0

if __name__ == '__main__':
    h2o.unit_main()
<|endoftext|>


<commit_before>import requests
import re
from urllib.parse import quote_plus

TYPE_MAP = {
    3: 'AR',
    1: 'SP',
}

class SentenceTooBigError(Exception):
    pass

def split_sentences(text):
    """Splits text into sentences, whilist keeping all original punctuation.
    Thus this is the expected output:

        I have an elephant! Its name is Bob. -> ["I have an elephant!", "Its name is Bob."]
    """
    tokens = re.split(r'([.!?]+)', text)
    # If the last token is '' it means the last sentence ends in punctuation.
    # So we remove the uneeded empty string in the end. Otherwise, we add an
    # empty string so the for ... in zip(...) works later on.
    if tokens[-1] == '':
        tokens.pop()
    else:
        tokens.append('')
    sentences = []
    for sentence, punct in zip(tokens[::2], tokens[1::2]):
        sentences.append(sentence + punct)
    return sentences

def merge_sentences(sentences, max_size=600):
    """Merges phrases together up to max_size. If a sentence is too big
    SentenceTooBigError is raised"""
    result = []
    current = ''
    for sent in sentences:
        if len(sent) > max_size:
            raise SentenceTooBigError()
        new_sent = current + sent
        if len(new_sent) > max_size:
            result.append(current)
            current = sent
        else:
            current += sent
    result.append(current)
    return result

def _get_type(correction, text):
    word = _get_word(correction, text)
    type_ = TYPE_MAP.get(correction["Type"], "undefined")
    if type_ == 'SP' and any(c['Text'].lower() == word for c in correction['Suggestions']):
        type_ = 'C'
    return type_

def _get_word(correction, text):
    return text[correction["From"]:correction["To"] + 1]

def _get_suggestions(correction):
    if correction['Suggestions']:
        return [suggestion['Text'] for suggestion in correction['Suggestions']]
    return []

def check(text):
    """Given a text, query ginger API and return a list of dicts in the same
    format as righter.check"""
    url = 'https://services.gingersoftware.com/Ginger/correct/jsonSecured/GingerTheTextFull'
    # Ginger splits texts on sentences. So, if you have a really big
    # sentence, even ginger plugin will fail. It does make sense to split by
    # sentences because there are sentence-wide analysis that could be made.
    #
    # The largest sentence ginger can handle is 600 characters.
    offset = 0
    params = {
        "isOnTheFly": "false",
        "securedPage": "true",
        "apiKey": "BrowserStandalone",
        "lang": "US",
        "clientVersion": "0.1.0.562",
        "_": "1460102265093"
    }
    changes = []
    for sentence in merge_sentences(split_sentences(text)):
        params["text"] = sentence
        resp = requests.get(url, params=params)
        resp.raise_for_status()
        response = resp.json()
        for correction in response["Corrections"]:
            change = {
                'selection': _get_word(correction, sentence),
                'start': correction['From'] + offset,
                'symbol': _get_type(correction, sentence),
                'originalSymbol': correction["Type"],
            }
            change['suggestions'] = _get_suggestions(correction)
            if change['suggestions']:
                change['correction'] = change['suggestions'][0]
            changes.append(change)
        offset += len(sentence)

    return changes
<commit_msg>Save requests to a file<commit_after>import re
import requests_cache
from urllib.parse import quote_plus

TYPE_MAP = {
    3: 'AR',
    1: 'SP',
}

session = requests_cache.CachedSession('ginger')

class SentenceTooBigError(Exception):
    pass

def split_sentences(text):
    """Splits text into sentences, whilist keeping all original punctuation.
    Thus this is the expected output:

        I have an elephant! Its name is Bob. -> ["I have an elephant!", "Its name is Bob."]
    """
    tokens = re.split(r'([.!?]+)', text)
    # If the last token is '' it means the last sentence ends in punctuation.
    # So we remove the uneeded empty string in the end. Otherwise, we add an
    # empty string so the for ... in zip(...) works later on.
    if tokens[-1] == '':
        tokens.pop()
    else:
        tokens.append('')
    sentences = []
    for sentence, punct in zip(tokens[::2], tokens[1::2]):
        sentences.append(sentence + punct)
    return sentences

def merge_sentences(sentences, max_size=600):
    """Merges phrases together up to max_size. If a sentence is too big
    SentenceTooBigError is raised"""
    result = []
    current = ''
    for sent in sentences:
        if len(sent) > max_size:
            raise SentenceTooBigError()
        new_sent = current + sent
        if len(new_sent) > max_size:
            result.append(current)
            current = sent
        else:
            current += sent
    result.append(current)
    return result

def _get_type(correction, text):
    word = _get_word(correction, text)
    type_ = TYPE_MAP.get(correction["Type"], "undefined")
    if type_ == 'SP' and any(c['Text'].lower() == word for c in correction['Suggestions']):
        type_ = 'C'
    return type_

def _get_word(correction, text):
    return text[correction["From"]:correction["To"] + 1]

def _get_suggestions(correction):
    if correction['Suggestions']:
        return [suggestion['Text'] for suggestion in correction['Suggestions']]
    return []

def check(text):
    """Given a text, query ginger API and return a list of dicts in the same
    format as righter.check"""
    url = 'https://services.gingersoftware.com/Ginger/correct/jsonSecured/GingerTheTextFull'
    # Ginger splits texts on sentences. So, if you have a really big
    # sentence, even ginger plugin will fail. It does make sense to split by
    # sentences because there are sentence-wide analysis that could be made.
    #
    # The largest sentence ginger can handle is 600 characters.
    offset = 0
    params = {
        "isOnTheFly": "false",
        "securedPage": "true",
        "apiKey": "BrowserStandalone",
        "lang": "US",
        "clientVersion": "0.1.0.562",
        "_": "1460102265093"
    }
    changes = []
    for sentence in merge_sentences(split_sentences(text)):
        params["text"] = sentence
        resp = session.get(url, params=params)
        resp.raise_for_status()
        response = resp.json()
        for correction in response["Corrections"]:
            change = {
                'selection': _get_word(correction, sentence),
                'start': correction['From'] + offset,
                'symbol': _get_type(correction, sentence),
                'originalSymbol': correction["Type"],
            }
            change['suggestions'] = _get_suggestions(correction)
            if change['suggestions']:
                change['correction'] = change['suggestions'][0]
            changes.append(change)
        offset += len(sentence)

    return changes
<|endoftext|>