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

from argparse import ArgumentParser
from contextlib import contextmanager
import logging
from textwrap import dedent
from subprocess import CalledProcessError
import sys

from jujupy import (
    EnvJujuClient,
    SimpleEnvironment,
    until_timeout,
    )
from deploy_stack import (
    BootstrapManager,
    check_token,
    get_random_string,
    )
from jujuci import add_credential_args
from utility import configure_logging

def prepare_dummy_env(client):
    """Use a client to prepare a dummy environment."""
    client.deploy('local:dummy-source')
    client.deploy('local:dummy-sink')
    token = get_random_string()
    client.juju('set', ('dummy-source', 'token=%s' % token))
    client.juju('add-relation', ('dummy-source', 'dummy-sink'))
    client.juju('expose', ('dummy-sink',))
    return token

def get_clients(initial, other, base_env, debug, agent_url):
    """Return the clients to use for testing."""
    environment = SimpleEnvironment.from_config(base_env)
    if agent_url is None:
        environment.config.pop('tools-metadata-url', None)
    initial_client = EnvJujuClient.by_version(environment, initial,
                                              debug=debug)
    other_client = EnvJujuClient.by_version(environment, other, debug=debug)
    # System juju is assumed to be released and the best choice for tearing
    # down environments reliably.  (For example, 1.18.x cannot tear down
    # environments with alpha agent-versions.)
    released_client = EnvJujuClient.by_version(environment, debug=debug)
    return initial_client, other_client, released_client

def assess_heterogeneous(initial, other, base_env, environment_name, log_dir,
                         upload_tools, debug, agent_url, agent_stream, series):
    """Top level function that prepares the clients and environment.

    initial and other are paths to the binary used initially, and a binary
    used later.  base_env is the name of the environment to base the
    environment on and environment_name is the new name for the environment.
    """
    initial_client, other_client, released_client = get_clients(
        initial, other, base_env, debug, agent_url)
    jes_enabled = initial_client.is_jes_enabled()
    bs_manager = BootstrapManager(
        environment_name, initial_client, released_client,
        bootstrap_host=None, machines=[], series=series, agent_url=agent_url,
        agent_stream=agent_stream, region=None, log_dir=log_dir,
        keep_env=False, permanent=jes_enabled, jes_enabled=jes_enabled)
    test_control_heterogeneous(bs_manager, other_client, upload_tools)

@contextmanager
def run_context(bs_manager, other, upload_tools):
    try:
        bs_manager.keep_env = True
        with bs_manager.booted_context(upload_tools):
            yield
        # Test clean shutdown of an environment.
        juju_with_fallback(other, bs_manager.tear_down_client,
                           'destroy-environment',
                           (other.env.environment, '-y'), include_e=False)
    except:
        bs_manager.tear_down()
        raise

def test_control_heterogeneous(bs_manager, other, upload_tools):
    """Test if one binary can control an environment set up by the other."""
    initial = bs_manager.client
    released = bs_manager.tear_down_client
    with run_context(bs_manager, other, upload_tools):
        token = prepare_dummy_env(initial)
        initial.wait_for_started()
        if sys.platform != "win32":
            # Currently, juju ssh is not working on Windows.
            check_token(initial, token)
            check_series(other)
            other.juju('run', ('--all', 'uname -a'))
        other.get_juju_output('get', 'dummy-source')
        other.get_juju_output('get-env')
        other.juju('remove-relation', ('dummy-source', 'dummy-sink'))
        status = other.get_status()
        other.juju('unexpose', ('dummy-sink',))
        status = other.get_status()
        juju_with_fallback(other, released, 'deploy',
                           ('local:dummy-sink', 'sink2'))
        other.wait_for_started()
        other.juju('add-relation', ('dummy-source', 'sink2'))
        status = other.get_status()
        other.juju('expose', ('sink2',))
        status = other.get_status()
        if 'sink2' not in status.status['services']:
            raise AssertionError('Sink2 missing')
        other.juju('destroy-service', ('sink2',))
        for ignored in until_timeout(30):
            status = other.get_status()
            if 'sink2' not in status.status['services']:
                break
        else:
            raise AssertionError('Sink2 not destroyed')
        other.juju('add-relation', ('dummy-source', 'dummy-sink'))
        status = other.get_status()
        relations = status.status['services']['dummy-sink']['relations']
        if not relations['source'] == ['dummy-source']:
            raise AssertionError('source is not dummy-source.')
        other.juju('expose', ('dummy-sink',))
        status = other.get_status()
        if not status.status['services']['dummy-sink']['exposed']:
            raise AssertionError('dummy-sink is not exposed')
        other.juju('add-unit', ('dummy-sink',))
        if not has_agent(other, 'dummy-sink/1'):
            raise AssertionError('dummy-sink/1 was not added.')
        other.juju('remove-unit', ('dummy-sink/1',))
        status = other.get_status()
        if has_agent(other, 'dummy-sink/1'):
            raise AssertionError('dummy-sink/1 was not removed.')
        other.juju('add-machine', ('lxc',))
        status = other.get_status()
        lxc_machine, = set(k for k, v in status.agent_items() if
                           k.endswith('/lxc/0'))
        lxc_holder = lxc_machine.split('/')[0]
        other.juju('remove-machine', (lxc_machine,))
        wait_until_removed(other, lxc_machine)
        other.juju('remove-machine', (lxc_holder,))
        wait_until_removed(other, lxc_holder)

def juju_with_fallback(other, released, command, args, include_e=True):
    """Fallback to released juju when 1.18 fails.

    Get as much test coverage of 1.18 as we can, by falling back to a released
    juju for commands that we expect to fail (due to unsupported agent version
    format).
    """
    for client in [other, released]:
        try:
            client.juju(command, args, include_e=include_e)
        except CalledProcessError:
            if not client.version.startswith('1.18.'):
                raise
        else:
            break

def has_agent(client, agent_id):
    return bool(agent_id in dict(client.get_status().agent_items()))

def wait_until_removed(client, agent_id):
    """Wait for an agent to be removed from the environment."""
    for ignored in until_timeout(240):
        if not has_agent(client, agent_id):
            return
    else:
        raise AssertionError('Machine not destroyed: {}.'.format(agent_id))

def check_series(client):
    """Use 'juju ssh' to check that the deployed series meets expectations."""
    result = client.get_juju_output('ssh', '0', 'lsb_release', '-c')
    label, codename = result.rstrip().split('\t')
    if label != 'Codename:':
        raise AssertionError()
    expected_codename = client.env.config['default-series']
    if codename != expected_codename:
        raise AssertionError(
            'Series is {}, not {}'.format(codename, expected_codename))

def parse_args(argv=None):
    parser = ArgumentParser(description=dedent("""\
        Determine whether one juju version can control an environment created
        by another version.
    """))
    parser.add_argument('initial', help='The initial juju binary.')
    parser.add_argument('other', help='A different juju binary.')
    parser.add_argument('base_environment', help='The environment to base on.')
    parser.add_argument('environment_name', help='The new environment name.')
    parser.add_argument('log_dir', help='The directory to dump logs to.')
    parser.add_argument(
        '--upload-tools', action='store_true', default=False,
        help='Upload local version of tools before bootstrapping.')
    parser.add_argument('--debug', help='Run juju with --debug',
                        action='store_true', default=False)
    parser.add_argument('--agent-url', default=None)
    parser.add_argument('--agent-stream', action='store',
                        help='URL for retrieving agent binaries.')
    parser.add_argument('--series', action='store',
                        help='Name of the Ubuntu series to use.')
    add_credential_args(parser)
    return parser.parse_args(argv)

def main():
    args = parse_args()
    configure_logging(logging.INFO)
    assess_heterogeneous(args.initial, args.other, args.base_environment,
                         args.environment_name, args.log_dir,
                         args.upload_tools, args.debug, args.agent_url,
                         args.agent_stream, args.series)

if __name__ == '__main__':
    main()
<commit_msg>Fix failing test.<commit_after>#!/usr/bin/env python

from argparse import ArgumentParser
from contextlib import contextmanager
import logging
from textwrap import dedent
from subprocess import CalledProcessError
import sys

from jujupy import (
    EnvJujuClient,
    SimpleEnvironment,
    until_timeout,
    )
from deploy_stack import (
    BootstrapManager,
    check_token,
    get_random_string,
    )
from jujuci import add_credential_args
from utility import configure_logging

def prepare_dummy_env(client):
    """Use a client to prepare a dummy environment."""
    client.deploy('local:dummy-source')
    client.deploy('local:dummy-sink')
    token = get_random_string()
    client.juju('set', ('dummy-source', 'token=%s' % token))
    client.juju('add-relation', ('dummy-source', 'dummy-sink'))
    client.juju('expose', ('dummy-sink',))
    return token

def get_clients(initial, other, base_env, debug, agent_url):
    """Return the clients to use for testing."""
    environment = SimpleEnvironment.from_config(base_env)
    if agent_url is None:
        environment.config.pop('tools-metadata-url', None)
    initial_client = EnvJujuClient.by_version(environment, initial,
                                              debug=debug)
    other_client = EnvJujuClient.by_version(environment, other, debug=debug)
    # System juju is assumed to be released and the best choice for tearing
    # down environments reliably.  (For example, 1.18.x cannot tear down
    # environments with alpha agent-versions.)
    released_client = EnvJujuClient.by_version(environment, debug=debug)
    return initial_client, other_client, released_client

def assess_heterogeneous(initial, other, base_env, environment_name, log_dir,
                         upload_tools, debug, agent_url, agent_stream, series):
    """Top level function that prepares the clients and environment.

    initial and other are paths to the binary used initially, and a binary
    used later.  base_env is the name of the environment to base the
    environment on and environment_name is the new name for the environment.
    """
    initial_client, other_client, released_client = get_clients(
        initial, other, base_env, debug, agent_url)
    jes_enabled = initial_client.is_jes_enabled()
    bs_manager = BootstrapManager(
        environment_name, initial_client, released_client,
        bootstrap_host=None, machines=[], series=series, agent_url=agent_url,
        agent_stream=agent_stream, region=None, log_dir=log_dir,
        keep_env=False, permanent=jes_enabled, jes_enabled=jes_enabled)
    test_control_heterogeneous(bs_manager, other_client, upload_tools)

@contextmanager
def run_context(bs_manager, other, upload_tools):
    try:
        bs_manager.keep_env = True
        with bs_manager.booted_context(upload_tools):
            other.juju_home = bs_manager.client.juju_home
            yield
        # Test clean shutdown of an environment.
        juju_with_fallback(other, bs_manager.tear_down_client,
                           'destroy-environment',
                           (other.env.environment, '-y'), include_e=False)
    except:
        bs_manager.tear_down()
        raise

def test_control_heterogeneous(bs_manager, other, upload_tools):
    """Test if one binary can control an environment set up by the other."""
    initial = bs_manager.client
    released = bs_manager.tear_down_client
    with run_context(bs_manager, other, upload_tools):
        token = prepare_dummy_env(initial)
        initial.wait_for_started()
        if sys.platform != "win32":
            # Currently, juju ssh is not working on Windows.
            check_token(initial, token)
            check_series(other)
            other.juju('run', ('--all', 'uname -a'))
        other.get_juju_output('get', 'dummy-source')
        other.get_juju_output('get-env')
        other.juju('remove-relation', ('dummy-source', 'dummy-sink'))
        status = other.get_status()
        other.juju('unexpose', ('dummy-sink',))
        status = other.get_status()
        juju_with_fallback(other, released, 'deploy',
                           ('local:dummy-sink', 'sink2'))
        other.wait_for_started()
        other.juju('add-relation', ('dummy-source', 'sink2'))
        status = other.get_status()
        other.juju('expose', ('sink2',))
        status = other.get_status()
        if 'sink2' not in status.status['services']:
            raise AssertionError('Sink2 missing')
        other.juju('destroy-service', ('sink2',))
        for ignored in until_timeout(30):
            status = other.get_status()
            if 'sink2' not in status.status['services']:
                break
        else:
            raise AssertionError('Sink2 not destroyed')
        other.juju('add-relation', ('dummy-source', 'dummy-sink'))
        status = other.get_status()
        relations = status.status['services']['dummy-sink']['relations']
        if not relations['source'] == ['dummy-source']:
            raise AssertionError('source is not dummy-source.')
        other.juju('expose', ('dummy-sink',))
        status = other.get_status()
        if not status.status['services']['dummy-sink']['exposed']:
            raise AssertionError('dummy-sink is not exposed')
        other.juju('add-unit', ('dummy-sink',))
        if not has_agent(other, 'dummy-sink/1'):
            raise AssertionError('dummy-sink/1 was not added.')
        other.juju('remove-unit', ('dummy-sink/1',))
        status = other.get_status()
        if has_agent(other, 'dummy-sink/1'):
            raise AssertionError('dummy-sink/1 was not removed.')
        other.juju('add-machine', ('lxc',))
        status = other.get_status()
        lxc_machine, = set(k for k, v in status.agent_items() if
                           k.endswith('/lxc/0'))
        lxc_holder = lxc_machine.split('/')[0]
        other.juju('remove-machine', (lxc_machine,))
        wait_until_removed(other, lxc_machine)
        other.juju('remove-machine', (lxc_holder,))
        wait_until_removed(other, lxc_holder)

def juju_with_fallback(other, released, command, args, include_e=True):
    """Fallback to released juju when 1.18 fails.

    Get as much test coverage of 1.18 as we can, by falling back to a released
    juju for commands that we expect to fail (due to unsupported agent version
    format).
    """
    for client in [other, released]:
        try:
            client.juju(command, args, include_e=include_e)
        except CalledProcessError:
            if not client.version.startswith('1.18.'):
                raise
        else:
            break

def has_agent(client, agent_id):
    return bool(agent_id in dict(client.get_status().agent_items()))

def wait_until_removed(client, agent_id):
    """Wait for an agent to be removed from the environment."""
    for ignored in until_timeout(240):
        if not has_agent(client, agent_id):
            return
    else:
        raise AssertionError('Machine not destroyed: {}.'.format(agent_id))

def check_series(client):
    """Use 'juju ssh' to check that the deployed series meets expectations."""
    result = client.get_juju_output('ssh', '0', 'lsb_release', '-c')
    label, codename = result.rstrip().split('\t')
    if label != 'Codename:':
        raise AssertionError()
    expected_codename = client.env.config['default-series']
    if codename != expected_codename:
        raise AssertionError(
            'Series is {}, not {}'.format(codename, expected_codename))

def parse_args(argv=None):
    parser = ArgumentParser(description=dedent("""\
        Determine whether one juju version can control an environment created
        by another version.
    """))
    parser.add_argument('initial', help='The initial juju binary.')
    parser.add_argument('other', help='A different juju binary.')
    parser.add_argument('base_environment', help='The environment to base on.')
    parser.add_argument('environment_name', help='The new environment name.')
    parser.add_argument('log_dir', help='The directory to dump logs to.')
    parser.add_argument(
        '--upload-tools', action='store_true', default=False,
        help='Upload local version of tools before bootstrapping.')
    parser.add_argument('--debug', help='Run juju with --debug',
                        action='store_true', default=False)
    parser.add_argument('--agent-url', default=None)
    parser.add_argument('--agent-stream', action='store',
                        help='URL for retrieving agent binaries.')
    parser.add_argument('--series', action='store',
                        help='Name of the Ubuntu series to use.')
    add_credential_args(parser)
    return parser.parse_args(argv)

def main():
    args = parse_args()
    configure_logging(logging.INFO)
    assess_heterogeneous(args.initial, args.other, args.base_environment,
                         args.environment_name, args.log_dir,
                         args.upload_tools, args.debug, args.agent_url,
                         args.agent_stream, args.series)

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


<commit_before>"""
Preprocess pipeline
"""
import datetime
import logging
import os.path

import numpy as np

from .. import read_config
from ..batch import BatchProcessorFactory

from .detect import threshold_detection
from .filter import whitening_matrix, whitening, localized_whitening_matrix, whitening_score, butterworth
from .score import get_score_pca, get_pca_suff_stat, get_pca_projection
from .standarize import standarize, sd
from ..neuralnetwork import NeuralNetDetector, NeuralNetTriage, nn_detection

# remove this
Q = None
Q_score = None

def run():
    """Execute preprocessing pipeline

    Returns
    -------
    score: list
        List of size n_channels, each list contains a (clear spikes x
        number of features x number of channels) multidimensional array
        score for every clear spike

    clear_index: list
        List of size n_channels, each list contains the indexes in
        spike_times (first column) where the spike was clear

    spike_times: list
        List with n_channels elements, each element contains spike times
        in the first column and [SECOND COLUMN?]

    Examples
    --------

    .. literalinclude:: ../examples/preprocess.py
    """
    logger = logging.getLogger(__name__)

    start_time = datetime.datetime.now()
    time = {'f': 0, 's': 0, 'd': 0, 'w': 0, 'b': 0, 'e': 0}

    # FIXME: remove this
    CONFIG = read_config()
    whiten_file = open(os.path.join(CONFIG.data.root_folder, 'tmp/whiten.bin'), 'wb')

    # initialize processor for raw data
    path = os.path.join(CONFIG.data.root_folder, CONFIG.data.recordings)
    dtype = CONFIG.recordings.dtype

    # initialize factory
    factory = BatchProcessorFactory(path_to_file=None,
                                    dtype=None,
                                    n_channels=CONFIG.recordings.n_channels,
                                    max_memory=CONFIG.resources.max_memory,
                                    buffer_size=None)

    if CONFIG.preprocess.filter == 1:

        _b = datetime.datetime.now()
        # make batch processor for raw data -> buterworth -> filtered
        bp = factory.make(path_to_file=path, dtype=dtype,
                          buffer_size=0)
        logger.info('Initialized butterworth batch processor: {}'
                    .format(bp))

        # run filtering
        path = os.path.join(CONFIG.data.root_folder,  'tmp/filtered.bin')
        dtype = bp.process_function(butterworth,
                                    path,
                                    CONFIG.filter.low_pass_freq,
                                    CONFIG.filter.high_factor,
                                    CONFIG.filter.order,
                                    CONFIG.recordings.sampling_rate)
        time['f'] += (datetime.datetime.now()-_b).total_seconds()

    # TODO: cache computations
    # make batch processor for filtered -> standarize -> standarized
    _b = datetime.datetime.now()
    bp = factory.make(path_to_file=path, dtype=dtype, buffer_size=0)

    # compute the standard deviation using the first batch only
    batch1 = next(bp)
    sd_ = sd(batch1, CONFIG.recordings.sampling_rate)

    # make another batch processor
    bp = factory.make(path_to_file=path, dtype=dtype, buffer_size=0)
    logger.info('Initialized standarization batch processor: {}'
                .format(bp))

    # run standarization
    path = os.path.join(CONFIG.data.root_folder,  'tmp/standarized.bin')
    dtype = bp.process_function(standarize,
                                path,
                                sd_)
    time['s'] += (datetime.datetime.now()-_b).total_seconds()

    # create another batch processor for the rest of the pipeline
    bp = factory.make(path_to_file=path, dtype=dtype,
                      buffer_size=CONFIG.BUFF)
    logger.info('Initialized preprocess batch processor: {}'
                .format(bp))

    # initialize output variables
    get_score = 1
    spike_index_clear = None
    spike_index_collision = None
    score = None
    pca_suff_stat = None
    spikes_per_channel = None

    for i, batch in enumerate(bp):

        # load nueral net detector if necessary:
        if CONFIG.spikes.detection == 'nn':
            nnDetector = NeuralNetDetector(CONFIG.neural_network_detector.filename,
                                           CONFIG.neural_network_autoencoder.filename)
            nnTriage = NeuralNetTriage(CONFIG.neural_network_triage.filename)

        else:
            nnDetector = None
            nnTriage = None

        if i > CONFIG.nPortion:
            get_score = 0

        # process batch
        # spike index is defined as a location in each minibatch
        (si_clr_batch, score_batch, si_col_batch,
         pss_batch, spc_batch,
         time) = process_batch(batch, get_score, CONFIG.BUFF, time,
                               nnDetector=nnDetector,
                               nnTriage=nnTriage, whiten_file=whiten_file)

        # spike time w.r.t. to the whole recording
        si_clr_batch[:,0] = si_clr_batch[:,0] + i*CONFIG.batch_size - CONFIG.BUFF
        si_col_batch[:,0] = si_col_batch[:,0] + i*CONFIG.batch_size - CONFIG.BUFF

        if i == 0:
            spike_index_clear = si_clr_batch
            spike_index_collision = si_col_batch
            score = score_batch

            pca_suff_stat = pss_batch
            spikes_per_channel = spc_batch

        else:
            spike_index_clear = np.vstack((spike_index_clear,
                si_clr_batch))
            spike_index_collision = np.vstack((spike_index_collision,
                si_col_batch))
            if get_score == 1 and CONFIG.spikes.detection == 'nn':
                score = np.concatenate((score, score_batch), axis = 0)
            pca_suff_stat += pss_batch
            spikes_per_channel += spc_batch

    whiten_file.close()

    if CONFIG.spikes.detection != 'nn':
        _b = datetime.datetime.now()
        rot = get_pca_projection(pca_suff_stat, spikes_per_channel,
                                 CONFIG.spikes.temporal_features, CONFIG.neighChannels)
        score = get_score_pca(spike_index_clear, rot, CONFIG.neighChannels,
                              CONFIG.geom, CONFIG.batch_size,
                              CONFIG.BUFF, CONFIG.nBatches,
                              os.path.join(CONFIG.data.root_folder,'tmp/whiten.bin'),
                              CONFIG.scaleToSave)

        time['e'] += (datetime.datetime.now()-_b).total_seconds()

    # timing
    current_time = datetime.datetime.now()
    logger.info("Preprocessing done in {0} seconds.".format(
                     (current_time-start_time).seconds))
    logger.info("\tfiltering:\t{0} seconds".format(time['f']))
    logger.info("\tstandardization:\t{0} seconds".format(time['s']))
    logger.info("\tdetection:\t{0} seconds".format(time['d']))
    logger.info("\twhitening:\t{0} seconds".format(time['w']))
    logger.info("\tsaving recording:\t{0} seconds".format(time['b']))
    logger.info("\tgetting waveforms:\t{0} seconds".format(time['e']))

    return score, spike_index_clear, spike_index_collision

def process_batch(rec, get_score, BUFF, time, nnDetector, nnTriage,
                  whiten_file):
    logger = logging.getLogger(__name__)
    CONFIG = read_config()

    global Q
    global Q_score

    # nn detection
    if CONFIG.spikes.detection == 'nn':

        # detect spikes
        _b = datetime.datetime.now()
        (spike_index_clear,
         spike_index_collision,
         score) = nn_detection(rec, 10000, BUFF,
                               CONFIG.neighChannels,
                               CONFIG.geom,
                               CONFIG.spikes.temporal_features,
                               3,
                               CONFIG.neural_network_detector.threshold_spike,
                               CONFIG.neural_network_triage.threshold_collision,
                               nnDetector,
                               nnTriage
                              )

        # since we alread have scores, no need to calculated sufficient
        # statistics for pca
        pca_suff_stat = 0
        spikes_per_channel = 0

        time['d'] += (datetime.datetime.now()-_b).total_seconds()

        if get_score ==0:
            spike_index_clear = np.zeros((0,2), 'int32')
            spike_index_collision = np.vstack((spike_index_collision,
                                       spike_index_clear))
            score = None

        else:
            # whiten signal
            _b = datetime.datetime.now()
            # get withening matrix per batch or onece in total
            if CONFIG.preprocess.whiten_batchwise or Q is None:
                Q_score = localized_whitening_matrix(rec,
                                               CONFIG.neighChannels,
                                               CONFIG.geom,
                                               CONFIG.spikeSize)
            score = whitening_score(score, spike_index_clear[:,1], Q_score)

            time['w'] += (datetime.datetime.now()-_b).total_seconds()

    # threshold detection
    elif CONFIG.spikes.detection == 'threshold':

        # detect spikes
        _b = datetime.datetime.now()
        spike_index = threshold_detection(rec,
                                          CONFIG.neighChannels,
                                          CONFIG.spikeSize,
                                          CONFIG.stdFactor)

        # every spikes are considered as clear spikes as no triage is done
        if get_score ==0:
            spike_index_clear = np.zeros((0,2), 'int32')
            spike_index_collision = spike_index
        else:
            spike_index_clear = spike_index
            spike_index_collision = np.zeros((0,2), 'int32')
        score = None

        # get sufficient statistics for pca if we don't have projection matrix
        pca_suff_stat, spikes_per_channel = get_pca_suff_stat(rec, spike_index,
                                                          CONFIG.spikeSize)

        time['d'] += (datetime.datetime.now()-_b).total_seconds()

        # whiten recording
        _b = datetime.datetime.now()
        if CONFIG.preprocess.whiten_batchwise or Q is None:
            Q = whitening_matrix(rec, CONFIG.neighChannels,
                                CONFIG.spikeSize)
        rec = whitening(rec, Q)

        time['w'] += (datetime.datetime.now()-_b).total_seconds()

    # Remove spikes detectted in buffer area
    spike_index_clear = spike_index_clear[np.logical_and(
      spike_index_clear[:, 0] > BUFF,
      spike_index_clear[:, 0] < (rec.shape[0] - BUFF))]
    spike_index_collision = spike_index_collision[np.logical_and(
      spike_index_collision[:, 0] > BUFF,
      spike_index_collision[:, 0] < (rec.shape[0] - BUFF))]

    _b = datetime.datetime.now()

    # save whiten data
    chunk = rec*CONFIG.scaleToSave
    chunk.reshape(chunk.shape[0]*chunk.shape[1])
    chunk.astype('int16').tofile(whiten_file)

    time['b'] += (datetime.datetime.now()-_b).total_seconds()

    return (spike_index_clear, score, spike_index_collision,
        pca_suff_stat, spikes_per_channel, time)
<commit_msg>saving partial results<commit_after>"""
Preprocess pipeline
"""
import datetime
import logging
import os.path

import numpy as np

from .. import read_config
from ..batch import BatchProcessorFactory

from .detect import threshold_detection
from .filter import whitening_matrix, whitening, localized_whitening_matrix, whitening_score, butterworth
from .score import get_score_pca, get_pca_suff_stat, get_pca_projection
from .standarize import standarize, sd
from ..neuralnetwork import NeuralNetDetector, NeuralNetTriage, nn_detection

SAVE_PARTIAL_RESULTS = True

# remove this
Q = None
Q_score = None

def run():
    """Execute preprocessing pipeline

    Returns
    -------
    score: list
        List of size n_channels, each list contains a (clear spikes x
        number of features x number of channels) multidimensional array
        score for every clear spike

    clear_index: list
        List of size n_channels, each list contains the indexes in
        spike_times (first column) where the spike was clear

    spike_times: list
        List with n_channels elements, each element contains spike times
        in the first column and [SECOND COLUMN?]

    Examples
    --------

    .. literalinclude:: ../examples/preprocess.py
    """
    logger = logging.getLogger(__name__)

    start_time = datetime.datetime.now()
    time = {'f': 0, 's': 0, 'd': 0, 'w': 0, 'b': 0, 'e': 0}

    # FIXME: remove this
    CONFIG = read_config()
    whiten_file = open(os.path.join(CONFIG.data.root_folder, 'tmp/whiten.bin'), 'wb')

    # initialize processor for raw data
    path = os.path.join(CONFIG.data.root_folder, CONFIG.data.recordings)
    dtype = CONFIG.recordings.dtype

    # initialize factory
    factory = BatchProcessorFactory(path_to_file=None,
                                    dtype=None,
                                    n_channels=CONFIG.recordings.n_channels,
                                    max_memory=CONFIG.resources.max_memory,
                                    buffer_size=None)

    if CONFIG.preprocess.filter == 1:

        _b = datetime.datetime.now()
        # make batch processor for raw data -> buterworth -> filtered
        bp = factory.make(path_to_file=path, dtype=dtype,
                          buffer_size=0)
        logger.info('Initialized butterworth batch processor: {}'
                    .format(bp))

        # run filtering
        path = os.path.join(CONFIG.data.root_folder,  'tmp/filtered.bin')
        dtype = bp.process_function(butterworth,
                                    path,
                                    CONFIG.filter.low_pass_freq,
                                    CONFIG.filter.high_factor,
                                    CONFIG.filter.order,
                                    CONFIG.recordings.sampling_rate)
        time['f'] += (datetime.datetime.now()-_b).total_seconds()

    # TODO: cache computations
    # make batch processor for filtered -> standarize -> standarized
    _b = datetime.datetime.now()
    bp = factory.make(path_to_file=path, dtype=dtype, buffer_size=0)

    # compute the standard deviation using the first batch only
    batch1 = next(bp)
    sd_ = sd(batch1, CONFIG.recordings.sampling_rate)

    # make another batch processor
    bp = factory.make(path_to_file=path, dtype=dtype, buffer_size=0)
    logger.info('Initialized standarization batch processor: {}'
                .format(bp))

    # run standarization
    path = os.path.join(CONFIG.data.root_folder,  'tmp/standarized.bin')
    dtype = bp.process_function(standarize,
                                path,
                                sd_)
    time['s'] += (datetime.datetime.now()-_b).total_seconds()

    # create another batch processor for the rest of the pipeline
    bp = factory.make(path_to_file=path, dtype=dtype,
                      buffer_size=CONFIG.BUFF)
    logger.info('Initialized preprocess batch processor: {}'
                .format(bp))

    # initialize output variables
    get_score = 1
    spike_index_clear = None
    spike_index_collision = None
    score = None
    pca_suff_stat = None
    spikes_per_channel = None

    for i, batch in enumerate(bp):

        # load nueral net detector if necessary:
        if CONFIG.spikes.detection == 'nn':
            nnDetector = NeuralNetDetector(CONFIG.neural_network_detector.filename,
                                           CONFIG.neural_network_autoencoder.filename)
            nnTriage = NeuralNetTriage(CONFIG.neural_network_triage.filename)

        else:
            nnDetector = None
            nnTriage = None

        if i > CONFIG.nPortion:
            get_score = 0

        # process batch
        # spike index is defined as a location in each minibatch
        (si_clr_batch, score_batch, si_col_batch,
         pss_batch, spc_batch,
         time) = process_batch(batch, get_score, CONFIG.BUFF, time,
                               nnDetector=nnDetector,
                               nnTriage=nnTriage, whiten_file=whiten_file)

        # spike time w.r.t. to the whole recording
        si_clr_batch[:,0] = si_clr_batch[:,0] + i*CONFIG.batch_size - CONFIG.BUFF
        si_col_batch[:,0] = si_col_batch[:,0] + i*CONFIG.batch_size - CONFIG.BUFF

        if i == 0:
            spike_index_clear = si_clr_batch
            spike_index_collision = si_col_batch
            score = score_batch

            pca_suff_stat = pss_batch
            spikes_per_channel = spc_batch

        else:
            spike_index_clear = np.vstack((spike_index_clear,
                si_clr_batch))
            spike_index_collision = np.vstack((spike_index_collision,
                si_col_batch))
            if get_score == 1 and CONFIG.spikes.detection == 'nn':
                score = np.concatenate((score, score_batch), axis = 0)
            pca_suff_stat += pss_batch
            spikes_per_channel += spc_batch

    whiten_file.close()

    if CONFIG.spikes.detection != 'nn':
        _b = datetime.datetime.now()
        rot = get_pca_projection(pca_suff_stat, spikes_per_channel,
                                 CONFIG.spikes.temporal_features, CONFIG.neighChannels)
        score = get_score_pca(spike_index_clear, rot, CONFIG.neighChannels,
                              CONFIG.geom, CONFIG.batch_size,
                              CONFIG.BUFF, CONFIG.nBatches,
                              os.path.join(CONFIG.data.root_folder,'tmp/whiten.bin'),
                              CONFIG.scaleToSave)

        time['e'] += (datetime.datetime.now()-_b).total_seconds()

    # timing
    current_time = datetime.datetime.now()
    logger.info("Preprocessing done in {0} seconds.".format(
                     (current_time-start_time).seconds))
    logger.info("\tfiltering:\t{0} seconds".format(time['f']))
    logger.info("\tstandardization:\t{0} seconds".format(time['s']))
    logger.info("\tdetection:\t{0} seconds".format(time['d']))
    logger.info("\twhitening:\t{0} seconds".format(time['w']))
    logger.info("\tsaving recording:\t{0} seconds".format(time['b']))
    logger.info("\tgetting waveforms:\t{0} seconds".format(time['e']))

    # save partial results...
    if SAVE_PARTIAL_RESULTS:
        logger.info('Saving partial results...')
        TMP = os.path.join(CONFIG.data.root_folder, 'tmp/')
        np.save(os.path.join(TMP, 'score.npy'), score)
        np.save(os.path.join(TMP, 'spike_index_clear.npy'), spike_index_clear)
        np.save(os.path.join(TMP, 'spike_index_collision.npy'),
                spike_index_collision)

    return score, spike_index_clear, spike_index_collision

def process_batch(rec, get_score, BUFF, time, nnDetector, nnTriage,
                  whiten_file):
    logger = logging.getLogger(__name__)
    CONFIG = read_config()

    global Q
    global Q_score

    # nn detection
    if CONFIG.spikes.detection == 'nn':

        # detect spikes
        _b = datetime.datetime.now()
        (spike_index_clear,
         spike_index_collision,
         score) = nn_detection(rec, 10000, BUFF,
                               CONFIG.neighChannels,
                               CONFIG.geom,
                               CONFIG.spikes.temporal_features,
                               3,
                               CONFIG.neural_network_detector.threshold_spike,
                               CONFIG.neural_network_triage.threshold_collision,
                               nnDetector,
                               nnTriage
                              )

        # since we alread have scores, no need to calculated sufficient
        # statistics for pca
        pca_suff_stat = 0
        spikes_per_channel = 0

        time['d'] += (datetime.datetime.now()-_b).total_seconds()

        if get_score ==0:
            spike_index_clear = np.zeros((0,2), 'int32')
            spike_index_collision = np.vstack((spike_index_collision,
                                       spike_index_clear))
            score = None

        else:
            # whiten signal
            _b = datetime.datetime.now()
            # get withening matrix per batch or onece in total
            if CONFIG.preprocess.whiten_batchwise or Q is None:
                Q_score = localized_whitening_matrix(rec,
                                               CONFIG.neighChannels,
                                               CONFIG.geom,
                                               CONFIG.spikeSize)
            score = whitening_score(score, spike_index_clear[:,1], Q_score)

            time['w'] += (datetime.datetime.now()-_b).total_seconds()

    # threshold detection
    elif CONFIG.spikes.detection == 'threshold':

        # detect spikes
        _b = datetime.datetime.now()
        spike_index = threshold_detection(rec,
                                          CONFIG.neighChannels,
                                          CONFIG.spikeSize,
                                          CONFIG.stdFactor)

        # every spikes are considered as clear spikes as no triage is done
        if get_score ==0:
            spike_index_clear = np.zeros((0,2), 'int32')
            spike_index_collision = spike_index
        else:
            spike_index_clear = spike_index
            spike_index_collision = np.zeros((0,2), 'int32')
        score = None

        # get sufficient statistics for pca if we don't have projection matrix
        pca_suff_stat, spikes_per_channel = get_pca_suff_stat(rec, spike_index,
                                                          CONFIG.spikeSize)

        time['d'] += (datetime.datetime.now()-_b).total_seconds()

        # whiten recording
        _b = datetime.datetime.now()
        if CONFIG.preprocess.whiten_batchwise or Q is None:
            Q = whitening_matrix(rec, CONFIG.neighChannels,
                                CONFIG.spikeSize)
        rec = whitening(rec, Q)

        time['w'] += (datetime.datetime.now()-_b).total_seconds()

    # Remove spikes detectted in buffer area
    spike_index_clear = spike_index_clear[np.logical_and(
      spike_index_clear[:, 0] > BUFF,
      spike_index_clear[:, 0] < (rec.shape[0] - BUFF))]
    spike_index_collision = spike_index_collision[np.logical_and(
      spike_index_collision[:, 0] > BUFF,
      spike_index_collision[:, 0] < (rec.shape[0] - BUFF))]

    _b = datetime.datetime.now()

    # save whiten data
    chunk = rec*CONFIG.scaleToSave
    chunk.reshape(chunk.shape[0]*chunk.shape[1])
    chunk.astype('int16').tofile(whiten_file)

    time['b'] += (datetime.datetime.now()-_b).total_seconds()

    return (spike_index_clear, score, spike_index_collision,
        pca_suff_stat, spikes_per_channel, time)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
import sys
import os

from PyQt4.QtGui import QApplication, QTableWidget, QTableWidgetItem
from PyQt4.QtCore import QUrl
from PyQt4.QtWebKit import QWebView, QWebPage
from PyQt4.QtGui import QGridLayout, QLineEdit, QWidget, QHeaderView
from PyQt4.QtNetwork import QNetworkAccessManager, QNetworkRequest

class UrlInput(QLineEdit):
    def __init__(self, browser):
        super(UrlInput, self).__init__()
        self.browser = browser
        self.returnPressed.connect(self._return_pressed)

    def _return_pressed(self):
        url = QUrl(self.text())
        browser.load(url)

class JavaScriptEvaluator(QLineEdit):
    def __init__(self, page):
        super(JavaScriptEvaluator, self).__init__()
        self.page = page
        self.returnPressed.connect(self._return_pressed)

    def _return_pressed(self):
        frame = self.page.currentFrame()
        result = frame.evaluateJavaScript(self.text())

class ActionInputBox(QLineEdit):
    def __init__(self, page):
        super(ActionInputBox, self).__init__()
        self.page = page
        self.returnPressed.connect(self._return_pressed)

    def _return_pressed(self):
        frame = self.page.currentFrame()
        action_string = str(self.text()).lower()
        if action_string == "b":
            self.page.triggerAction(QWebPage.Back)
        elif action_string == "f":
            self.page.triggerAction(QWebPage.Forward)
        elif action_string == "s":
            self.page.triggerAction(QWebPage.Stop)
        elif action_string == "r":
            self.page.triggerAction(QWebPage.Reload)

class RequestsTable(QTableWidget):
    header = ["url", "status", "content-type"]

    def __init__(self):
        super(RequestsTable, self).__init__()
        self.setColumnCount(3)
        self.setHorizontalHeaderLabels(self.header)
        header = self.horizontalHeader()
        header.setStretchLastSection(True)
        header.setResizeMode(QHeaderView.ResizeToContents)

    def update(self, data):
        last_row = self.rowCount()
        next_row = last_row + 1
        self.setRowCount(next_row)
        for col, dat in enumerate(data, 0):
            if not dat:
                continue
            self.setItem(last_row, col, QTableWidgetItem(dat))

class Manager(QNetworkAccessManager):
    def __init__(self, table):
        QNetworkAccessManager.__init__(self)
        self.finished.connect(self._finished)
        self.table = table

    def _finished(self, reply):
        headers = reply.rawHeaderPairs()
        headers = {str(k):str(v) for k,v in headers}
        content_type = headers.get("Content-Type")
        url = reply.url().toString()
        status = reply.attribute(QNetworkRequest.HttpStatusCodeAttribute)
        status, ok = status.toInt()
        self.table.update([url, str(status), content_type])

if __name__ == "__main__":
    app = QApplication(sys.argv)

    grid = QGridLayout()
    browser = QWebView()
    #url_input = UrlInput(browser)
    #requests_table = RequestsTable()

    #manager = Manager(requests_table)
    page = QWebPage()
    #page.setNetworkAccessManager(manager)
    browser.setPage(page)

    browser.load(QUrl('file://%s' % os.path.abspath(sys.argv[1])))

    #js_eval = JavaScriptEvaluator(page)
    #action_box = ActionInputBox(page)

    #grid.addWidget(url_input, 1, 0)
    #grid.addWidget(action_box, 2, 0)
    grid.addWidget(browser, 0, 0)
    #grid.addWidget(requests_table, 4, 0)
    #grid.addWidget(js_eval, 5, 0)

    main_frame = QWidget()
    main_frame.setLayout(grid)
    main_frame.show()

    sys.exit(app.exec_())<commit_msg>Add support for load from Zip file<commit_after>#!/usr/bin/env python
import sys
import os

from zipfile import ZipFile, is_zipfile

from PyQt4.QtGui import QApplication, QTableWidget, QTableWidgetItem
from PyQt4.QtCore import QUrl, QVariant, QTimer, SIGNAL
from PyQt4.QtWebKit import QWebView, QWebPage
from PyQt4.QtGui import QGridLayout, QLineEdit, QWidget, QHeaderView
from PyQt4.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkReply

class UrlInput(QLineEdit):
    def __init__(self, browser):
        super(UrlInput, self).__init__()
        self.browser = browser
        self.returnPressed.connect(self._return_pressed)

    def _return_pressed(self):
        url = QUrl(self.text())
        browser.load(url)

class JavaScriptEvaluator(QLineEdit):
    def __init__(self, page):
        super(JavaScriptEvaluator, self).__init__()
        self.page = page
        self.returnPressed.connect(self._return_pressed)

    def _return_pressed(self):
        frame = self.page.currentFrame()
        result = frame.evaluateJavaScript(self.text())

class ActionInputBox(QLineEdit):
    def __init__(self, page):
        super(ActionInputBox, self).__init__()
        self.page = page
        self.returnPressed.connect(self._return_pressed)

    def _return_pressed(self):
        frame = self.page.currentFrame()
        action_string = str(self.text()).lower()
        if action_string == "b":
            self.page.triggerAction(QWebPage.Back)
        elif action_string == "f":
            self.page.triggerAction(QWebPage.Forward)
        elif action_string == "s":
            self.page.triggerAction(QWebPage.Stop)
        elif action_string == "r":
            self.page.triggerAction(QWebPage.Reload)

class RequestsTable(QTableWidget):
    header = ["url", "status", "content-type"]

    def __init__(self):
        super(RequestsTable, self).__init__()
        self.setColumnCount(3)
        self.setHorizontalHeaderLabels(self.header)
        header = self.horizontalHeader()
        header.setStretchLastSection(True)
        header.setResizeMode(QHeaderView.ResizeToContents)

    def update(self, data):
        last_row = self.rowCount()
        next_row = last_row + 1
        self.setRowCount(next_row)
        for col, dat in enumerate(data, 0):
            if not dat:
                continue
            self.setItem(last_row, col, QTableWidgetItem(dat))

class WamFileReply(QNetworkReply):
    def __init__(self, parent, url):
        super(WamFileReply, self).__init__(parent)
        #self.setHeader(QNetworkRequest.ContentTypeHeader, QVariant('text/html; charset=utf-8'))
        QTimer.singleShot(0, self, SIGNAL("readyRead()"))
        QTimer.singleShot(0, self, SIGNAL("finished()"))
        self.open(self.ReadOnly | self.Unbuffered)
        self.setUrl(url)
        self.content = parent.wam.read(str(url.path()[1:]))

    def __getattribute__(self, name):
        print(name)
        return object.__getattribute__(self, name)

    def abort(self):
        pass

    def isSequential(self):
        return True

    def bytesAvailable(self):
        return len(self.content)

    def readData(self, maxSize):
        _content = self.content[:min(maxSize, self.bytesAvailable)]
        self.content = _content
        return _content

class Manager(QNetworkAccessManager):
    def __init__(self, table, wam=None):
        QNetworkAccessManager.__init__(self)
        self.finished.connect(self._finished)
        self.table = table
        self.wam = wam

    def _finished(self, reply):
        headers = reply.rawHeaderPairs()
        headers = {str(k):str(v) for k,v in headers}
        content_type = headers.get("Content-Type")
        url = reply.url().toString()
        status = reply.attribute(QNetworkRequest.HttpStatusCodeAttribute)
        status, ok = status.toInt()
        self.table.update([url, str(status), content_type])

    def createRequest(self, operation, request, data):
        if request.url().scheme() != "wam" or self.wam is None:
            return super(Manager, self).createRequest(operation, request, data)
        else:
            return WamFileReply(self, request.url())

if __name__ == "__main__":
    app = QApplication(sys.argv)

    grid = QGridLayout()
    browser = QWebView()
    url_input = UrlInput(browser)
    requests_table = RequestsTable()
    wam_file = ZipFile(os.path.abspath(sys.argv[1]))
    manager = Manager(requests_table, wam_file)
    page = QWebPage()
    page.setNetworkAccessManager(manager)
    browser.setPage(page)
    browser.load(QUrl('wam:///test.html'))

    js_eval = JavaScriptEvaluator(page)
    #action_box = ActionInputBox(page)

    grid.addWidget(url_input, 1, 0)
    #grid.addWidget(action_box, 2, 0)
    grid.addWidget(browser, 0, 0)
    grid.addWidget(requests_table, 4, 0)
    grid.addWidget(js_eval, 5, 0)

    main_frame = QWidget()
    main_frame.setLayout(grid)
    main_frame.show()

    sys.exit(app.exec_())<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright 2016 Carl Evans
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import thousandeyes
import solarwinds
import errors
import logging
import configparser

"""Synchronises nodes in Solarwinds with the path monitoring platform Thousand Eyes.

Nodes in Solarwinds that should be created as a test in Thousand Eyes should have a custom boolean property
(configurable in the settings file) set to True.
"""
__author__ = "Carl Evans"
__copyright__ = "Copyright 2016 Carl Evans"
__license__ = "Apache 2.0"
__title__ = "SolarEyes"
__version__ = "1.0"

class SolarEyes(object):
    """SolarEyes class.

    Handles the synchronisation between SolarWinds & ThousandEyes.

    Attributes:
        te_api: A ThousandEyes API instance.
        sw_api: A SolarWinds API instance.
        settings: A SolarEyesSettings instance.
    """
    def __init__(self, te_api=None, sw_api=None, settings=None):
        if te_api is None or sw_api is None or settings is None:
            raise errors.Error("[%s.%s] - You must provide instances of SolarEyesSettings,"
                               "ThousandEyes & SolarWinds APIs." % (__name__, self.__class__.__name__))
        else:
            self.te_api = te_api
            self.sw_api = sw_api
            self.settings = settings

    def get_agent_ids(self, agent_list):
        """Get a list of enterprise agent ids.

        Filter a full list of agents for just the enterprise types and return a list of their ids.

        Args:
            agent_list: A list of thousandeyes.Agent instances.

        Returns:
            A list of integers representing all enterprise agent ids.
        """
        agent_ids = []

        for agent in agent_list:
            if agent.type == "Enterprise":
                agent_ids.append(agent.id)

        return agent_ids

    def get_se_tests(self, test_list):
        """Get a list of tests that were previously created with SolarEyes.

        Filter a full list of tests for those names starting with a specific prefix. This prefix
        is specified in the settings.ini file.

        Args:
            test_list: A list of thousandeyes.NetworkTest instances.

        Returns:
            A filtered list of thousandeyes.NetworkTest instances.
        """
        se_tests = []

        for test in test_list:
            if test.name[0:len(self.settings.te_test_prefix)] == self.settings.te_test_prefix:
                se_tests.append(test)

        return se_tests

    def check_apis(self):
        """Check the availability of both SolarWinds and ThousandEyes APIs.

        Returns:
            A boolean of True if both APIs are ready, or False if an error occurred.
        """
        if self.te_api.status() and self.sw_api.status():
            return True
        else:
            return False

    def create_test(self, test_name=None, test_server=None):
        """Create a test with defaults taken from settings.ini and from passed values.

        Returns:
            A boolean of True if the creation was successful or False if an error occurred.
        """
        if test_name is None or test_server is None:
            raise errors.Error("[%s.%s] - You must provide a test name and server."
                               % (__name__, self.__class__.__name__))
        else:
            test = thousandeyes.NetworkTest()
            test.name = self.settings.te_test_prefix + " " + test_name
            test.server = test_server
            test.alerts_enabled = bool(self.settings.te_test_alerts)
            test.bandwidth_measurements = False
            test.mtu_measurements = True
            test.network_measurements = True
            test.bgp_measurements = True
            test.protocol = self.settings.te_test_protocol
            test.port = int(self.settings.te_test_port)
            test.interval = int(self.settings.te_test_interval)

        for agent_id in self.get_agent_ids(self.te_api.get_agents()):
            test.agents.append({"agentId": agent_id})

        if self.te_api.create_network_test(test):
            return True
        else:
            return False

    def sync(self):
        """Perform a synchronisation between SolarWinds and ThousandEyes.

        Returns:
            A boolean of True if the synchronisation was successful, or False if an error occurred.
        """
        # if self.check_apis():
        #     # Synchronise between SolarWinds and ThousandEyes.
        #
        #     #se_tests = self.get_se_tests(self.te_api.get_network_tests())
        #
        #     # for se_test in se_tests:
        #     #     print(se_test.name)
        #
        #     #if self.create_test("Test", "8.8.8.8"):
        #         #Log test creation here. Need to pass in logger.
        #
        #     self.sw_api.get_path_nodes()
        #
        #     return True
        # else:
        #     return False
        self.sw_api.get_path_nodes()
        return True

class SolarEyesSettings(object):
    """All external settings required by the SolarEyes class.

    Attributes:
        sw_custom_bool: A string containing the base API url such as https://api.thousandeyes.com.
        te_test_protocol: A string containing the ThousandEyes default test protocol.
        te_test_port: A string containing the ThousandEyes default test port.
        te_test_alerts: A string containing the ThousandEyes default test alert requirements.
        te_test_interval: A string containing the ThousandEyes default test interval.
        te_test_prefix: A string containing the ThousandEyes default test name prefix.
    """
    def __init__(self, sw_custom_bool=None, te_test_protocol=None, te_test_port=None,
                 te_test_alerts=None, te_test_interval=None, te_test_prefix=None):

        if sw_custom_bool is None or te_test_protocol is None or te_test_port is None \
                or te_test_alerts is None or te_test_interval is None or te_test_prefix is None:
                    raise errors.Error("[%s.%s] - You must provide all required settings to SolarEyesSettings."
                                       % (__name__, self.__class__.__name__))
        else:
            self.sw_custom_bool = sw_custom_bool
            self.te_test_protocol = te_test_protocol
            self.te_test_port = te_test_port
            self.te_test_alerts = te_test_alerts
            self.te_test_interval = te_test_interval
            self.te_test_prefix = te_test_prefix

def main():
    # Read settings file.
    config = configparser.ConfigParser()
    config.read('settings.ini')
    settings = config['Settings']

    # Configure logging.
    logging.basicConfig(level=logging.DEBUG,
                        filename=settings['log_file'], filemode=settings['log_mode'],
                        format='%(asctime)s %(levelname)s %(message)s')

    try:
        # Create various instances required by the SolarEyes constructor.
        te_api_request = thousandeyes.ApiRequest(settings['te_api'], settings['te_auth_email'], settings['te_auth_token'])
        te_api = thousandeyes.Api(te_api_request)
        sw_api = solarwinds.Api(settings['sw_api'], settings['sw_username'], settings['sw_password'])
        se_settings = SolarEyesSettings(settings['sw_custom_bool'],
                                        settings['te_test_protocol'],
                                        settings['te_test_port'],
                                        settings['te_test_alerts'],
                                        settings['te_test_interval'],
                                        settings['te_test_prefix'])

        # Pass the above API and settings instances to the constructor of SolarEyes.
        solareyes = SolarEyes(te_api, sw_api, se_settings)
    except errors.Error as error:
        logging.error(error)
    else:
        try:
            # Run synchronisation.
            if solareyes.sync():
                logging.info("[%s] - Synchronisation successful." % __name__)
            else:
                logging.error("[%s] - There was a problem during synchronisation."
                              % __name__)
        except errors.Error as error:
            logging.error(error)

# Only execute the main script if passed to an interpreter. Do not execute if imported.
if __name__ == "__main__":
    main()
<commit_msg>Started get_path_nodes()<commit_after>#!/usr/bin/env python
#
# Copyright 2016 Carl Evans
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import thousandeyes
import solarwinds
import errors
import logging
import configparser

"""Synchronises nodes in Solarwinds with the path monitoring platform Thousand Eyes.

Nodes in Solarwinds that should be created as a test in Thousand Eyes should have a custom boolean property
(configurable in the settings file) set to True.
"""
__author__ = "Carl Evans"
__copyright__ = "Copyright 2016 Carl Evans"
__license__ = "Apache 2.0"
__title__ = "SolarEyes"
__version__ = "1.0"

class SolarEyes(object):
    """SolarEyes class.

    Handles the synchronisation between SolarWinds & ThousandEyes.

    Attributes:
        te_api: A ThousandEyes API instance.
        sw_api: A SolarWinds API instance.
        settings: A SolarEyesSettings instance.
    """
    def __init__(self, te_api=None, sw_api=None, settings=None):
        if te_api is None or sw_api is None or settings is None:
            raise errors.Error("[%s.%s] - You must provide instances of SolarEyesSettings,"
                               "ThousandEyes & SolarWinds APIs." % (__name__, self.__class__.__name__))
        else:
            self.te_api = te_api
            self.sw_api = sw_api
            self.settings = settings

    def get_agent_ids(self, agent_list):
        """Get a list of enterprise agent ids.

        Filter a full list of agents for just the enterprise types and return a list of their ids.

        Args:
            agent_list: A list of thousandeyes.Agent instances.

        Returns:
            A list of integers representing all enterprise agent ids.
        """
        agent_ids = []

        for agent in agent_list:
            if agent.type == "Enterprise":
                agent_ids.append(agent.id)

        return agent_ids

    def get_se_tests(self, test_list):
        """Get a list of tests that were previously created with SolarEyes.

        Filter a full list of tests for those names starting with a specific prefix. This prefix
        is specified in the settings.ini file.

        Args:
            test_list: A list of thousandeyes.NetworkTest instances.

        Returns:
            A filtered list of thousandeyes.NetworkTest instances.
        """
        se_tests = []

        for test in test_list:
            if test.name[0:len(self.settings.te_test_prefix)] == self.settings.te_test_prefix:
                se_tests.append(test)

        return se_tests

    def get_path_nodes(self):
        """Get a list of nodes that have been marked as requiring path data.

        Returns:
            A list of node names and their IP addresses for those nodes requiring path data.
        """
        nodes = self.sw_api.query("SELECT NodeName, IPAddress FROM Orion.Nodes WHERE "
                                  "Nodes.CustomProperties.%s = TRUE" % self.settings.sw_custom_bool)

        for node in nodes['results']:
            print(node['NodeName'])
            print(node['IPAddress'])
            print("---------------")

    def check_apis(self):
        """Check the availability of both SolarWinds and ThousandEyes APIs.

        Returns:
            A boolean of True if both APIs are ready, or False if an error occurred.
        """
        if self.te_api.status() and self.sw_api.status():
            return True
        else:
            return False

    def create_test(self, test_name=None, test_server=None):
        """Create a test with defaults taken from settings.ini and from passed values.

        Args:
            test_name: The name of the test as displayed in the ThousandEyes web interface.
            test_server: The IP address that will be tested in ThousandEyes.

        Returns:
            A boolean of True if the creation was successful or False if an error occurred.
        """
        if test_name is None or test_server is None:
            raise errors.Error("[%s.%s] - You must provide a test name and server."
                               % (__name__, self.__class__.__name__))
        else:
            test = thousandeyes.NetworkTest()
            test.name = self.settings.te_test_prefix + " " + test_name
            test.server = test_server
            test.alerts_enabled = bool(self.settings.te_test_alerts)
            test.bandwidth_measurements = False
            test.mtu_measurements = True
            test.network_measurements = True
            test.bgp_measurements = True
            test.protocol = self.settings.te_test_protocol
            test.port = int(self.settings.te_test_port)
            test.interval = int(self.settings.te_test_interval)

        for agent_id in self.get_agent_ids(self.te_api.get_agents()):
            test.agents.append({"agentId": agent_id})

        if self.te_api.create_network_test(test):
            return True
        else:
            return False

    def sync(self):
        """Perform a synchronisation between SolarWinds and ThousandEyes.

        Returns:
            A boolean of True if the synchronisation was successful, or False if an error occurred.
        """
        if self.check_apis():
            # Synchronise between SolarWinds and ThousandEyes.

            #se_tests = self.get_se_tests(self.te_api.get_network_tests())

            # for se_test in se_tests:
            #     print(se_test.)

            #if self.create_test("Test", "8.8.8.8"):
                #Log test creation here. Need to pass in logger.
            self.get_path_nodes()

            return True
        else:
            return False

class SolarEyesSettings(object):
    """All external settings required by the SolarEyes class.

    Attributes:
        sw_custom_bool: A string containing the base API url such as https://api.thousandeyes.com.
        te_test_protocol: A string containing the ThousandEyes default test protocol.
        te_test_port: A string containing the ThousandEyes default test port.
        te_test_alerts: A string containing the ThousandEyes default test alert requirements.
        te_test_interval: A string containing the ThousandEyes default test interval.
        te_test_prefix: A string containing the ThousandEyes default test name prefix.
    """
    def __init__(self, sw_custom_bool=None, te_test_protocol=None, te_test_port=None,
                 te_test_alerts=None, te_test_interval=None, te_test_prefix=None):

        if sw_custom_bool is None or te_test_protocol is None or te_test_port is None \
                or te_test_alerts is None or te_test_interval is None or te_test_prefix is None:
                    raise errors.Error("[%s.%s] - You must provide all required settings to SolarEyesSettings."
                                       % (__name__, self.__class__.__name__))
        else:
            self.sw_custom_bool = sw_custom_bool
            self.te_test_protocol = te_test_protocol
            self.te_test_port = te_test_port
            self.te_test_alerts = te_test_alerts
            self.te_test_interval = te_test_interval
            self.te_test_prefix = te_test_prefix

def main():
    # Read settings file.
    config = configparser.ConfigParser()
    config.read('settings.ini')
    settings = config['Settings']

    # Configure logging.
    logging.basicConfig(level=logging.DEBUG,
                        filename=settings['log_file'], filemode=settings['log_mode'],
                        format='%(asctime)s %(levelname)s %(message)s')

    try:
        # Create various instances required by the SolarEyes constructor.
        te_api_request = thousandeyes.ApiRequest(settings['te_api'], settings['te_auth_email'], settings['te_auth_token'])
        te_api = thousandeyes.Api(te_api_request)
        sw_api = solarwinds.Api(settings['sw_api'], settings['sw_username'], settings['sw_password'])
        se_settings = SolarEyesSettings(settings['sw_custom_bool'],
                                        settings['te_test_protocol'],
                                        settings['te_test_port'],
                                        settings['te_test_alerts'],
                                        settings['te_test_interval'],
                                        settings['te_test_prefix'])

        # Pass the above API and settings instances to the constructor of SolarEyes.
        solareyes = SolarEyes(te_api, sw_api, se_settings)
    except errors.Error as error:
        logging.error(error)
    else:
        try:
            # Run synchronisation.
            if solareyes.sync():
                logging.info("[%s] - Synchronisation successful." % __name__)
            else:
                logging.error("[%s] - There was a problem during synchronisation."
                              % __name__)
        except errors.Error as error:
            logging.error(error)

# Only execute the main script if passed to an interpreter. Do not execute if imported.
if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>#Name: solvers0D.py
#Author: Will Barnes
#Purpose: Implement several basic ODE methods

#Import needed modules
import numpy as np

class Solvers0D(object):

    def __init__(self,func,func_params,rka_err=1e-6,safety_1=0.9,safety_2=1.1,safety_3=4.0):
        """Initialize 0D Solvers class."""

        if type(func).__name__ != 'function':
            raise TypeError("Input function is not of type function.")
        else:
            self.func = func
        self.func_params = func_params
        #save adaptive stepper parameters
        self.rka_err = rka_err
        self.safety_1 = safety_1
        self.safety_2 = safety_2
        self.safety_3 = safety_3

    def euler_solve(self,state,time,tau):
        """Use first order Euler method to step equation forward in time."""

        #Step forward in time
        derivs = self.func(state,time,tau,**self.func_params)

        #Update parameters
        new_state = np.zeros(len(state))
        if len(derivs) != len(state):
            raise ValueError("Dimension mismatch between state and derivs vectors.")
        for i in range(len(derivs)):
            new_state[i] = state[i] + derivs[i]*tau

        return new_state

    def rk4_solve(self,state,time,tau):
        """Use fourth-order Runge-Kutta solver to step equation forward in time."""

        #Initialize different time and timestep lists
        tau_temp = [.5*tau]*3 + [tau]
        time_temp = [time] + [time+.5*tau]*2 + [time+tau]

        initial_state = state
        f_rk = []

        #Calculate RK f_1--f_4 functions
        for i in range(len(tau_temp)):
            f_rk.append(self.func(state,time_temp[i],tau_temp[i],**self.func_params))

            if len(f_rk[-1]) != len(state):
                raise ValueError("Dimension mismatch between state and derivs vectors.")

            for j in range(len(f_rk[-1])):
                new_state[j] = state[j] + f_rk[-1][j]*tau_temp[i]

            state = new_state

        #Update state vector
        for i in range(len(initial_state)):
            new_state[i] = initial_state[i] + 1.0/6.0*tau*(f_rk[0][i] + f_rk[3][i] + 2.0*(f_rk[1][i] + f_rk[2][i]))

        return new_state

    #def adaptive_timestep(self,)<commit_msg>vectorizing RK solver<commit_after>#Name: solvers0D.py
#Author: Will Barnes
#Purpose: Implement several basic ODE methods

#Import needed modules
import numpy as np

class Solvers0D(object):

    def __init__(self,func,func_params,rka_err=1e-6,safety_1=0.9,safety_2=1.1,safety_3=4.0):
        """Initialize 0D Solvers class."""

        if type(func).__name__ != 'function':
            raise TypeError("Input function is not of type function.")
        else:
            self.func = func
        self.func_params = func_params
        #save adaptive stepper parameters
        self.rka_err = rka_err
        self.safety_1 = safety_1
        self.safety_2 = safety_2
        self.safety_3 = safety_3

    def euler_solve(self,state,time,tau):
        """Use first order Euler method to step equation forward in time."""

        #Step forward in time
        derivs = self.func(state,time,tau,**self.func_params)

        #Update parameters
        new_state = np.zeros(len(state))
        if len(derivs) != len(state):
            raise ValueError("Dimension mismatch between state and derivs vectors.")
        for i in range(len(derivs)):
            new_state[i] = state[i] + derivs[i]*tau

        return new_state

    def rk4_solve(self,state,time,tau):
        """Use fourth-order Runge-Kutta solver to step equation forward in time."""

        #calculate f functions
        #   cast as numpy array just in case this has not been done in the function
        f1 = np.array(self.func(state,time,tau,**self.func_params))
        f2 = np.array(self.func(state+tau/2*f1,time+tau/2,tau/2,**self.func_params))
        f3 = np.array(self.func(state+tau/2*f2,time+tau/2,tau/2,**self.func_params))
        f4 = np.array(self.func(state+tau*f3,time+tau,tau,**self.func_params))

        return np.array(state) + tau/6.0*(f1 + f2 + f3 + f4)

    #def adaptive_timestep(self,)<|endoftext|>


<commit_before>"""
South's fake ORM; lets you not have to write SQL inside migrations.
Roughly emulates the real Django ORM, to a point.
"""

import inspect
import datetime

from django.db import models
from django.db.models.loading import cache
from django.core.exceptions import ImproperlyConfigured

from south.db import db
from south.utils import ask_for_it_by_name
from south.hacks import hacks

class ModelsLocals(object):

    """
    Custom dictionary-like class to be locals();
    falls back to lowercase search for items that don't exist
    (because we store model names as lowercase).
    """

    def __init__(self, data):
        self.data = data

    def __getitem__(self, key):
        try:
            return self.data[key]
        except KeyError:
            return self.data[key.lower()]

# Stores already-created ORMs.
_orm_cache = {}

def FakeORM(*args):
    """
    Creates a Fake Django ORM.
    This is actually a memoised constructor; the real class is _FakeORM.
    """
    if not args in _orm_cache:
        _orm_cache[args] = _FakeORM(*args)
    return _orm_cache[args]

class LazyFakeORM(object):
    """
    In addition to memoising the ORM call, this function lazily generates them
    for a Migration class. Assign the result of this to (for example)
    .orm, and as soon as .orm is accessed the ORM will be created.
    """

    def __init__(self, *args):
        self._args = args
        self.orm = None

    def __get__(self, obj, type=None):
        if not self.orm:
            self.orm = FakeORM(*self._args)
        return self.orm

class _FakeORM(object):

    """
    Simulates the Django ORM at some point in time,
    using a frozen definition on the Migration class.
    """

    def __init__(self, cls, app):
        self.default_app = app
        self.cls = cls
        # Try loading the models off the migration class; default to no models.
        self.models = {}
        try:
            self.models_source = cls.models
        except AttributeError:
            return

        # Start a 'new' AppCache
        hacks.clear_app_cache()

        # Now, make each model's data into a FakeModel
        # We first make entries for each model that are just its name
        # This allows us to have circular model dependency loops
        model_names = []
        for name, data in self.models_source.items():
            # Make sure there's some kind of Meta
            if "Meta" not in data:
                data['Meta'] = {}
            try:
                app_label, model_name = name.split(".", 1)
            except ValueError:
                app_label = self.default_app
                model_name = name
                name = "%s.%s" % (app_label, model_name)

            # If there's an object_name in the Meta, use it and remove it
            if "object_name" in data['Meta']:
                model_name = data['Meta']['object_name']
                del data['Meta']['object_name']

            name = name.lower()
            self.models[name] = name
            model_names.append((name, app_label, model_name, data))

        for name, app_label, model_name, data in model_names:
            self.models[name] = self.make_model(app_label, model_name, data)

        # And perform the second run to iron out any circular/backwards depends.
        self.retry_failed_fields()

        # Force evaluation of relations on the models now
        for model in self.models.values():
            model._meta.get_all_field_names()

        # Reset AppCache
        hacks.unclear_app_cache()

    def __iter__(self):
        return iter(self.models.values())

    def __getattr__(self, key):
        fullname = (self.default_app+"."+key).lower()
        try:
            return self.models[fullname]
        except KeyError:
            raise AttributeError("The model '%s' from the app '%s' is not available in this migration." % (key, self.default_app))

    def __getitem__(self, key):
        # Detect if they asked for a field on a model or not.
        if ":" in key:
            key, fname = key.split(":")
        else:
            fname = None
        # Now, try getting the model
        key = key.lower()
        try:
            model = self.models[key]
        except KeyError:
            try:
                app, model = key.split(".", 1)
            except ValueError:
                raise KeyError("The model '%s' is not in appname.modelname format." % key)
            else:
                raise KeyError("The model '%s' from the app '%s' is not available in this migration." % (model, app))
        # If they asked for a field, get it.
        if fname:
            return model._meta.get_field_by_name(fname)[0]
        else:
            return model

    def eval_in_context(self, code, app, extra_imports={}):
        "Evaluates the given code in the context of the migration file."

        # Drag in the migration module's locals (hopefully including models.py)
        fake_locals = dict(inspect.getmodule(self.cls).__dict__)

        # Remove all models from that (i.e. from modern models.py), to stop pollution
        for key, value in fake_locals.items():
            if isinstance(value, type) and issubclass(value, models.Model) and hasattr(value, "_meta"):
                del fake_locals[key]

        # We add our models into the locals for the eval
        fake_locals.update(dict([
            (name.split(".")[-1], model)
            for name, model in self.models.items()
        ]))

        # Make sure the ones for this app override.
        fake_locals.update(dict([
            (name.split(".")[-1], model)
            for name, model in self.models.items()
            if name.split(".")[0] == app
        ]))

        # Ourselves as orm, to allow non-fail cross-app referencing
        fake_locals['orm'] = self

        # And a fake _ function
        fake_locals['_'] = lambda x: x

        # Datetime; there should be no datetime direct accesses
        fake_locals['datetime'] = datetime

        # Now, go through the requested imports and import them.
        for name, value in extra_imports.items():
            # First, try getting it out of locals.
            parts = value.split(".")
            try:
                obj = fake_locals[parts[0]]
                for part in parts[1:]:
                    obj = getattr(obj, part)
            except (KeyError, AttributeError):
                pass
            else:
                fake_locals[name] = obj
                continue
            # OK, try to import it directly
            try:
                fake_locals[name] = ask_for_it_by_name(value)
            except ImportError:
                if name == "SouthFieldClass":
                    raise ValueError("Cannot import the required field '%s'" % value)
                else:
                    print "WARNING: Cannot import '%s'" % value

        # Use ModelsLocals to make lookups work right for CapitalisedModels
        fake_locals = ModelsLocals(fake_locals)

        return eval(code, globals(), fake_locals)

    def make_meta(self, app, model, data, stub=False):
        "Makes a Meta class out of a dict of eval-able arguments."
        results = {'app_label': app}
        for key, code in data.items():
            # Some things we never want to use.
            if key in ["_bases"]:
                continue
            # Some things we don't want with stubs.
            if stub and key in ["order_with_respect_to"]:
                continue
            # OK, add it.
            try:
                results[key] = self.eval_in_context(code, app)
            except (NameError, AttributeError), e:
                raise ValueError("Cannot successfully create meta field '%s' for model '%s.%s': %s." % (
                    key, app, model, e
                ))
        return type("Meta", tuple(), results)

    def make_model(self, app, name, data):
        "Makes a Model class out of the given app name, model name and pickled data."

        # Extract any bases out of Meta
        if "_bases" in data['Meta']:
            bases = data['Meta']['_bases']
        else:
            bases = ['django.db.models.Model']

        # Turn the Meta dict into a basic class
        meta = self.make_meta(app, name, data['Meta'], data.get("_stub", False))

        failed_fields = {}
        fields = {}
        stub = False

        # Now, make some fields!
        for fname, params in data.items():
            # If it's the stub marker, ignore it.
            if fname == "_stub":
                stub = bool(params)
                continue
            elif fname == "Meta":
                continue
            elif not params:
                raise ValueError("Field '%s' on model '%s.%s' has no definition." % (fname, app, name))
            elif isinstance(params, (str, unicode)):
                # It's a premade definition string! Let's hope it works...
                code = params
                extra_imports = {}
            else:
                # If there's only one parameter (backwards compat), make it 3.
                if len(params) == 1:
                    params = (params[0], [], {})
                # There should be 3 parameters. Code is a tuple of (code, what-to-import)
                if len(params) == 3:
                    code = "SouthFieldClass(%s)" % ", ".join(
                        params[1] +
                        ["%s=%s" % (n, v) for n, v in params[2].items()]
                    )
                    extra_imports = {"SouthFieldClass": params[0]}
                else:
                    raise ValueError("Field '%s' on model '%s.%s' has a weird definition length (should be 1 or 3 items)." % (fname, app, name))

            try:
                # Execute it in a probably-correct context.
                field = self.eval_in_context(code, app, extra_imports)
            except (NameError, AttributeError, AssertionError, KeyError):
                # It might rely on other models being around. Add it to the
                # model for the second pass.
                failed_fields[fname] = (code, extra_imports)
            else:
                fields[fname] = field

        # Find the app in the Django core, and get its module
        more_kwds = {}
        try:
            app_module = models.get_app(app)
            more_kwds['__module__'] = app_module.__name__
        except ImproperlyConfigured:
            # The app this belonged to has vanished, but thankfully we can still
            # make a mock model, so ignore the error.
            more_kwds['__module__'] = '_south_mock'

        more_kwds['Meta'] = meta

        # Make our model
        fields.update(more_kwds)

        model = type(
            str(name),
            tuple(map(ask_for_it_by_name, bases)),
            fields,
        )

        # If this is a stub model, change Objects to a whiny class
        if stub:
            model.objects = WhinyManager()
            # Also, make sure they can't instantiate it
            model.__init__ = whiny_method
        else:
            model.objects = NoDryRunManager(model.objects)

        if failed_fields:
            model._failed_fields = failed_fields

        return model

    def retry_failed_fields(self):
        "Tries to re-evaluate the _failed_fields for each model."
        for modelkey, model in self.models.items():
            app, modelname = modelkey.split(".", 1)
            if hasattr(model, "_failed_fields"):
                for fname, (code, extra_imports) in model._failed_fields.items():
                    try:
                        field = self.eval_in_context(code, app, extra_imports)
                    except (NameError, AttributeError, AssertionError, KeyError), e:
                        # It's failed again. Complain.
                        raise ValueError("Cannot successfully create field '%s' for model '%s': %s." % (
                            fname, modelname, e
                        ))
                    else:
                        # Startup that field.
                        model.add_to_class(fname, field)

class WhinyManager(object):
    "A fake manager that whines whenever you try to touch it. For stub models."

    def __getattr__(self, key):
        raise AttributeError("You cannot use items from a stub model.")

class NoDryRunManager(object):
    """
    A manager that always proxies through to the real manager,
    unless a dry run is in progress.
    """

    def __init__(self, real):
        self.real = real

    def __getattr__(self, name):
        if db.dry_run:
            raise AttributeError("You are in a dry run, and cannot access the ORM.\nWrap ORM sections in 'if not db.dry_run:', or if the whole migration is only a data migration, set no_dry_run = True on the Migration class.")
        return getattr(self.real, name)

def whiny_method(*a, **kw):
    raise ValueError("You cannot instantiate a stub model.")
<commit_msg>Fix the hopefully last one of those 'str has no _meta' errors<commit_after>"""
South's fake ORM; lets you not have to write SQL inside migrations.
Roughly emulates the real Django ORM, to a point.
"""

import inspect
import datetime

from django.db import models
from django.db.models.loading import cache
from django.core.exceptions import ImproperlyConfigured

from south.db import db
from south.utils import ask_for_it_by_name
from south.hacks import hacks

class ModelsLocals(object):

    """
    Custom dictionary-like class to be locals();
    falls back to lowercase search for items that don't exist
    (because we store model names as lowercase).
    """

    def __init__(self, data):
        self.data = data

    def __getitem__(self, key):
        try:
            return self.data[key]
        except KeyError:
            return self.data[key.lower()]

# Stores already-created ORMs.
_orm_cache = {}

def FakeORM(*args):
    """
    Creates a Fake Django ORM.
    This is actually a memoised constructor; the real class is _FakeORM.
    """
    if not args in _orm_cache:
        _orm_cache[args] = _FakeORM(*args)
    return _orm_cache[args]

class LazyFakeORM(object):
    """
    In addition to memoising the ORM call, this function lazily generates them
    for a Migration class. Assign the result of this to (for example)
    .orm, and as soon as .orm is accessed the ORM will be created.
    """

    def __init__(self, *args):
        self._args = args
        self.orm = None

    def __get__(self, obj, type=None):
        if not self.orm:
            self.orm = FakeORM(*self._args)
        return self.orm

class _FakeORM(object):

    """
    Simulates the Django ORM at some point in time,
    using a frozen definition on the Migration class.
    """

    def __init__(self, cls, app):
        self.default_app = app
        self.cls = cls
        # Try loading the models off the migration class; default to no models.
        self.models = {}
        try:
            self.models_source = cls.models
        except AttributeError:
            return

        # Start a 'new' AppCache
        hacks.clear_app_cache()

        # Now, make each model's data into a FakeModel
        # We first make entries for each model that are just its name
        # This allows us to have circular model dependency loops
        model_names = []
        for name, data in self.models_source.items():
            # Make sure there's some kind of Meta
            if "Meta" not in data:
                data['Meta'] = {}
            try:
                app_label, model_name = name.split(".", 1)
            except ValueError:
                app_label = self.default_app
                model_name = name

            # If there's an object_name in the Meta, use it and remove it
            if "object_name" in data['Meta']:
                model_name = data['Meta']['object_name']
                del data['Meta']['object_name']

            name = "%s.%s" % (app_label, model_name)
            self.models[name.lower()] = name
            model_names.append((name.lower(), app_label, model_name, data))

        for name, app_label, model_name, data in model_names:
            self.models[name] = self.make_model(app_label, model_name, data)

        # And perform the second run to iron out any circular/backwards depends.
        self.retry_failed_fields()

        # Force evaluation of relations on the models now
        for model in self.models.values():
            model._meta.get_all_field_names()

        # Reset AppCache
        hacks.unclear_app_cache()

    def __iter__(self):
        return iter(self.models.values())

    def __getattr__(self, key):
        fullname = (self.default_app+"."+key).lower()
        try:
            return self.models[fullname]
        except KeyError:
            raise AttributeError("The model '%s' from the app '%s' is not available in this migration." % (key, self.default_app))

    def __getitem__(self, key):
        # Detect if they asked for a field on a model or not.
        if ":" in key:
            key, fname = key.split(":")
        else:
            fname = None
        # Now, try getting the model
        key = key.lower()
        try:
            model = self.models[key]
        except KeyError:
            try:
                app, model = key.split(".", 1)
            except ValueError:
                raise KeyError("The model '%s' is not in appname.modelname format." % key)
            else:
                raise KeyError("The model '%s' from the app '%s' is not available in this migration." % (model, app))
        # If they asked for a field, get it.
        if fname:
            return model._meta.get_field_by_name(fname)[0]
        else:
            return model

    def eval_in_context(self, code, app, extra_imports={}):
        "Evaluates the given code in the context of the migration file."

        # Drag in the migration module's locals (hopefully including models.py)
        fake_locals = dict(inspect.getmodule(self.cls).__dict__)

        # Remove all models from that (i.e. from modern models.py), to stop pollution
        for key, value in fake_locals.items():
            if isinstance(value, type) and issubclass(value, models.Model) and hasattr(value, "_meta"):
                del fake_locals[key]

        # We add our models into the locals for the eval
        fake_locals.update(dict([
            (name.split(".")[-1], model)
            for name, model in self.models.items()
        ]))

        # Make sure the ones for this app override.
        fake_locals.update(dict([
            (name.split(".")[-1], model)
            for name, model in self.models.items()
            if name.split(".")[0] == app
        ]))

        # Ourselves as orm, to allow non-fail cross-app referencing
        fake_locals['orm'] = self

        # And a fake _ function
        fake_locals['_'] = lambda x: x

        # Datetime; there should be no datetime direct accesses
        fake_locals['datetime'] = datetime

        # Now, go through the requested imports and import them.
        for name, value in extra_imports.items():
            # First, try getting it out of locals.
            parts = value.split(".")
            try:
                obj = fake_locals[parts[0]]
                for part in parts[1:]:
                    obj = getattr(obj, part)
            except (KeyError, AttributeError):
                pass
            else:
                fake_locals[name] = obj
                continue
            # OK, try to import it directly
            try:
                fake_locals[name] = ask_for_it_by_name(value)
            except ImportError:
                if name == "SouthFieldClass":
                    raise ValueError("Cannot import the required field '%s'" % value)
                else:
                    print "WARNING: Cannot import '%s'" % value

        # Use ModelsLocals to make lookups work right for CapitalisedModels
        fake_locals = ModelsLocals(fake_locals)

        return eval(code, globals(), fake_locals)

    def make_meta(self, app, model, data, stub=False):
        "Makes a Meta class out of a dict of eval-able arguments."
        results = {'app_label': app}
        for key, code in data.items():
            # Some things we never want to use.
            if key in ["_bases"]:
                continue
            # Some things we don't want with stubs.
            if stub and key in ["order_with_respect_to"]:
                continue
            # OK, add it.
            try:
                results[key] = self.eval_in_context(code, app)
            except (NameError, AttributeError), e:
                raise ValueError("Cannot successfully create meta field '%s' for model '%s.%s': %s." % (
                    key, app, model, e
                ))
        return type("Meta", tuple(), results)

    def make_model(self, app, name, data):
        "Makes a Model class out of the given app name, model name and pickled data."

        # Extract any bases out of Meta
        if "_bases" in data['Meta']:
            bases = data['Meta']['_bases']
        else:
            bases = ['django.db.models.Model']

        # Turn the Meta dict into a basic class
        meta = self.make_meta(app, name, data['Meta'], data.get("_stub", False))

        failed_fields = {}
        fields = {}
        stub = False

        # Now, make some fields!
        for fname, params in data.items():
            # If it's the stub marker, ignore it.
            if fname == "_stub":
                stub = bool(params)
                continue
            elif fname == "Meta":
                continue
            elif not params:
                raise ValueError("Field '%s' on model '%s.%s' has no definition." % (fname, app, name))
            elif isinstance(params, (str, unicode)):
                # It's a premade definition string! Let's hope it works...
                code = params
                extra_imports = {}
            else:
                # If there's only one parameter (backwards compat), make it 3.
                if len(params) == 1:
                    params = (params[0], [], {})
                # There should be 3 parameters. Code is a tuple of (code, what-to-import)
                if len(params) == 3:
                    code = "SouthFieldClass(%s)" % ", ".join(
                        params[1] +
                        ["%s=%s" % (n, v) for n, v in params[2].items()]
                    )
                    extra_imports = {"SouthFieldClass": params[0]}
                else:
                    raise ValueError("Field '%s' on model '%s.%s' has a weird definition length (should be 1 or 3 items)." % (fname, app, name))

            try:
                # Execute it in a probably-correct context.
                field = self.eval_in_context(code, app, extra_imports)
            except (NameError, AttributeError, AssertionError, KeyError):
                # It might rely on other models being around. Add it to the
                # model for the second pass.
                failed_fields[fname] = (code, extra_imports)
            else:
                fields[fname] = field

        # Find the app in the Django core, and get its module
        more_kwds = {}
        try:
            app_module = models.get_app(app)
            more_kwds['__module__'] = app_module.__name__
        except ImproperlyConfigured:
            # The app this belonged to has vanished, but thankfully we can still
            # make a mock model, so ignore the error.
            more_kwds['__module__'] = '_south_mock'

        more_kwds['Meta'] = meta

        # Make our model
        fields.update(more_kwds)

        model = type(
            str(name),
            tuple(map(ask_for_it_by_name, bases)),
            fields,
        )

        # If this is a stub model, change Objects to a whiny class
        if stub:
            model.objects = WhinyManager()
            # Also, make sure they can't instantiate it
            model.__init__ = whiny_method
        else:
            model.objects = NoDryRunManager(model.objects)

        if failed_fields:
            model._failed_fields = failed_fields

        return model

    def retry_failed_fields(self):
        "Tries to re-evaluate the _failed_fields for each model."
        for modelkey, model in self.models.items():
            app, modelname = modelkey.split(".", 1)
            if hasattr(model, "_failed_fields"):
                for fname, (code, extra_imports) in model._failed_fields.items():
                    try:
                        field = self.eval_in_context(code, app, extra_imports)
                    except (NameError, AttributeError, AssertionError, KeyError), e:
                        # It's failed again. Complain.
                        raise ValueError("Cannot successfully create field '%s' for model '%s': %s." % (
                            fname, modelname, e
                        ))
                    else:
                        # Startup that field.
                        model.add_to_class(fname, field)

class WhinyManager(object):
    "A fake manager that whines whenever you try to touch it. For stub models."

    def __getattr__(self, key):
        raise AttributeError("You cannot use items from a stub model.")

class NoDryRunManager(object):
    """
    A manager that always proxies through to the real manager,
    unless a dry run is in progress.
    """

    def __init__(self, real):
        self.real = real

    def __getattr__(self, name):
        if db.dry_run:
            raise AttributeError("You are in a dry run, and cannot access the ORM.\nWrap ORM sections in 'if not db.dry_run:', or if the whole migration is only a data migration, set no_dry_run = True on the Migration class.")
        return getattr(self.real, name)

def whiny_method(*a, **kw):
    raise ValueError("You cannot instantiate a stub model.")
<|endoftext|>