<commit_before>"""Main PAC methods.

This file include the following methods :
- Mean Vector Length (Canolty, 2006)
- Kullback Leibler Distance (Tort, 2010)
- Heights Ratio (Lakata, 2005)
- Normalized direct Pac (Ozkurt, 2012)
- Phase Synchrony
"""

import numpy as np
from scipy.special import erfinv
from .stats import circ_corrcc

__all__ = ['ComputePac']

def ComputePac(pha, amp, idp, nbins, p):
    """Copute real Phase-Amplitude coupling.

    Each method take at least a pha and amp array with the respective
    dimensions:
    pha.shape = (npha, ..., npts)
    amp.shape = (namp, ..., npts)
    And each method should return a (namp, npha, ...)
    """
    # Mean Vector Length (Canolty, 2006)
    if idp == 1:
        return MVL(pha, amp)

    # Kullback-Leiber divergence (Tort, 2010)
    elif idp == 2:
        return klDivergence(pha, amp, nbins)

    # Heights ratio
    elif idp == 3:
        return HeightsRatio(pha, amp, nbins)

    # ndPac (Ozkurt, 2012)
    elif idp == 4:
        return ndPac(pha, amp, p)

    # Phase-Synchrony (or adapted PLV)
    elif idp == 5:
        return PhaseSync(pha, amp)

    # Event-Related Phase amplitude Coupling
    elif idp == 6:
        # return erpac(pha, amp, p)
        raise(NotImplementedError)

    else:
        raise ValueError(str(idp) + " is not recognized as a valid pac "
                         "method.")

def MVL(pha, amp):
    """Mean Vector Length (Canolty, 2006).

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Number of time points :
    npts = pha.shape[-1]
    return np.abs(np.einsum('i...j, k...j->ik...', amp, np.exp(1j*pha)))/npts

def klDistance(pha, amp, nbins):
    """Kullback Leibler Distance (Tort, 2010).

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

        nbins: int
            Number of bins in which the phase in cut in bins.

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Get the phase locked binarized amplitude :
    abin = _kl_hr(pha, amp, nbins)
    # Divide the binned amplitude by the mean over the bins :
    abin = np.divide(abin, np.sum(abin, axis=0))
    abin[abin == 0] = 1
    abin = abin * np.log2(abin)

    return (1 + abin.sum(axis=0)/np.log2(nbins))

def HeightsRatio(pha, amp, nbins):
    """Pac Heights Ratio.

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

        nbins: int
            Number of bins in which the phase in cut in bins.

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Get the phase locked binarized amplitude :
    abin = _kl_hr(pha, amp, nbins)
    M, m = abin.max(axis=0), abin.min(axis=0)
    MDown = M.copy()
    MDown[MDown == 0] = 1

    return (M-m)/MDown

def _kl_hr(pha, amp, nbins):
    """Binarize the amplitude according to phase values.

    This function is shared by the Kullback-Leibler Distance and the
    Height Ratio.
    """
    # Build the default binned phase vector :
    step = 2*np.pi/nbins
    x = np.arange(-np.pi, np.pi+step, step)
    vecbin = np.c_[x[0:-1], x[1::]]

    abin = []
    for i in vecbin:
        # Find where phase take vecbin values :
        idx = np.logical_and(pha >= i[0], pha < i[1])
        # Take the sum of amplitude inside the bin :
        abin.append(np.einsum('i...j, k...j->ik...', amp, idx))

    return np.array(abin)

def ndPac(pha, amp, p):
    """Normalized direct Pac (Ozkurt, 2012).

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

        p: float
            The p-value to use.

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    npts = amp.shape[-1]
    # Normalize amplitude :
    np.subtract(amp, np.mean(amp, axis=-1, keepdims=True), out=amp)
    np.divide(amp, np.std(amp, axis=-1, keepdims=True), out=amp)
    # Compute pac :
    pac = np.abs(np.einsum('i...j, k...j->ik...', amp, np.exp(1j*pha)))
    pac *= pac/npts
    # Set to zero non-significant values:
    xlim = erfinv(1-p)**2
    pac[pac <= 2 * xlim] = 0.
    return pac

def PhaseSync(pha, amp):
    """Phase Synchrony.

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Number of time points :
    npts = pha.shape[-1]
    pac = np.einsum('i...j, k...j->ik...', np.exp(-1j*amp), np.exp(1j*pha))
    return np.abs(pac)/npts

def erpac(pha, amp, p):
    pha = np.swapaxes(pha, -1, -2)
    amp = np.swapaxes(amp, -1, -2)
    return circ_corrcc(pha, amp)
<commit_msg>Divergence -> distance<commit_after>"""Main PAC methods.

This file include the following methods :
- Mean Vector Length (Canolty, 2006)
- Kullback Leibler Distance (Tort, 2010)
- Heights Ratio (Lakata, 2005)
- Normalized direct Pac (Ozkurt, 2012)
- Phase Synchrony
"""

import numpy as np
from scipy.special import erfinv
from .stats import circ_corrcc

__all__ = ['ComputePac']

def ComputePac(pha, amp, idp, nbins, p):
    """Copute real Phase-Amplitude coupling.

    Each method take at least a pha and amp array with the respective
    dimensions:
    pha.shape = (npha, ..., npts)
    amp.shape = (namp, ..., npts)
    And each method should return a (namp, npha, ...)
    """
    # Mean Vector Length (Canolty, 2006)
    if idp == 1:
        return MVL(pha, amp)

    # Kullback-Leiber distance (Tort, 2010)
    elif idp == 2:
        return klDistance(pha, amp, nbins)

    # Heights ratio
    elif idp == 3:
        return HeightsRatio(pha, amp, nbins)

    # ndPac (Ozkurt, 2012)
    elif idp == 4:
        return ndPac(pha, amp, p)

    # Phase-Synchrony (or adapted PLV)
    elif idp == 5:
        return PhaseSync(pha, amp)

    # Event-Related Phase amplitude Coupling
    elif idp == 6:
        # return erpac(pha, amp, p)
        raise(NotImplementedError)

    else:
        raise ValueError(str(idp) + " is not recognized as a valid pac "
                         "method.")

def MVL(pha, amp):
    """Mean Vector Length (Canolty, 2006).

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Number of time points :
    npts = pha.shape[-1]
    return np.abs(np.einsum('i...j, k...j->ik...', amp, np.exp(1j*pha)))/npts

def klDistance(pha, amp, nbins):
    """Kullback Leibler Distance (Tort, 2010).

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

        nbins: int
            Number of bins in which the phase in cut in bins.

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Get the phase locked binarized amplitude :
    abin = _kl_hr(pha, amp, nbins)
    # Divide the binned amplitude by the mean over the bins :
    abin = np.divide(abin, np.sum(abin, axis=0))
    abin[abin == 0] = 1
    abin = abin * np.log2(abin)

    return (1 + abin.sum(axis=0)/np.log2(nbins))

def HeightsRatio(pha, amp, nbins):
    """Pac Heights Ratio.

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

        nbins: int
            Number of bins in which the phase in cut in bins.

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Get the phase locked binarized amplitude :
    abin = _kl_hr(pha, amp, nbins)
    M, m = abin.max(axis=0), abin.min(axis=0)
    MDown = M.copy()
    MDown[MDown == 0] = 1

    return (M-m)/MDown

def _kl_hr(pha, amp, nbins):
    """Binarize the amplitude according to phase values.

    This function is shared by the Kullback-Leibler Distance and the
    Height Ratio.
    """
    # Build the default binned phase vector :
    step = 2*np.pi/nbins
    x = np.arange(-np.pi, np.pi+step, step)
    vecbin = np.c_[x[0:-1], x[1::]]

    abin = []
    for i in vecbin:
        # Find where phase take vecbin values :
        idx = np.logical_and(pha >= i[0], pha < i[1])
        # Take the sum of amplitude inside the bin :
        abin.append(np.einsum('i...j, k...j->ik...', amp, idx))

    return np.array(abin)

def ndPac(pha, amp, p):
    """Normalized direct Pac (Ozkurt, 2012).

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

        p: float
            The p-value to use.

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    npts = amp.shape[-1]
    # Normalize amplitude :
    np.subtract(amp, np.mean(amp, axis=-1, keepdims=True), out=amp)
    np.divide(amp, np.std(amp, axis=-1, keepdims=True), out=amp)
    # Compute pac :
    pac = np.abs(np.einsum('i...j, k...j->ik...', amp, np.exp(1j*pha)))
    pac *= pac/npts
    # Set to zero non-significant values:
    xlim = erfinv(1-p)**2
    pac[pac <= 2 * xlim] = 0.
    return pac

def PhaseSync(pha, amp):
    """Phase Synchrony.

    Args:
        pha: np.ndarray
            Array of phases of shapes (npha, ..., npts)

        amp: np.ndarra
            Array of amplitudes of shapes (namp, ..., npts)

    Return:
        pac: np.ndarray
            PAC of shape (npha, namp, ...)
    """
    # Number of time points :
    npts = pha.shape[-1]
    pac = np.einsum('i...j, k...j->ik...', np.exp(-1j*amp), np.exp(1j*pha))
    return np.abs(pac)/npts

def erpac(pha, amp, p):
    pha = np.swapaxes(pha, -1, -2)
    amp = np.swapaxes(amp, -1, -2)
    return circ_corrcc(pha, amp)
<|endoftext|>


<commit_before>from django.conf import settings
from django.db import models, connection, transaction
from django.db.models.signals import post_delete
from django.dispatch import receiver
from django.core.management import call_command
from django.utils.functional import lazy
from tenant_schemas.postgresql_backend.base import _check_identifier
from tenant_schemas.signals import post_schema_sync
from tenant_schemas.utils import django_is_in_test_mode, schema_exists, get_tenant_model
from tenant_schemas.utils import get_public_schema_name

class TenantMixin(models.Model):
    auto_drop_schema = False   # USE THIS WITH CAUTION!
                               # set this flag to true on a parent class if
                               # you want the schema to be automatically
                               # removed after tenant remove.

    auto_create_schema = True  # set this flag to false on a parent class if
                               # you dont want the schema to be automatically
                               # created upon save.

    domain_url = models.CharField(max_length=128, unique=True)
    schema_name = models.CharField(max_length=63)

    class Meta:
        abstract = True

    def save(self, verbosity=1, *args, **kwargs):
        is_new = self.pk is None

        if is_new and connection.get_schema() != get_public_schema_name():
            raise Exception("Can't update tenant outside the public schema. Current schema is %s."
                            % connection.get_schema())
        elif not is_new and connection.get_schema() not in (self.schema_name, get_public_schema_name()):
            raise Exception("Can't update tenant outside it's own schema or the public schema. Current schema is %s."
                            % connection.get_schema())

        super(TenantMixin, self).save(*args, **kwargs)

        if is_new and self.auto_create_schema:
            self.create_schema(check_if_exists=True, verbosity=verbosity)
            post_schema_sync.send(sender=TenantMixin, tenant=self)

        transaction.commit_unless_managed()

    def delete(self, *args, **kwargs):
        if connection.get_schema() not in (self.schema_name, get_public_schema_name()):
            raise Exception("Can't delete tenant outside it's own schema or the public schema. Current schema is %s."
                            % connection.get_schema())

        super(TenantMixin, self).delete(*args, **kwargs)

    def create_schema(self, check_if_exists=False, sync_schema=True, verbosity=1):
        """
        Creates the schema 'schema_name' for this tenant. Optionally checks if the schema
        already exists before creating it. Returns true if the schema was created, false
        otherwise.
        """

        # safety check
        _check_identifier(self.schema_name)
        cursor = connection.cursor()

        if check_if_exists and schema_exists(self.schema_name):
            return False

        # create the schema
        cursor.execute('CREATE SCHEMA %s' % self.schema_name)

        if sync_schema:
            call_command('sync_schemas',
                         schema_name=self.schema_name,
                         tenant=True,
                         public=False,
                         interactive=False,  # don't ask to create an admin user
                         migrate_all=True,  # migrate all apps directly to last version
                         verbosity=verbosity,
                         )

            # fake all migrations
            if 'south' in settings.INSTALLED_APPS and not django_is_in_test_mode():
                call_command('migrate_schemas', fake=True, schema_name=self.schema_name, verbosity=verbosity)

        connection.set_schema_to_public()

        return True

@receiver(post_delete, sender=lazy(get_tenant_model))
def drop_schema(sender, instance, **kwargs):
    """
    Called in post_delete signal.
    Drops the schema related to the tenant instance. Just drop the schema if the parent
    class model has the attribute auto_drop_schema set to True.
    """
    if schema_exists(instance.schema_name) and instance.auto_drop_schema:
        cursor = connection.cursor()
        cursor.execute('DROP SCHEMA %s CASCADE' % instance.schema_name)
<commit_msg>updated string message when saving tenant outside public<commit_after>from django.conf import settings
from django.db import models, connection, transaction
from django.db.models.signals import post_delete
from django.dispatch import receiver
from django.core.management import call_command
from django.utils.functional import lazy
from tenant_schemas.postgresql_backend.base import _check_identifier
from tenant_schemas.signals import post_schema_sync
from tenant_schemas.utils import django_is_in_test_mode, schema_exists, get_tenant_model
from tenant_schemas.utils import get_public_schema_name

class TenantMixin(models.Model):
    auto_drop_schema = False   # USE THIS WITH CAUTION!
                               # set this flag to true on a parent class if
                               # you want the schema to be automatically
                               # removed after tenant remove.

    auto_create_schema = True  # set this flag to false on a parent class if
                               # you dont want the schema to be automatically
                               # created upon save.

    domain_url = models.CharField(max_length=128, unique=True)
    schema_name = models.CharField(max_length=63)

    class Meta:
        abstract = True

    def save(self, verbosity=1, *args, **kwargs):
        is_new = self.pk is None

        if is_new and connection.get_schema() != get_public_schema_name():
            raise Exception("Can't create tenant outside the public schema. Current schema is %s."
                            % connection.get_schema())
        elif not is_new and connection.get_schema() not in (self.schema_name, get_public_schema_name()):
            raise Exception("Can't update tenant outside it's own schema or the public schema. Current schema is %s."
                            % connection.get_schema())

        super(TenantMixin, self).save(*args, **kwargs)

        if is_new and self.auto_create_schema:
            self.create_schema(check_if_exists=True, verbosity=verbosity)
            post_schema_sync.send(sender=TenantMixin, tenant=self)

        transaction.commit_unless_managed()

    def delete(self, *args, **kwargs):
        if connection.get_schema() not in (self.schema_name, get_public_schema_name()):
            raise Exception("Can't delete tenant outside it's own schema or the public schema. Current schema is %s."
                            % connection.get_schema())

        super(TenantMixin, self).delete(*args, **kwargs)

    def create_schema(self, check_if_exists=False, sync_schema=True, verbosity=1):
        """
        Creates the schema 'schema_name' for this tenant. Optionally checks if the schema
        already exists before creating it. Returns true if the schema was created, false
        otherwise.
        """

        # safety check
        _check_identifier(self.schema_name)
        cursor = connection.cursor()

        if check_if_exists and schema_exists(self.schema_name):
            return False

        # create the schema
        cursor.execute('CREATE SCHEMA %s' % self.schema_name)

        if sync_schema:
            call_command('sync_schemas',
                         schema_name=self.schema_name,
                         tenant=True,
                         public=False,
                         interactive=False,  # don't ask to create an admin user
                         migrate_all=True,  # migrate all apps directly to last version
                         verbosity=verbosity,
                         )

            # fake all migrations
            if 'south' in settings.INSTALLED_APPS and not django_is_in_test_mode():
                call_command('migrate_schemas', fake=True, schema_name=self.schema_name, verbosity=verbosity)

        connection.set_schema_to_public()

        return True

@receiver(post_delete, sender=lazy(get_tenant_model))
def drop_schema(sender, instance, **kwargs):
    """
    Called in post_delete signal.
    Drops the schema related to the tenant instance. Just drop the schema if the parent
    class model has the attribute auto_drop_schema set to True.
    """
    if schema_exists(instance.schema_name) and instance.auto_drop_schema:
        cursor = connection.cursor()
        cursor.execute('DROP SCHEMA %s CASCADE' % instance.schema_name)
<|endoftext|>


<commit_before>from copy import deepcopy

import settingshelper as helper
from settings import *

USING_CITUS = any(db.get('ROLE') == 'citus_master' for db in DATABASES.values())

# note: the only reason these are prepended to INSTALLED_APPS is because of
# a weird travis issue with kafka. if for any reason this order causes problems
# it can be reverted whenever that's figured out.
# https://github.com/dimagi/commcare-hq/pull/10034#issuecomment-174868270
INSTALLED_APPS = (
    'django_nose',
    'testapps.test_elasticsearch',
    'testapps.test_pillowtop',
) + tuple(INSTALLED_APPS)

if USING_CITUS:
    if 'testapps.citus_master' not in INSTALLED_APPS:
        INSTALLED_APPS = (
            'testapps.citus_master',
            'testapps.citus_worker',
        ) + tuple(INSTALLED_APPS)

    if 'testapps.citus_master.citus_router.CitusDBRouter' not in DATABASE_ROUTERS:
        # this router must go first
        DATABASE_ROUTERS = ['testapps.citus_master.citus_router.CitusDBRouter'] + DATABASE_ROUTERS

TEST_RUNNER = 'django_nose.BasicNoseRunner'
NOSE_ARGS = [
    #'--no-migrations' # trim ~120s from test run with db tests
    #'--with-fixture-bundling',
]
NOSE_PLUGINS = [
    'corehq.tests.nose.HqTestFinderPlugin',
    'corehq.tests.noseplugins.dividedwerun.DividedWeRunPlugin',
    'corehq.tests.noseplugins.djangomigrations.DjangoMigrationsPlugin',
    'corehq.tests.noseplugins.cmdline_params.CmdLineParametersPlugin',
    'corehq.tests.noseplugins.patches.PatchesPlugin',
    'corehq.tests.noseplugins.redislocks.RedisLockTimeoutPlugin',
    'corehq.tests.noseplugins.uniformresult.UniformTestResultPlugin',

    # The following are not enabled by default
    'corehq.tests.noseplugins.logfile.LogFilePlugin',
    'corehq.tests.noseplugins.timing.TimingPlugin',
    'corehq.tests.noseplugins.output.OutputPlugin',

    # Uncomment to debug tests. Plugins have nice hooks for inspecting state
    # before/after each test or context setup/teardown, etc.
    #'corehq.tests.noseplugins.debug.DebugPlugin',
]

# these settings can be overridden with environment variables
for key, value in {
    'NOSE_DB_TEST_CONTEXT': 'corehq.tests.nose.HqdbContext',
    'NOSE_NON_DB_TEST_CONTEXT': 'corehq.tests.nose.ErrorOnDbAccessContext',

    'NOSE_IGNORE_FILES': '^localsettings',

    'NOSE_EXCLUDE_DIRS': ';'.join([
        'scripts',
    ]),
}.items():
    os.environ.setdefault(key, value)
del key, value

if "SKIP_TESTS_REQUIRING_EXTRA_SETUP" not in globals():
    SKIP_TESTS_REQUIRING_EXTRA_SETUP = False

CELERY_TASK_ALWAYS_EAGER = True
# keep a copy of the original PILLOWTOPS setting around in case other tests want it.
_PILLOWTOPS = PILLOWTOPS
PILLOWTOPS = {}

# required by auditcare tests
AUDIT_MODEL_SAVE = ['django.contrib.auth.models.User']
AUDIT_ADMIN_VIEWS = False

PHONE_TIMEZONES_HAVE_BEEN_PROCESSED = True
PHONE_TIMEZONES_SHOULD_BE_PROCESSED = True

ENABLE_PRELOGIN_SITE = True

# override dev_settings
CACHE_REPORTS = True

def _set_logging_levels(levels):
    import logging
    for path, level in levels.items():
        logging.getLogger(path).setLevel(level)
_set_logging_levels({
    # Quiet down noisy loggers. Selective removal can be handy for debugging.
    'alembic': 'WARNING',
    'auditcare': 'INFO',
    'boto3': 'WARNING',
    'botocore': 'INFO',
    'couchdbkit.request': 'INFO',
    'couchdbkit.designer': 'WARNING',
    'datadog': 'WARNING',
    'elasticsearch': 'ERROR',
    'kafka.conn': 'WARNING',
    'kafka.client': 'WARNING',
    'kafka.consumer.kafka': 'WARNING',
    'kafka.metrics': 'WARNING',
    'kafka.protocol.parser': 'WARNING',
    'kafka.producer': 'WARNING',
    'quickcache': 'INFO',
    'requests.packages.urllib3': 'WARNING',
    's3transfer': 'INFO',
    'urllib3': 'WARNING',
})

# use empty LOGGING dict with --debug=nose,nose.plugins to debug test discovery
# TODO empty logging config (and fix revealed deprecation warnings)
LOGGING = {
    'disable_existing_loggers': False,
    'version': 1,
    'loggers': {},
}

# Default custom databases to use the same configuration as the default
# This is so that all devs don't have to run citus locally
if 'icds-ucr' not in DATABASES:
    DATABASES['icds-ucr'] = deepcopy(DATABASES['default'])
    # use a different name otherwise migrations don't get run
    DATABASES['icds-ucr']['NAME'] = 'commcarehq_icds_ucr'
    del DATABASES['icds-ucr']['TEST']['NAME']  # gets set by `helper.assign_test_db_names`

helper.assign_test_db_names(DATABASES)

# See comment under settings.SMS_QUEUE_ENABLED
SMS_QUEUE_ENABLED = False

# use all providers in tests
METRICS_PROVIDERS = [
    'corehq.util.metrics.datadog.DatadogMetrics',
    'corehq.util.metrics.prometheus.PrometheusMetrics',
]

# timeout faster in tests
ES_SEARCH_TIMEOUT = 5

if os.path.exists("custom/icds"):
    STATIC_DATA_SOURCES.extend([
        "custom/icds_reports/ucr/data_sources/*.json",
        "custom/icds_reports/ucr/data_sources/dashboard/*.json"
    ])
    STATIC_DATA_SOURCES = list(set(STATIC_DATA_SOURCES))

    STATIC_UCR_REPORTS.extend([
        "custom/icds_reports/ucr/reports/dashboard/*.j,son",
        "custom/icds_reports/ucr/reports/asr/*.json",
        "custom/icds_reports/ucr/reports/asr/ucr_v2/*.json",
        "custom/icds_reports/ucr/reports/mpr/*.json",
        "custom/icds_reports/ucr/reports/mpr/dashboard/*.json",
        "custom/icds_reports/ucr/reports/ls/*.json",
        "custom/icds_reports/ucr/reports/other/*.json",
    ])

    CUSTOM_UCR_EXPRESSION_LISTS.append("custom.icds_reports.ucr.expressions.CUSTOM_UCR_EXPRESSIONS")
    CUSTOM_UCR_EXPRESSION_LISTS = list(set(CUSTOM_UCR_EXPRESSION_LISTS))

    CUSTOM_UCR_REPORT_FILTERS.append(
        ('village_choice_list', 'custom.icds_reports.ucr.filter_spec.build_village_choice_list_filter_spec'),
    )
    CUSTOM_UCR_REPORT_FILTERS = list(set(CUSTOM_UCR_REPORT_FILTERS))
    CUSTOM_UCR_REPORT_FILTER_VALUES.append(
        ('village_choice_list', 'custom.icds_reports.ucr.filter_value.VillageFilterValue'),
    )
    CUSTOM_UCR_REPORT_FILTER_VALUES = list(set(CUSTOM_UCR_REPORT_FILTER_VALUES))

    icds_apps = [
        "custom.icds",
        "custom.icds.data_management",
        "custom.icds_reports"
    ]
    for app in icds_apps:
        if app not in INSTALLED_APPS:
            INSTALLED_APPS = (app,) + tuple(INSTALLED_APPS)

    CUSTOM_DOMAIN_SPECIFIC_URL_MODULES.extend([
        'custom.icds_reports.urls',
        'custom.icds.urls',
        'custom.icds.data_management.urls'
    ])
    CUSTOM_DOMAIN_SPECIFIC_URL_MODULES = list(set(CUSTOM_DOMAIN_SPECIFIC_URL_MODULES))
<commit_msg>remove old setting<commit_after>from copy import deepcopy

import settingshelper as helper
from settings import *

USING_CITUS = any(db.get('ROLE') == 'citus_master' for db in DATABASES.values())

# note: the only reason these are prepended to INSTALLED_APPS is because of
# a weird travis issue with kafka. if for any reason this order causes problems
# it can be reverted whenever that's figured out.
# https://github.com/dimagi/commcare-hq/pull/10034#issuecomment-174868270
INSTALLED_APPS = (
    'django_nose',
    'testapps.test_elasticsearch',
    'testapps.test_pillowtop',
) + tuple(INSTALLED_APPS)

if USING_CITUS:
    if 'testapps.citus_master' not in INSTALLED_APPS:
        INSTALLED_APPS = (
            'testapps.citus_master',
            'testapps.citus_worker',
        ) + tuple(INSTALLED_APPS)

    if 'testapps.citus_master.citus_router.CitusDBRouter' not in DATABASE_ROUTERS:
        # this router must go first
        DATABASE_ROUTERS = ['testapps.citus_master.citus_router.CitusDBRouter'] + DATABASE_ROUTERS

TEST_RUNNER = 'django_nose.BasicNoseRunner'
NOSE_ARGS = [
    #'--no-migrations' # trim ~120s from test run with db tests
    #'--with-fixture-bundling',
]
NOSE_PLUGINS = [
    'corehq.tests.nose.HqTestFinderPlugin',
    'corehq.tests.noseplugins.dividedwerun.DividedWeRunPlugin',
    'corehq.tests.noseplugins.djangomigrations.DjangoMigrationsPlugin',
    'corehq.tests.noseplugins.cmdline_params.CmdLineParametersPlugin',
    'corehq.tests.noseplugins.patches.PatchesPlugin',
    'corehq.tests.noseplugins.redislocks.RedisLockTimeoutPlugin',
    'corehq.tests.noseplugins.uniformresult.UniformTestResultPlugin',

    # The following are not enabled by default
    'corehq.tests.noseplugins.logfile.LogFilePlugin',
    'corehq.tests.noseplugins.timing.TimingPlugin',
    'corehq.tests.noseplugins.output.OutputPlugin',

    # Uncomment to debug tests. Plugins have nice hooks for inspecting state
    # before/after each test or context setup/teardown, etc.
    #'corehq.tests.noseplugins.debug.DebugPlugin',
]

# these settings can be overridden with environment variables
for key, value in {
    'NOSE_DB_TEST_CONTEXT': 'corehq.tests.nose.HqdbContext',
    'NOSE_NON_DB_TEST_CONTEXT': 'corehq.tests.nose.ErrorOnDbAccessContext',

    'NOSE_IGNORE_FILES': '^localsettings',

    'NOSE_EXCLUDE_DIRS': ';'.join([
        'scripts',
    ]),
}.items():
    os.environ.setdefault(key, value)
del key, value

if "SKIP_TESTS_REQUIRING_EXTRA_SETUP" not in globals():
    SKIP_TESTS_REQUIRING_EXTRA_SETUP = False

CELERY_TASK_ALWAYS_EAGER = True
# keep a copy of the original PILLOWTOPS setting around in case other tests want it.
_PILLOWTOPS = PILLOWTOPS
PILLOWTOPS = {}

# required by auditcare tests
AUDIT_MODEL_SAVE = ['django.contrib.auth.models.User']
AUDIT_ADMIN_VIEWS = False

PHONE_TIMEZONES_HAVE_BEEN_PROCESSED = True
PHONE_TIMEZONES_SHOULD_BE_PROCESSED = True

ENABLE_PRELOGIN_SITE = True

# override dev_settings
CACHE_REPORTS = True

def _set_logging_levels(levels):
    import logging
    for path, level in levels.items():
        logging.getLogger(path).setLevel(level)
_set_logging_levels({
    # Quiet down noisy loggers. Selective removal can be handy for debugging.
    'alembic': 'WARNING',
    'auditcare': 'INFO',
    'boto3': 'WARNING',
    'botocore': 'INFO',
    'couchdbkit.request': 'INFO',
    'couchdbkit.designer': 'WARNING',
    'datadog': 'WARNING',
    'elasticsearch': 'ERROR',
    'kafka.conn': 'WARNING',
    'kafka.client': 'WARNING',
    'kafka.consumer.kafka': 'WARNING',
    'kafka.metrics': 'WARNING',
    'kafka.protocol.parser': 'WARNING',
    'kafka.producer': 'WARNING',
    'quickcache': 'INFO',
    'requests.packages.urllib3': 'WARNING',
    's3transfer': 'INFO',
    'urllib3': 'WARNING',
})

# use empty LOGGING dict with --debug=nose,nose.plugins to debug test discovery
# TODO empty logging config (and fix revealed deprecation warnings)
LOGGING = {
    'disable_existing_loggers': False,
    'version': 1,
    'loggers': {},
}

# Default custom databases to use the same configuration as the default
# This is so that all devs don't have to run citus locally
if 'icds-ucr' not in DATABASES:
    DATABASES['icds-ucr'] = deepcopy(DATABASES['default'])
    # use a different name otherwise migrations don't get run
    DATABASES['icds-ucr']['NAME'] = 'commcarehq_icds_ucr'
    del DATABASES['icds-ucr']['TEST']['NAME']  # gets set by `helper.assign_test_db_names`

helper.assign_test_db_names(DATABASES)

# See comment under settings.SMS_QUEUE_ENABLED
SMS_QUEUE_ENABLED = False

# use all providers in tests
METRICS_PROVIDERS = [
    'corehq.util.metrics.datadog.DatadogMetrics',
    'corehq.util.metrics.prometheus.PrometheusMetrics',
]

# timeout faster in tests
ES_SEARCH_TIMEOUT = 5

if os.path.exists("custom/icds"):
    STATIC_DATA_SOURCES.extend([
        "custom/icds_reports/ucr/data_sources/*.json",
        "custom/icds_reports/ucr/data_sources/dashboard/*.json"
    ])
    STATIC_DATA_SOURCES = list(set(STATIC_DATA_SOURCES))

    STATIC_UCR_REPORTS.extend([
        "custom/icds_reports/ucr/reports/dashboard/*.j,son",
        "custom/icds_reports/ucr/reports/asr/*.json",
        "custom/icds_reports/ucr/reports/asr/ucr_v2/*.json",
        "custom/icds_reports/ucr/reports/mpr/*.json",
        "custom/icds_reports/ucr/reports/mpr/dashboard/*.json",
        "custom/icds_reports/ucr/reports/ls/*.json",
        "custom/icds_reports/ucr/reports/other/*.json",
    ])

    CUSTOM_UCR_EXPRESSION_LISTS.append("custom.icds_reports.ucr.expressions.CUSTOM_UCR_EXPRESSIONS")
    CUSTOM_UCR_EXPRESSION_LISTS = list(set(CUSTOM_UCR_EXPRESSION_LISTS))

    CUSTOM_UCR_REPORT_FILTERS.append(
        ('village_choice_list', 'custom.icds_reports.ucr.filter_spec.build_village_choice_list_filter_spec'),
    )
    CUSTOM_UCR_REPORT_FILTERS = list(set(CUSTOM_UCR_REPORT_FILTERS))
    CUSTOM_UCR_REPORT_FILTER_VALUES.append(
        ('village_choice_list', 'custom.icds_reports.ucr.filter_value.VillageFilterValue'),
    )
    CUSTOM_UCR_REPORT_FILTER_VALUES = list(set(CUSTOM_UCR_REPORT_FILTER_VALUES))

    icds_apps = [
        "custom.icds",
        "custom.icds.data_management",
        "custom.icds_reports"
    ]
    for app in icds_apps:
        if app not in INSTALLED_APPS:
            INSTALLED_APPS = (app,) + tuple(INSTALLED_APPS)
<|endoftext|>


<commit_before>from settings import *

INSTALLED_APPS += (
    'django_nose',
)

TEST_RUNNER = 'django_nose.NoseTestSuiteRunner'
NOSE_ARGS = [
    #'--with-migrations' # adds ~30s to test run; TODO travis should use it
    #'--with-doctest', # adds 5s to discovery (before tests start); TODO travis should use it
    '--with-fixture-bundling',
    '--logging-clear-handlers',
]
NOSE_PLUGINS = [
    # Disable migrations by default. Use --with-migrations to enable them.
    'corehq.util.nose.DjangoMigrationsPlugin',
    'corehq.util.nose.OmitDjangoInitModuleTestsPlugin',
]

# these settings can be overridden with environment variables
for key, value in {
    'NOSE_DB_TEST_CONTEXT': 'corehq.util.nose.HqdbContext',
    'NOSE_NON_DB_TEST_CONTEXT': 'corehq.util.nose.ErrorOnDbAccessContext',

    # ignore record_deploy_success.py because datadog may not be installed
    # (only matters when running --with-doctests)
    'NOSE_IGNORE_FILES': '^(localsettings|record_deploy_success\.py)',

    'NOSE_EXCLUDE_DIRS': ';'.join([
        'corehq/apps/cloudcare/tests/selenium',
        'corehq/apps/reports/tests/selenium',
        'scripts',
        'testapps',

        # excludes for --with-doctest
        # these cause gevent.threading to be imported, which causes this error:
        # DatabaseError: DatabaseWrapper objects created in a thread can only
        # be used in that same thread. ...
        'corehq/apps/hqcase/management/commands',
        'corehq/preindex/management/commands',
        'deployment/gunicorn',
    ]),
}.items():
    os.environ.setdefault(key, value)
del key, value

# HqTestSuiteRunner settings
INSTALLED_APPS = INSTALLED_APPS + list(TEST_APPS)
CELERY_ALWAYS_EAGER = True
PILLOWTOPS = {}

# required by auditcare tests
AUDIT_MODEL_SAVE = ['django.contrib.auth.models.User']
AUDIT_ADMIN_VIEWS=False

PHONE_TIMEZONES_HAVE_BEEN_PROCESSED = True
PHONE_TIMEZONES_SHOULD_BE_PROCESSED = True

ENABLE_PRELOGIN_SITE = True

def _set_couchdb_test_settings():
    import settingshelper

    def get_test_db_name(dbname):
        return "%s_test" % dbname

    global COUCH_DATABASE_NAME, EXTRA_COUCHDB_DATABASES

    COUCH_DATABASE_NAME = get_test_db_name(COUCH_DATABASE_NAME)
    globals().update(settingshelper.get_dynamic_db_settings(
        COUCH_SERVER_ROOT,
        COUCH_USERNAME,
        COUCH_PASSWORD,
        COUCH_DATABASE_NAME,
    ))

    EXTRA_COUCHDB_DATABASES = {
        db_name: get_test_db_name(url)
        for db_name, url in EXTRA_COUCHDB_DATABASES.items()
    }

_set_couchdb_test_settings()

def _clean_up_logging_output():
    import logging
    logging.getLogger('raven').setLevel('WARNING')

    # make all loggers propagate to prevent
    # "No handlers could be found for logger ..."
    # (a side effect of --logging-clear-handlers)
    for item in LOGGING["loggers"].values():
        if not item.get("propagate", True):
            item["propagate"] = True

_clean_up_logging_output()
<commit_msg>Style cleanup [ci skip]<commit_after>from settings import *

INSTALLED_APPS += (
    'django_nose',
)

TEST_RUNNER = 'django_nose.NoseTestSuiteRunner'
NOSE_ARGS = [
    #'--with-migrations' # adds ~30s to test run; TODO travis should use it
    #'--with-doctest', # adds 5s to discovery (before tests start); TODO travis should use it
    '--with-fixture-bundling',
    '--logging-clear-handlers',
]
NOSE_PLUGINS = [
    # Disable migrations by default. Use --with-migrations to enable them.
    'corehq.util.nose.DjangoMigrationsPlugin',
    'corehq.util.nose.OmitDjangoInitModuleTestsPlugin',
]

# these settings can be overridden with environment variables
for key, value in {
    'NOSE_DB_TEST_CONTEXT': 'corehq.util.nose.HqdbContext',
    'NOSE_NON_DB_TEST_CONTEXT': 'corehq.util.nose.ErrorOnDbAccessContext',

    # ignore record_deploy_success.py because datadog may not be installed
    # (only matters when running --with-doctests)
    'NOSE_IGNORE_FILES': '^(localsettings|record_deploy_success\.py)',

    'NOSE_EXCLUDE_DIRS': ';'.join([
        'corehq/apps/cloudcare/tests/selenium',
        'corehq/apps/reports/tests/selenium',
        'scripts',
        'testapps',

        # excludes for --with-doctest
        # these cause gevent.threading to be imported, which causes this error:
        # DatabaseError: DatabaseWrapper objects created in a thread can only
        # be used in that same thread. ...
        'corehq/apps/hqcase/management/commands',
        'corehq/preindex/management/commands',
        'deployment/gunicorn',
    ]),
}.items():
    os.environ.setdefault(key, value)
del key, value

# HqTestSuiteRunner settings
INSTALLED_APPS = INSTALLED_APPS + list(TEST_APPS)
CELERY_ALWAYS_EAGER = True
PILLOWTOPS = {}

# required by auditcare tests
AUDIT_MODEL_SAVE = ['django.contrib.auth.models.User']
AUDIT_ADMIN_VIEWS = False

PHONE_TIMEZONES_HAVE_BEEN_PROCESSED = True
PHONE_TIMEZONES_SHOULD_BE_PROCESSED = True

ENABLE_PRELOGIN_SITE = True

def _set_couchdb_test_settings():
    import settingshelper

    def get_test_db_name(dbname):
        return "%s_test" % dbname

    global COUCH_DATABASE_NAME, EXTRA_COUCHDB_DATABASES

    COUCH_DATABASE_NAME = get_test_db_name(COUCH_DATABASE_NAME)
    globals().update(settingshelper.get_dynamic_db_settings(
        COUCH_SERVER_ROOT,
        COUCH_USERNAME,
        COUCH_PASSWORD,
        COUCH_DATABASE_NAME,
    ))

    EXTRA_COUCHDB_DATABASES = {
        db_name: get_test_db_name(url)
        for db_name, url in EXTRA_COUCHDB_DATABASES.items()
    }

_set_couchdb_test_settings()

def _clean_up_logging_output():
    import logging
    logging.getLogger('raven').setLevel('WARNING')

    # make all loggers propagate to prevent
    # "No handlers could be found for logger ..."
    # (a side effect of --logging-clear-handlers)
    for item in LOGGING["loggers"].values():
        if not item.get("propagate", True):
            item["propagate"] = True

_clean_up_logging_output()
<|endoftext|>


<commit_before>from ts3.TS3Connection import TS3QueryException

__author__ = 'Fabian'
import afkmover
import ClientInfo
import ts3.Events as Events
from importlib import reload
import threading
import logging
import Bot
import re
import Quotes

logger = logging.getLogger("bot")

def send_msg(ts3conn, msg):
    ts3conn.sendtextmessage(targetmode=2, target=1, msg=msg)

class CommandHandler:
    def __init__(self, ts3conn):
        self.ts3conn = ts3conn
        self.afkStopper = threading.Event()
        self.afkMover = None
        self.logger = logging.getLogger("textMsg")
        self.logger.setLevel(logging.INFO)
        file_handler = logging.FileHandler("msg.log", mode='a+')
        formatter = logging.Formatter('MSG Logger %(asctime)s %(message)s')
        file_handler.setFormatter(formatter)
        self.logger.addHandler(file_handler)
        self.logger.propagate = 0

    def start_afkmover(self):
        if self.afkMover is None:
            self.afkMover = afkmover.AfkMover(self.afkStopper, self.ts3conn)
            self.afkStopper.clear()
            self.afkMover.start()
            logger.info("Started afkmover!")

    def handle_command(self, msg, sender=0):
        logger.debug("Handling message " + msg)
        if "!reload" in msg:
            reload(afkmover)
            reload(ClientInfo)
            reload(Events)
        elif msg.startswith("!multimove "):
            self.multi_move(msg, sender)
        elif "!hello" in msg:
            Bot.send_msg_to_client(self.ts3conn, sender, "Hello World")
        elif '!stop' == msg:
            self.afkStopper.set()
            self.afkMover.join()
            self.ts3conn.quit()
            logger.warning("Bot was quit!")
        elif '!restart' == msg:
            self.afkStopper.set()
            self.afkMover = None
            self.ts3conn.quit()
            logger.warning("Bot was quit!")
            import main
            main.restart_program()
        elif '!client' == msg:
            logger.info("client command issued")
            Bot.send_msg_to_client(self.ts3conn, sender, "Client!")
        elif '!afkmove' in msg:
            self.start_afkmover()
        elif "!afkstop" in msg:
            logger.info("Afkstop command received")
            self.afkStopper.set()
            self.afkMover = None
            logger.warning("Stopped afkmover")
        elif re.match('!addquote', msg, re.I):
            if len(msg) > len("!addQuote "):
                Quotes.add(msg[len("!addQuote "):])
                Bot.send_msg_to_client(self.ts3conn, sender, "Quote '" + msg[len("!addQuote "):] + "' was added.")
        else:
            Bot.send_msg_to_client(self.ts3conn, sender, "I cannot interpret your command. I am very sorry. :(")
            logger.info("Unknown command " + msg + " received!")

    def multi_move(self, msg, sender=0):
        channels = msg[len("!multimove "):].split()
        source_name = ""
        dest_name = ""
        source = None
        dest = None
        if len(channels) < 2:
            if sender != 0:
                Bot.send_msg_to_client(self.ts3conn, sender, "Usage: multimove source destination")
                return
        elif len(channels) > 2:
            channel_name_list = self.ts3conn.channel_name_list()
            for channel_name in channel_name_list:
                if msg[len("!multimove "):].startswith(channel_name):
                    source_name = channel_name
                    dest_name = msg[len("!multimove ") + len(source_name)+1:]
        else:
            source_name = channels[0]
            dest_name = channels[1]
        if source_name == "":
            Bot.send_msg_to_client(self.ts3conn, sender, "Source channel not found")
            return
        if dest_name == "":
            Bot.send_msg_to_client(self.ts3conn, sender, "Destination channel not found")
            return
        try:
            channel_candidates = self.ts3conn.channelfind_by_name(source_name)
            if len(channel_candidates) > 0:
                source = channel_candidates[0].get("cid", '-1')
            if source is None or source == "-1":
                Bot.send_msg_to_client(self.ts3conn, sender, "Source channel not found")
                return
        except TS3QueryException:
            Bot.send_msg_to_client(self.ts3conn, sender, "Source channel not found")
        try:
            channel_candidates = self.ts3conn.channelfind_by_name(dest_name)
            if len(channel_candidates) > 0:
                dest = channel_candidates[0].get("cid", '-1')
            if dest is None or dest == "-1":
                Bot.send_msg_to_client(self.ts3conn, sender, "Destination channel not found")
                return
        except TS3QueryException:
            Bot.send_msg_to_client(self.ts3conn, sender, "Destination channel not found")
        try:
            clientlist = self.ts3conn.clientlist()
            self.logger.info(str(clientlist))
            self.logger.info("Moving all from " + source + " to " + dest)
            clientlist = [client for client in clientlist if client.get("cid", '-1') == source]
            for client in clientlist:
                clid = client.get("clid", '-1')
                self.logger.info("Found client in channel: " + client.get("client_nickname", "") + " id = " + clid)
                self.ts3conn.clientmove(int(dest), int(clid))
        except TS3QueryException as e:
            Bot.send_msg_to_client(self.ts3conn, sender, "Error moving clients: id = " +
                                   str(e.id) + e.message)

    def inform(self, event):
        if type(event) is Events.TextMessageEvent:
            if event.targetmode == "Private":
                if event.invoker_id != int(self.ts3conn.whoami()["client_id"]):  # Don't talk to yourself ...
                    ci = ClientInfo.ClientInfo(event.invoker_id, self.ts3conn)
                    self.logger.info("Message: " + event.message + " from: " + ci.name)
                    if ci.is_in_servergroups("Server Admin"):
                        self.handle_command(event.message, sender=event.invoker_id)
                    else:
                        Bot.send_msg_to_client(self.ts3conn, event.invoker_id, "Sorry, but I will only talk to admins!")
<commit_msg>Allow moderators to send textcommands<commit_after>from ts3.TS3Connection import TS3QueryException

__author__ = 'Fabian'
import afkmover
import ClientInfo
import ts3.Events as Events
from importlib import reload
import threading
import logging
import Bot
import re
import Quotes

logger = logging.getLogger("bot")

def send_msg(ts3conn, msg):
    ts3conn.sendtextmessage(targetmode=2, target=1, msg=msg)

class CommandHandler:
    def __init__(self, ts3conn):
        self.ts3conn = ts3conn
        self.afkStopper = threading.Event()
        self.afkMover = None
        self.logger = logging.getLogger("textMsg")
        self.logger.setLevel(logging.INFO)
        file_handler = logging.FileHandler("msg.log", mode='a+')
        formatter = logging.Formatter('MSG Logger %(asctime)s %(message)s')
        file_handler.setFormatter(formatter)
        self.logger.addHandler(file_handler)
        self.logger.propagate = 0

    def start_afkmover(self):
        if self.afkMover is None:
            self.afkMover = afkmover.AfkMover(self.afkStopper, self.ts3conn)
            self.afkStopper.clear()
            self.afkMover.start()
            logger.info("Started afkmover!")

    def handle_command(self, msg, sender=0):
        logger.debug("Handling message " + msg)
        if "!reload" in msg:
            reload(afkmover)
            reload(ClientInfo)
            reload(Events)
        elif msg.startswith("!multimove "):
            self.multi_move(msg, sender)
        elif "!hello" in msg:
            Bot.send_msg_to_client(self.ts3conn, sender, "Hello World")
        elif '!stop' == msg:
            self.afkStopper.set()
            self.afkMover.join()
            self.ts3conn.quit()
            logger.warning("Bot was quit!")
        elif '!restart' == msg:
            self.afkStopper.set()
            self.afkMover = None
            self.ts3conn.quit()
            logger.warning("Bot was quit!")
            import main
            main.restart_program()
        elif '!client' == msg:
            logger.info("client command issued")
            Bot.send_msg_to_client(self.ts3conn, sender, "Client!")
        elif '!afkmove' in msg:
            self.start_afkmover()
        elif "!afkstop" in msg:
            logger.info("Afkstop command received")
            self.afkStopper.set()
            self.afkMover = None
            logger.warning("Stopped afkmover")
        elif re.match('!addquote', msg, re.I):
            if len(msg) > len("!addQuote "):
                Quotes.add(msg[len("!addQuote "):])
                Bot.send_msg_to_client(self.ts3conn, sender, "Quote '" + msg[len("!addQuote "):] + "' was added.")
        else:
            Bot.send_msg_to_client(self.ts3conn, sender, "I cannot interpret your command. I am very sorry. :(")
            logger.info("Unknown command " + msg + " received!")

    def multi_move(self, msg, sender=0):
        channels = msg[len("!multimove "):].split()
        source_name = ""
        dest_name = ""
        source = None
        dest = None
        if len(channels) < 2:
            if sender != 0:
                Bot.send_msg_to_client(self.ts3conn, sender, "Usage: multimove source destination")
                return
        elif len(channels) > 2:
            channel_name_list = self.ts3conn.channel_name_list()
            for channel_name in channel_name_list:
                if msg[len("!multimove "):].startswith(channel_name):
                    source_name = channel_name
                    dest_name = msg[len("!multimove ") + len(source_name)+1:]
        else:
            source_name = channels[0]
            dest_name = channels[1]
        if source_name == "":
            Bot.send_msg_to_client(self.ts3conn, sender, "Source channel not found")
            return
        if dest_name == "":
            Bot.send_msg_to_client(self.ts3conn, sender, "Destination channel not found")
            return
        try:
            channel_candidates = self.ts3conn.channelfind_by_name(source_name)
            if len(channel_candidates) > 0:
                source = channel_candidates[0].get("cid", '-1')
            if source is None or source == "-1":
                Bot.send_msg_to_client(self.ts3conn, sender, "Source channel not found")
                return
        except TS3QueryException:
            Bot.send_msg_to_client(self.ts3conn, sender, "Source channel not found")
        try:
            channel_candidates = self.ts3conn.channelfind_by_name(dest_name)
            if len(channel_candidates) > 0:
                dest = channel_candidates[0].get("cid", '-1')
            if dest is None or dest == "-1":
                Bot.send_msg_to_client(self.ts3conn, sender, "Destination channel not found")
                return
        except TS3QueryException:
            Bot.send_msg_to_client(self.ts3conn, sender, "Destination channel not found")
        try:
            clientlist = self.ts3conn.clientlist()
            self.logger.info(str(clientlist))
            self.logger.info("Moving all from " + source + " to " + dest)
            clientlist = [client for client in clientlist if client.get("cid", '-1') == source]
            for client in clientlist:
                clid = client.get("clid", '-1')
                self.logger.info("Found client in channel: " + client.get("client_nickname", "") + " id = " + clid)
                self.ts3conn.clientmove(int(dest), int(clid))
        except TS3QueryException as e:
            Bot.send_msg_to_client(self.ts3conn, sender, "Error moving clients: id = " +
                                   str(e.id) + e.message)

    def inform(self, event):
        if type(event) is Events.TextMessageEvent:
            if event.targetmode == "Private":
                if event.invoker_id != int(self.ts3conn.whoami()["client_id"]):  # Don't talk to yourself ...
                    ci = ClientInfo.ClientInfo(event.invoker_id, self.ts3conn)
                    self.logger.info("Message: " + event.message + " from: " + ci.name)
                    if ci.is_in_servergroups("Server Admin") or ci.is_in_servergroups("Moderator"):
                        self.handle_command(event.message, sender=event.invoker_id)
                    else:
                        Bot.send_msg_to_client(self.ts3conn, event.invoker_id, "Sorry, but I will only talk to admins!")
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#######################################################################
# Name: export.py
# Purpose: Export of textX based models and metamodels to dot file
# Author: Igor R. Dejanovic <igor DOT dejanovic AT gmail DOT com>
# Copyright: (c) 2014 Igor R. Dejanovic <igor DOT dejanovic AT gmail DOT com>
# License: MIT License
#######################################################################
from __future__ import unicode_literals
import codecs
import sys
if sys.version < '3':
    text = unicode
else:
    text = str

from .const import MULT_ZEROORMORE, MULT_ONEORMORE, MULT_ONE, RULE_MATCH
from .textx import BASE_TYPE_NAMES, PRIMITIVE_PYTHON_TYPES

HEADER = '''
    digraph xtext {
    fontname = "Bitstream Vera Sans"
    fontsize = 8
    node[
        shape=record,
        style=filled,
        fillcolor=aliceblue
    ]
    edge[dir=black,arrowtail=empty]

'''

def metamodel_export(metamodel, file_name):

    with codecs.open(file_name, 'w', encoding="utf-8") as f:
        f.write(HEADER)

        for cls in metamodel:
            name = cls.__name__
            attrs = ""
            if cls._tx_type == RULE_MATCH:
                attrs = cls._match_str.replace("\\", "\\\\")\
                                      .replace("|", "\\|")\
                                      .replace('"', '\\"')\
                                      .replace('{', '\\{')\
                                      .replace('}', '\\}')
            else:
                for attr in cls._tx_attrs.values():
                    arrowtail = "arrowtail=diamond, dir=both, " \
                        if attr.cont else ""
                    mult_list = attr.mult in [MULT_ZEROORMORE, MULT_ONEORMORE]
                    required = "+" if attr.mult in \
                        [MULT_ONE, MULT_ONEORMORE] else ""
                    attr_type = "list[{}]".format(attr.cls.__name__) \
                        if mult_list else attr.cls.__name__
                    if attr.ref:
                        # If attribute is a reference
                        mult = attr.mult if not attr.mult == MULT_ONE else ""
                        f.write('{} -> {}[{}headlabel="{} {}"]\n'
                                .format(id(cls), id(attr.cls), arrowtail,
                                        attr.name, mult))
                    else:
                        # If it is plain type
                        attrs += "{}{}:{}\\l".format(required,
                                                     attr.name, attr_type)

            f.write('{}[ label="{{{}|{}}}"]\n'.format(
                    id(cls), name, attrs))

            for inherited_by in cls._tx_inh_by:
                f.write('{} -> {} [dir=back]\n'
                        .format(id(cls), id(inherited_by)))

            f.write("\n")

        f.write("\n}\n")

def model_export(model, file_name):

    processed_set = set()

    with codecs.open(file_name, 'w', encoding="utf-8") as f:
        f.write(HEADER)

        def _export(obj):

            if obj is None or obj in processed_set or type(obj) \
                    in PRIMITIVE_PYTHON_TYPES:
                return

            processed_set.add(obj)

            attrs = ""
            obj_cls = obj.__class__
            name = ""
            for attr_name, attr in obj_cls._tx_attrs.items():

                attr_value = getattr(obj, attr_name)

                endmark = 'arrowtail=diamond dir=both' if attr.cont else ""
                required = "+" if attr.mult in \
                    [MULT_ONE, MULT_ONEORMORE] else ""

                if attr.mult in [MULT_ONEORMORE, MULT_ZEROORMORE]:
                    if not attr.ref:
                        attrs += "{}{}:list=[".format(required, attr_name)
                        attrs += ",".join([str(x) for x in attr_value])
                        attrs += "]\\l"
                    else:
                        for idx, list_obj in enumerate(attr_value):
                            if list_obj is not None:
                                if type(list_obj) in PRIMITIVE_PYTHON_TYPES:
                                    f.write(
                                        '{} -> "{}:{}" [label="{}:{}" {}]\n'
                                        .format(id(obj), list_obj,
                                                type(list_obj).__name__,
                                                attr_name, idx, endmark))
                                else:
                                    f.write('{} -> {} [label="{}:{}" {}]\n'
                                            .format(id(obj), id(list_obj),
                                                    attr_name, idx, endmark))
                                    _export(list_obj)
                else:
                    # Plain attributes
                    if not attr.ref:
                        if attr_name == 'name':
                            name = attr_value
                        else:
                            if attr.cls.__name__ in BASE_TYPE_NAMES \
                                    or attr.cls._tx_type == RULE_MATCH:
                                if type(attr_value) in [str, text]:
                                    attr_value = \
                                        attr_value.replace('\n', r'\n')\
                                        .replace('"', '\\"')
                                attrs += "{}{}:{}={}\\l".format(
                                    required, attr_name, type(attr_value)
                                    .__name__, attr_value)
                            else:
                                attrs += "{}{}:{}\\l".format(
                                    required, attr_name, type(attr_value)
                                    .__name__)
                    else:
                        # Object references
                        if attr_value is not None:
                            f.write('{} -> {} [label="{}" {}]\n'.format(
                                id(obj), id(attr_value),
                                attr_name, endmark))
                            _export(attr_value)

            name = "{}:{}".format(name, obj_cls.__name__)

            f.write('{}[label="{{{}|{}}}"]\n'.format(id(obj), name, attrs))

        _export(model)

        f.write('\n}\n')
<commit_msg>fix #19 Proper model rendering for mixed rules (abstract/match).<commit_after># -*- coding: utf-8 -*-
#######################################################################
# Name: export.py
# Purpose: Export of textX based models and metamodels to dot file
# Author: Igor R. Dejanovic <igor DOT dejanovic AT gmail DOT com>
# Copyright: (c) 2014-2016 Igor R. Dejanovic <igor DOT dejanovic AT gmail DOT com>
# License: MIT License
#######################################################################
from __future__ import unicode_literals
from .const import MULT_ZEROORMORE, MULT_ONEORMORE, MULT_ONE, RULE_MATCH
from .textx import PRIMITIVE_PYTHON_TYPES
import codecs
import sys
if sys.version < '3':
    text = unicode
else:
    text = str

HEADER = '''
    digraph xtext {
    fontname = "Bitstream Vera Sans"
    fontsize = 8
    node[
        shape=record,
        style=filled,
        fillcolor=aliceblue
    ]
    edge[dir=black,arrowtail=empty]

'''

def dot_escape(s):
    return s.replace('\n', r'\n').replace('"', r'\"').replace(':', r'\:')

def metamodel_export(metamodel, file_name):

    with codecs.open(file_name, 'w', encoding="utf-8") as f:
        f.write(HEADER)

        for cls in metamodel:
            name = cls.__name__
            attrs = ""
            if cls._tx_type == RULE_MATCH:
                attrs = cls._match_str.replace("\\", "\\\\")\
                                      .replace("|", "\\|")\
                                      .replace('"', '\\"')\
                                      .replace('{', '\\{')\
                                      .replace('}', '\\}')
            else:
                for attr in cls._tx_attrs.values():
                    arrowtail = "arrowtail=diamond, dir=both, " \
                        if attr.cont else ""
                    mult_list = attr.mult in [MULT_ZEROORMORE, MULT_ONEORMORE]
                    required = "+" if attr.mult in \
                        [MULT_ONE, MULT_ONEORMORE] else ""
                    attr_type = "list[{}]".format(attr.cls.__name__) \
                        if mult_list else attr.cls.__name__
                    if attr.ref:
                        # If attribute is a reference
                        mult = attr.mult if not attr.mult == MULT_ONE else ""
                        f.write('{} -> {}[{}headlabel="{} {}"]\n'
                                .format(id(cls), id(attr.cls), arrowtail,
                                        attr.name, mult))
                    else:
                        # If it is plain type
                        attrs += "{}{}:{}\\l".format(required,
                                                     attr.name, attr_type)

            f.write('{}[ label="{{{}|{}}}"]\n'.format(
                    id(cls), name, attrs))

            for inherited_by in cls._tx_inh_by:
                f.write('{} -> {} [dir=back]\n'
                        .format(id(cls), id(inherited_by)))

            f.write("\n")

        f.write("\n}\n")

def model_export(model, file_name):

    processed_set = set()

    with codecs.open(file_name, 'w', encoding="utf-8") as f:
        f.write(HEADER)

        def _export(obj):

            if obj is None or obj in processed_set or type(obj) \
                    in PRIMITIVE_PYTHON_TYPES:
                return

            processed_set.add(obj)

            attrs = ""
            obj_cls = obj.__class__
            name = ""
            for attr_name, attr in obj_cls._tx_attrs.items():

                attr_value = getattr(obj, attr_name)

                endmark = 'arrowtail=diamond dir=both' if attr.cont else ""
                required = "+" if attr.mult in \
                    [MULT_ONE, MULT_ONEORMORE] else ""

                if attr.mult in [MULT_ONEORMORE, MULT_ZEROORMORE]:
                    if all([type(x) in PRIMITIVE_PYTHON_TYPES for x in attr_value]):
                        attrs += "{}{}:list=[".format(required, attr_name)
                        attrs += ",".join([str(x) for x in attr_value])
                        attrs += "]\\l"
                    else:
                        for idx, list_obj in enumerate(attr_value):
                            if list_obj is not None:
                                if type(list_obj) in PRIMITIVE_PYTHON_TYPES:
                                    f.write(
                                        '{} -> "{}:{}" [label="{}:{}" {}]\n'
                                        .format(id(obj), list_obj,
                                                type(list_obj).__name__,
                                                attr_name, idx, endmark))
                                else:
                                    f.write('{} -> {} [label="{}:{}" {}]\n'
                                            .format(id(obj), id(list_obj),
                                                    attr_name, idx, endmark))
                                    _export(list_obj)
                else:

                    # Plain attributes
                    if type(attr_value) in [str, text]:
                        attr_value = dot_escape(attr_value)

                    if type(attr_value) in PRIMITIVE_PYTHON_TYPES:
                        if attr_name == 'name':
                            name = attr_value
                        else:
                            attrs += "{}{}:{}={}\\l".format(
                                required, attr_name, type(attr_value)
                                .__name__, attr_value)
                    else:
                        # Object references
                        if attr_value is not None:
                            f.write('{} -> {} [label="{}" {}]\n'.format(
                                id(obj), id(attr_value),
                                attr_name, endmark))
                            _export(attr_value)

            name = "{}:{}".format(name, obj_cls.__name__)

            f.write('{}[label="{{{}|{}}}"]\n'.format(id(obj), name, attrs))

        _export(model)

        f.write('\n}\n')
<|endoftext|>


<commit_before># 4/18/2014
# Charles O. Goddard

from __future__ import print_function

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from bs4 import BeautifulSoup

import os
import sys
import pickle

def striptext(text):
    words = []
    soup = BeautifulSoup(text, 'html5lib')
    elems = [elem for elem in soup]
    while elems:
        elem = elems.pop(0)
        if elem.string != None:
            words.append(elem.string)
        if hasattr(elem, 'contents'):
            elems += elem.contents
    res = '\n'.join(words)
    while '\n\n' in res:
        res = res.replace('\n\n', '\n')
    return res

def read_doc(fn):
    with open(fn, 'r') as f:
        text = f.read()
        return striptext(text)

def search(text, documents=None):
    if documents is None:
        documents = pickle.load(open('documents.pickle', 'rb'))

    #print('Query:', text)

    docnames = ["query"] + [d[0] for d in documents]
    documents = [text] + [d[1] for d in documents]

    vectorizer = TfidfVectorizer(stop_words='english')
    tfidf = vectorizer.fit_transform(documents)
    num_samples, num_features = tfidf.shape
    print('{0} samples, {1} features'.format(num_samples, num_features))

    search_vector = tfidf[0].A[0]
    match_scores = [(sum(search_vector * tfidf[i].A[0]), i) for i in range(1, num_samples)]
    match_scores.sort(reverse=True)

    return [(score, docnames[i]) for (score, i) in match_scores[:10]]

if __name__ == '__main__':
    if len(sys.argv) > 1:
        if sys.argv[1] == 'reparse':
            documents = []
            docnames = os.listdir('scraped')
            for docname in docnames:
                documents.append((docname, read_doc('scraped/'+docname)))
            print('Read in all documents')
            pickle.dump(documents, open('documents.pickle', 'wb'), -1)
            sys.exit(0)
    query = "New scorecard ranks states on their 'fertility friendliness'. For National Infertility Week, RESOLVE: The National Infertility Association has released its annual Fertility Scorecard - a map ranking each state by how easy it is for citizens to gain access to fertility support resources and treatments in that area."
    match_scores = search(query, documents)
    for score, i in match_scores[:10]:
        print('{0} -> {1}'.format(score, docnames[i]))
<commit_msg>Fix striptext() to remove comments and produce less whitespace<commit_after># 4/18/2014
# Charles O. Goddard

from __future__ import print_function

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from bs4 import BeautifulSoup, Comment

import os
import sys
import pickle

def striptext(text):
    words = []
    soup = BeautifulSoup(text, 'html5lib')
    elems = [elem for elem in soup.html]
    while elems:
        elem = elems.pop(0)
        if not isinstance(elem, Comment) and elem.string is not None:
            words.append(elem.string)
        if hasattr(elem, 'contents'):
            elems = elem.contents + elems
    res = '\n'.join(words)
    res = '\n'.join(w.strip() for w in res.split())
    while '\n\n' in res:
        res = res.replace('\n\n', '\n')
    return res

def read_doc(fn):
    with open(fn, 'r') as f:
        text = f.read().decode('utf-8')
        return striptext(text)

def search(text, documents=None):
    if documents is None:
        documents = pickle.load(open('documents.pickle', 'rb'))

    #print('Query:', text)

    docnames = ["query"] + [d[0] for d in documents]
    documents = [text] + [d[1] for d in documents]

    vectorizer = TfidfVectorizer(stop_words='english')
    tfidf = vectorizer.fit_transform(documents)
    num_samples, num_features = tfidf.shape
    print('{0} samples, {1} features'.format(num_samples, num_features))

    search_vector = tfidf[0].A[0]
    match_scores = [(sum(search_vector * tfidf[i].A[0]), i) for i in range(1, num_samples)]
    match_scores.sort(reverse=True)

    return [(score, docnames[i]) for (score, i) in match_scores[:10]]

if __name__ == '__main__':
    if len(sys.argv) > 1:
        if sys.argv[1] == 'reparse':
            documents = []
            docnames = os.listdir('scraped')
            for docname in docnames:
                documents.append((docname, read_doc('scraped/'+docname)))
            print('Read in all documents')
            pickle.dump(documents, open('documents.pickle', 'wb'), -1)
            sys.exit(0)
    query = "New scorecard ranks states on their 'fertility friendliness'. For National Infertility Week, RESOLVE: The National Infertility Association has released its annual Fertility Scorecard - a map ranking each state by how easy it is for citizens to gain access to fertility support resources and treatments in that area."
    match_scores = search(query, documents)
    for score, i in match_scores[:10]:
        print('{0} -> {1}'.format(score, docnames[i]))
<|endoftext|>


<commit_before>'''
@author: shylent
'''
from tftp.datagram import (ACKDatagram, ERRORDatagram, ERR_TID_UNKNOWN,
    TFTPDatagramFactory, split_opcode, OP_DATA, OP_ERROR, ERR_ILLEGAL_OP,
    ERR_DISK_FULL, OP_ACK, DATADatagram, ERR_NOT_DEFINED)
from twisted.internet import reactor
from twisted.internet.defer import maybeDeferred
from twisted.internet.protocol import DatagramProtocol
from twisted.python import log
from tftp.util import SequentialCall

class WriteSession(DatagramProtocol):

    block_size = 512
    timeout = 10

    def __init__(self, remote, writer, _clock=None):
        self.writer = writer
        self.remote = remote
        self.blocknum = 0
        self.completed = False
        self.timeout_watchdog = None
        if _clock is None:
            self._clock = reactor
        else:
            self._clock = _clock

    def _resetWatchdog(self, timeout):
        if self.timeout_watchdog is not None:
            self.timeout_watchdog.reset(timeout)
        else:
            self.timeout_watchdog = self._clock.callLater(timeout, self.timedOut)

    def cancel(self):
        if self.timeout_watchdog is not None:
            self.timeout_watchdog.cancel()
        self.writer.cancel()
        self.transport.stopListening()

    def startProtocol(self):
        addr = self.transport.getHost()
        log.msg("Write session started on %s, remote: %s" % (addr, self.remote))
        self.transport.connect(*self.remote)
        self._resetWatchdog(self.timeout)

    def datagramReceived(self, datagram, addr):
        if self.remote[1] != addr[1]:
            self.transport.write(ERRORDatagram.from_code(ERR_TID_UNKNOWN).to_wire())
            return # Does not belong to this transfer
        datagram = TFTPDatagramFactory(*split_opcode(datagram))
        log.msg("Datagram received from %s: %s" % (addr, datagram))
        if datagram.opcode == OP_DATA:
            return self.tftp_DATA(datagram)
        elif datagram.opcode == OP_ERROR:
            log.msg("Got error: " % datagram)
            self.cancel()

    def tftp_DATA(self, datagram):
        next_blocknum = self.blocknum + 1
        if datagram.blocknum < next_blocknum:
            self.transport.write(ACKDatagram(datagram.blocknum).to_wire())
        elif datagram.blocknum == next_blocknum:
            if self.completed:
                self.transport.write(ERRORDatagram.from_code(
                    ERR_ILLEGAL_OP, "Transfer already finished").to_wire())
            else:
                return self.nextBlock(datagram)
        else:
            self.transport.write(ERRORDatagram.from_code(
                ERR_ILLEGAL_OP, "Block number mismatch").to_wire())

    def nextBlock(self, datagram):
        self._resetWatchdog(self.timeout)
        self.blocknum += 1
        d = maybeDeferred(self.writer.write, datagram.data)
        d.addCallbacks(callback=self.blockWriteSuccess, callbackArgs=[datagram, ],
                       errback=self.blockWriteFailure)
        return d

    def blockWriteFailure(self, failure):
        log.err(failure)
        self.transport.write(ERRORDatagram.from_code(ERR_DISK_FULL).to_wire())
        self.cancel()

    def blockWriteSuccess(self, ign, datagram):
        self.transport.write(ACKDatagram(datagram.blocknum).to_wire())
        if len(datagram.data) < self.block_size:
            self.completed = True
            self.writer.finish()

    def timedOut(self):
        if not self.completed:
            log.msg("Timed out while waiting for next block")
            self.writer.cancel()
        else:
            log.msg("Timed out after a successful transfer")
        self.transport.stopListening()

class LocalOriginWriteSession(WriteSession):

    def __init__(self, remote, writer, _clock=None):
        self.handshake_timeout_watchdog = None
        WriteSession.__init__(self, remote, writer, _clock)

    def startProtocol(self):
        if self.handshake_timeout_watchdog is not None:
            self.handshake_timeout_watchdog.start()
        WriteSession.startProtocol(self)

    def nextBlock(self, datagram):
        if self.handshake_timeout_watchdog.active():
            self.handshake_timeout_watchdog.cancel()
        return WriteSession.nextBlock(self, datagram)

class RemoteOriginWriteSession(WriteSession):

    def startProtocol(self):
        WriteSession.startProtocol(self)
        self.transport.write(ACKDatagram(self.blocknum).to_wire())

class ReadSession(DatagramProtocol):
    block_size = 512
    timeout = (3, 5, 10)

    def __init__(self, remote, reader, _clock=None):
        self.remote = remote
        self.reader = reader
        self.blocknum = 0
        self.completed = False
        self.timeout_watchdog = None
        if _clock is None:
            self._clock = reactor
        else:
            self._clock = _clock

    def cancel(self):
        self.reader.finish()
        if self.timeout_watchdog is not None:
            self.timeout_watchdog.cancel()
        self.transport.stopListening()

    def startProtocol(self):
        addr = self.transport.getHost()
        log.msg("Read session started on %s, remote: %s" % (addr, self.remote))
        self.transport.connect(*self.remote)

    def datagramReceived(self, datagram, addr):
        if self.remote[1] != addr[1]:
            self.transport.write(ERRORDatagram.from_code(ERR_TID_UNKNOWN).to_wire())
            return
        datagram = TFTPDatagramFactory(*split_opcode(datagram))
        log.msg("Datagram received from %s: %s" % (addr, datagram))
        if datagram.opcode == OP_ACK:
            return self.tftp_ACK(datagram)
        elif datagram.opcode == OP_ERROR:
            log.msg("Got error: " % datagram)
            self.cancel()

    def tftp_ACK(self, datagram):
        if datagram.blocknum < self.blocknum:
            log.msg("Duplicate ACK for blocknum %s" % datagram.blocknum)
        elif datagram.blocknum == self.blocknum:
            if self.timeout_watchdog is not None:
                self.timeout_watchdog.cancel()
            if self.completed:
                log.msg("Final ACK received, transfer successful")
                self.cancel()
            else:
                return self.nextBlock()
        else:
            self.transport.write(ERRORDatagram.from_code(
                ERR_ILLEGAL_OP, "Block number mismatch").to_wire())

    def nextBlock(self):
        self.blocknum += 1
        d = maybeDeferred(self.reader.read, self.block_size)
        d.addCallbacks(callback=self.dataFromReader, errback=self.readFailed)
        return d

    def dataFromReader(self, data):
        if len(data) < self.block_size:
            self.completed = True
        bytes = DATADatagram(self.blocknum, data).to_wire()
        self.timeout_watchdog = SequentialCall.run(self.timeout[:-1],
            callable=self.sendData, callable_args=[bytes, ],
            on_timeout=lambda: self._clock.callLater(self.timeout[-1], self.timedOut),
            run_now=True,
            _clock=self._clock
        )

    def readFailed(self, fail):
        log.err(fail)
        self.transport.write(ERRORDatagram.from_code(ERR_NOT_DEFINED, "Read failed").to_wire())
        self.cancel()

    def timedOut(self):
        log.msg("Session timed out, last wait was %s seconds long" % self.timeout[-1])
        self.cancel()

    def sendData(self, bytes):
        self.transport.write(bytes)

class LocalOriginReadSession(ReadSession):

    def __init__(self, remote, reader, _clock=None):
        self.handshake_timeout_watchdog = None
        ReadSession.__init__(self, remote, reader, _clock)

    def startProtocol(self):
        if self.handshake_timeout_watchdog is not None:
            self.handshake_timeout_watchdog.start()
        ReadSession.startProtocol(self)

    def nextBlock(self):
        if (self.handshake_timeout_watchdog is not None and
                self.handshake_timeout_watchdog.active()):
            self.handshake_timeout_watchdog.cancel()
        return ReadSession.nextBlock(self)

class RemoteOriginReadSession(ReadSession):

    def startProtocol(self):
        ReadSession.startProtocol(self)
        return self.nextBlock()
<commit_msg>Docstrings for sessions<commit_after>'''
@author: shylent
'''
from tftp.datagram import (ACKDatagram, ERRORDatagram, ERR_TID_UNKNOWN,
    TFTPDatagramFactory, split_opcode, OP_DATA, OP_ERROR, ERR_ILLEGAL_OP,
    ERR_DISK_FULL, OP_ACK, DATADatagram, ERR_NOT_DEFINED)
from twisted.internet import reactor
from twisted.internet.defer import maybeDeferred
from twisted.internet.protocol import DatagramProtocol
from twisted.python import log
from tftp.util import SequentialCall

class WriteSession(DatagramProtocol):
    """Represents a transfer, during which we write to a local file. If we are a
    server, this means, that we received a WRQ (write request). If we are a client,
    this means, that we have requested a read from a remote server.

    @cvar block_size: Expected block size. If a data chunk is received and its length
    is less, than C{block_size}, it is assumed that that data chunk is the last in the
    transfer. Default: 512 (as per U{RFC1350<http://tools.ietf.org/html/rfc1350>})
    @type block_size: C{int}.

    @cvar timeout: a number of seconds to wait for the next data chunk. Default: 10
    @type timeout: C{int}

    """

    block_size = 512
    timeout = 10

    def __init__(self, remote, writer, _clock=None):
        self.writer = writer
        self.remote = remote
        self.blocknum = 0
        self.completed = False
        self.timeout_watchdog = None
        if _clock is None:
            self._clock = reactor
        else:
            self._clock = _clock

    def cancel(self):
        """Cancel this session, discard any data, that was collected
        and give up the connector.

        """
        if self.timeout_watchdog is not None:
            self.timeout_watchdog.cancel()
        self.writer.cancel()
        self.transport.stopListening()

    def startProtocol(self):
        """Enter "connected" mode and start the timeout watchdog"""
        addr = self.transport.getHost()
        log.msg("Write session started on %s, remote: %s" % (addr, self.remote))
        self.transport.connect(*self.remote)
        self._resetWatchdog(self.timeout)

    def connectionRefused(self):
        if not self.completed:
            self.writer.cancel()
        self.transport.stopListening()

    def datagramReceived(self, datagram, addr):
        if self.remote[1] != addr[1]:
            self.transport.write(ERRORDatagram.from_code(ERR_TID_UNKNOWN).to_wire())
            return # Does not belong to this transfer
        datagram = TFTPDatagramFactory(*split_opcode(datagram))
        log.msg("Datagram received from %s: %s" % (addr, datagram))
        if datagram.opcode == OP_DATA:
            return self.tftp_DATA(datagram)
        elif datagram.opcode == OP_ERROR:
            log.msg("Got error: " % datagram)
            self.cancel()

    def tftp_DATA(self, datagram):
        """Handle incoming DATA TFTP datagram

        @type datagram: L{DATADatagram}

        """
        next_blocknum = self.blocknum + 1
        if datagram.blocknum < next_blocknum:
            self.transport.write(ACKDatagram(datagram.blocknum).to_wire())
        elif datagram.blocknum == next_blocknum:
            if self.completed:
                self.transport.write(ERRORDatagram.from_code(
                    ERR_ILLEGAL_OP, "Transfer already finished").to_wire())
            else:
                return self.nextBlock(datagram)
        else:
            self.transport.write(ERRORDatagram.from_code(
                ERR_ILLEGAL_OP, "Block number mismatch").to_wire())

    def nextBlock(self, datagram):
        """Handle fresh data, attempt to write it to backend

        @type datagram: L{DATADatagram}

        """
        self._resetWatchdog(self.timeout)
        self.blocknum += 1
        d = maybeDeferred(self.writer.write, datagram.data)
        d.addCallbacks(callback=self.blockWriteSuccess, callbackArgs=[datagram, ],
                       errback=self.blockWriteFailure)
        return d

    def blockWriteSuccess(self, ign, datagram):
        """The write was successful, respond with ACK for current block number

        If this is the last chunk (received data length < block size), the protocol
        will keep running until the end of current timeout period, so we can respond
        to any duplicates.

        @type datagram: L{DATADatagram}

        """
        self.transport.write(ACKDatagram(datagram.blocknum).to_wire())
        if len(datagram.data) < self.block_size:
            self.completed = True
            self.writer.finish()

    def blockWriteFailure(self, failure):
        """Write failed"""
        log.err(failure)
        self.transport.write(ERRORDatagram.from_code(ERR_DISK_FULL).to_wire())
        self.cancel()

    def timedOut(self):
        """Called when the protocol has timed out. Let the backend know, if the
        the transfer was successful.

        """
        if not self.completed:
            log.msg("Timed out while waiting for next block")
            self.writer.cancel()
        else:
            log.msg("Timed out after a successful transfer")
        self.transport.stopListening()

    def _resetWatchdog(self, timeout):
        if self.timeout_watchdog is not None:
            self.timeout_watchdog.reset(timeout)
        else:
            self.timeout_watchdog = self._clock.callLater(timeout, self.timedOut)

class LocalOriginWriteSession(WriteSession):
    """Bootstraps a L{WriteSession}, that was initiated locally, - we've requested
    a read from a remote server

    @ivar handshake_timeout_watchdog: an object, responsible for terminating this
    session, if the initial data chunk has not been received. Set by the owner
    before this protocol's L{startProtocol} is called

    """
    def __init__(self, remote, writer, _clock=None):
        self.handshake_timeout_watchdog = None
        WriteSession.__init__(self, remote, writer, _clock)

    def startProtocol(self):
        """Start the L{handshake_timeout_watchdog} and hand over control to the
        actual L{WriteSession} protocol.

        """
        if self.handshake_timeout_watchdog is not None:
            self.handshake_timeout_watchdog.start()
        WriteSession.startProtocol(self)

    def nextBlock(self, datagram):
        """Initial DATA datagram has been received and the transfer is considered
        to have started successfully. Cancel the L{handshake_timeout_watchdog}

        """
        if self.handshake_timeout_watchdog.active():
            self.handshake_timeout_watchdog.cancel()
        return WriteSession.nextBlock(self, datagram)

class RemoteOriginWriteSession(WriteSession):
    """Bootstraps a L{WriteSession}, that was originated remotely, - we've
    received a WRQ from a client.

    """

    def startProtocol(self):
        """Respond with an initial ACK"""
        WriteSession.startProtocol(self)
        self.transport.write(ACKDatagram(self.blocknum).to_wire())

class ReadSession(DatagramProtocol):
    """Represents a transfer, during which we read from a local file
    (and write to the network). If we are a server, this means, that we've received
    a RRQ (read request). If we are a client, this means that we've requested to
    write to a remote server.

    @cvar block_size: The data will be sent in chunks of this size. If we send
    a chunk with the size < C{block_size}, the transfer will end.
    Default: 512 (as per U{RFC1350<http://tools.ietf.org/html/rfc1350>})
    @type block_size: C{int}

    @cvar timeout: An iterable, that yields timeout values for every subsequent
    unacknowledged DATADatagram, that we've sent. When (if) the iterable is exhausted,
    the transfer is considered failed.
    @type timeout: any iterable

    """
    block_size = 512
    timeout = (3, 5, 10)

    def __init__(self, remote, reader, _clock=None):
        self.remote = remote
        self.reader = reader
        self.blocknum = 0
        self.completed = False
        self.timeout_watchdog = None
        if _clock is None:
            self._clock = reactor
        else:
            self._clock = _clock

    def cancel(self):
        """Tell the reader to give up the resources. Stop the timeout cycle
        and disconnect the transport.

        """
        self.reader.finish()
        if self.timeout_watchdog is not None:
            self.timeout_watchdog.cancel()
        self.transport.stopListening()

    def startProtocol(self):
        """Enter "connected" mode"""
        addr = self.transport.getHost()
        log.msg("Read session started on %s, remote: %s" % (addr, self.remote))
        self.transport.connect(*self.remote)

    def connectionRefused(self):
        self.finish()

    def datagramReceived(self, datagram, addr):
        if self.remote[1] != addr[1]:
            self.transport.write(ERRORDatagram.from_code(ERR_TID_UNKNOWN).to_wire())
            return
        datagram = TFTPDatagramFactory(*split_opcode(datagram))
        log.msg("Datagram received from %s: %s" % (addr, datagram))
        if datagram.opcode == OP_ACK:
            return self.tftp_ACK(datagram)
        elif datagram.opcode == OP_ERROR:
            log.msg("Got error: " % datagram)
            self.cancel()

    def tftp_ACK(self, datagram):
        """Handle the incoming ACK TFTP datagram.

        @type datagram: L{ACKDatagram}

        """
        if datagram.blocknum < self.blocknum:
            log.msg("Duplicate ACK for blocknum %s" % datagram.blocknum)
        elif datagram.blocknum == self.blocknum:
            if self.timeout_watchdog is not None:
                self.timeout_watchdog.cancel()
            if self.completed:
                log.msg("Final ACK received, transfer successful")
                self.cancel()
            else:
                return self.nextBlock()
        else:
            self.transport.write(ERRORDatagram.from_code(
                ERR_ILLEGAL_OP, "Block number mismatch").to_wire())

    def nextBlock(self):
        """ACK datagram for the previous block has been received. Attempt to read
        the next block, that will be sent.

        """
        self.blocknum += 1
        d = maybeDeferred(self.reader.read, self.block_size)
        d.addCallbacks(callback=self.dataFromReader, errback=self.readFailed)
        return d

    def dataFromReader(self, data):
        """Got data from the reader. Send it to the network and start the timeout
        cycle.

        """
        if len(data) < self.block_size:
            self.completed = True
        bytes = DATADatagram(self.blocknum, data).to_wire()
        self.timeout_watchdog = SequentialCall.run(self.timeout[:-1],
            callable=self.sendData, callable_args=[bytes, ],
            on_timeout=lambda: self._clock.callLater(self.timeout[-1], self.timedOut),
            run_now=True,
            _clock=self._clock
        )

    def readFailed(self, fail):
        """The reader reported an error. Notify the remote end and cancel the transfer"""
        log.err(fail)
        self.transport.write(ERRORDatagram.from_code(ERR_NOT_DEFINED, "Read failed").to_wire())
        self.cancel()

    def timedOut(self):
        """Timeout iterable has been exhausted. End the transfer"""
        log.msg("Session timed out, last wait was %s seconds long" % self.timeout[-1])
        self.cancel()

    def sendData(self, bytes):
        self.transport.write(bytes)

class LocalOriginReadSession(ReadSession):
    """Bootstraps a L{ReadSession}, that was originated locally, - we've requested
    a write to a remote server.

    @ivar handshake_timeout_watchdog: an object, responsible for terminating this
    session, if the initial data chunk has not been received. Set by the owner
    before this protocol's L{startProtocol} is called

    """
    def __init__(self, remote, reader, _clock=None):
        self.handshake_timeout_watchdog = None
        ReadSession.__init__(self, remote, reader, _clock)

    def startProtocol(self):
        """Start the L{handshake_timeout_watchdog} and hand over control to
        L{ReadSession}.

        """
        if self.handshake_timeout_watchdog is not None:
            self.handshake_timeout_watchdog.start()
        ReadSession.startProtocol(self)

    def nextBlock(self):
        """An initial ACK datagram has been received, the transfer has started.
        Cancel the L{handshake_timeout_watchdog}

        """
        if (self.handshake_timeout_watchdog is not None and
                self.handshake_timeout_watchdog.active()):
            self.handshake_timeout_watchdog.cancel()
        return ReadSession.nextBlock(self)

class RemoteOriginReadSession(ReadSession):
    """Bootstraps a L{ReadSession}, that was started remotely, - we've received
    a RRQ.

    """

    def startProtocol(self):
        """Initialize the L{ReadSession} and make it send the first data chunk."""
        ReadSession.startProtocol(self)
        return self.nextBlock()
<|endoftext|>


<commit_before><commit_msg>add code<commit_after># get colors into lioshi.tmTheme file
# list colors to replace by colrs-rgb.txt palette<|endoftext|>