# Copyright 2015 Rackspace
# All Rights Reserved.
#
#    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 abc
import collections
import os
import shlex
import subprocess

import six

from timid import steps
from timid import utils

@six.add_metaclass(abc.ABCMeta)
class SpecialVariable(object):
    """
    A mix-in class for "special" environment variables.  These are
    variables which can be represented as lists or sets.
    """

    def __init__(self, env, var, sep=os.pathsep, value=utils.unset):
        """
        Initialize a ``SpecialVariable`` instance.

        :param env: The ``Environment`` containing the variable.
        :param var: The name of the environment variable.
        :param sep: The separator used.  Defaults to the value of
                    ``os.pathsep``.
        """

        # Save the basic data we need to function
        self._env = env
        self._var = var
        self._sep = sep

        # Select the current value
        if value is utils.unset:
            value = env._data.get(var)

        # Set up the value cache
        self._update(value)

    def __str__(self):
        """
        Obtain a string form of the ``EnvListVariable`` instance.

        :returns: The string form of the variable.
        """

        return self._env._data.get(self._var) or ''

    def __repr__(self):
        """
        Obtain a representation of the ``EnvListVariable`` instance.

        :returns: A representation of the variable.
        """

        return repr(self._value)

    def __len__(self):
        """
        Obtain the length of an ``EnvListVariable`` instance.

        :returns: The number of elements in the variable.
        """

        return len(self._value)

    def _rebuild(self):
        """
        Helper method used to alter the variable's value in a linked
        ``Environment`` instance when a change is made to the
        ``SpecialVariable`` instance.
        """

        self._env._data[self._var] = self._sep.join(self._value)

    def _update(self, value):
        """
        Alert method used when the variable's value in an ``Environment``
        instance has been altered.

        :param value: The new value to set the value to.
        """

        if not value:
            # Empty string or None, so use the empty value
            value = self._coerce()
        elif isinstance(value, six.string_types):
            # It's a string, so split and coerce to the right type
            value = self._coerce(value.split(self._sep))
        else:
            # It's already a compatible type, so just coerce it
            value = self._coerce(value)

        # Set the value
        self._value = value

    @abc.abstractproperty
    def _type(self):
        """
        The abstract type the value should be represented as, e.g.,
        ``collections.Set`` or ``collections.Sequence``.
        """

        pass  # pragma: no cover

    @abc.abstractproperty
    def _coerce(self):
        """
        The underlying type the value should be represented as, e.g.,
        ``set`` or ``list``.  Used to convert a value into the
        appropriate type.
        """

        pass  # pragma: no cover

class ListVariable(SpecialVariable, collections.MutableSequence):
    """
    Represent the value of an environment variable which is list form.
    This is used for representing the PATH environment variable, and
    may be used for other similar variables as well.
    """

    _type = collections.Sequence
    _coerce = list

    def __getitem__(self, idx):
        """
        Retrieve an item from a ``ListVariable`` instance.

        :param idx: The index or index-like object (e.g., slice) to
                    retrieve.

        :returns: The value at the designated index.
        """

        return self._value[idx]

    def __setitem__(self, idx, value):
        """
        Set an item on the ``ListVariable`` instance.

        :param idx: The index or index-like object (e.g., slice) to
                    set.
        :param value: The value to set it to.
        """

        self._value[idx] = value
        self._rebuild()

    def __delitem__(self, idx):
        """
        Delete an item from the ``ListVariable`` instance.

        :param idx: The index or index-like object (e.g., slice) to
                    delete.
        """

        del self._value[idx]
        self._rebuild()

    def insert(self, idx, value):
        """
        Inserts a value in the ``ListVariable`` at an appropriate index.

        :param idx: The index before which to insert the new value.
        :param value: The value to insert.
        """

        self._value.insert(idx, value)
        self._rebuild()

class SetVariable(SpecialVariable, collections.MutableSet):
    """
    Represent the value of an environment variable which is set form.
    This is used for representing the TIMID_SENSITIVE environment
    variable, and may be used for other similar variables as well.
    """

    _type = collections.Set
    _coerce = set

    def __contains__(self, item):
        """
        Check if an item is a member of the ``SetVariable`` instance.

        :param item: The item to check.

        :returns: A ``True`` value if the item is a member, ``False``
                  otherwise.
        """

        return item in self._value

    def __iter__(self):
        """
        Iterate over the items in the ``SetVariable`` instance.

        :returns: An iterator over the items.
        """

        return iter(self._value)

    def add(self, item):
        """
        Add a new item to the ``SetVariable`` instance.

        :param item: The item to add.
        """

        self._value.add(item)
        self._rebuild()

    def discard(self, item):
        """
        Discard an item from the ``SetVariable`` instance.

        :param item: The item to discard.
        """

        self._value.discard(item)
        self._rebuild()

class Environment(utils.SensitiveDict):
    """
    Represents a calling environment for a process.  This contains
    environment variables and the working directory that processes
    should be executed with.
    """

    def __init__(self, environ=None, sensitive=None, cwd=None):
        """
        Initialize a new ``Environment`` instance.

        :param environ: An optional dictionary containing the initial
                        set of environment variables.  If omitted,
                        ``os.environ`` will be used.
        :param sensitive: An optional set of "sensitive" variables,
                          environment variables whose values should
                          not be printed out.
        :param cwd: The working directory processes should be executed
                    in.  If not given, the current working directory
                    will be used.  If a relative path, it will be
                    interpreted relative to the current working
                    directory.
        """

        # Select the starting environment
        environ = environ or os.environ.copy()

        # Build the 'sensitive' set
        sensitive_var = SetVariable(self, 'TIMID_SENSITIVE',
                                    value=environ.get('TIMID_SENSITIVE'))
        if sensitive:
            sensitive_var |= sensitive

        # Initialize the SensitiveDict
        super(Environment, self).__init__(environ, sensitive_var)

        # Track special variables
        self._special = {
            'PATH': ListVariable(self, 'PATH'),
            'TIMID_SENSITIVE': sensitive_var,
        }

        # Initialize the working directory
        self._cwd = utils.canonicalize_path(os.getcwd(), cwd or os.curdir)

    def __getitem__(self, name):
        """
        Retrieve the value of a variable from the environment.  If the
        variable has been declared as a list-like variable, the return
        value will be a list of strings; otherwise, it will be a
        string.

        :param name: The name of the variable to retrieve.

        :returns: The value of the designated variable.
        """

        # Raise a KeyError if needed
        if name not in self._data:
            raise KeyError(name)

        # If it's special, return the special instance
        if name in self._special:
            return self._special[name]

        return self._data[name]

    def __setitem__(self, name, value):
        """
        Set the value of a variable in the environment.  If the variable
        has been declared as a list-like variable, an intelligent
        conversion of the value will be performed.

        :param name: The name of the variable to set.
        :param value: The value to set.  Must be a string or, for
                      list-like variables, a sequence.  A ``None``
                      value is interpreted as a delete.
        """

        # Get the special variable, if any
        special = self._special.get(name)

        # Handle the delete case first
        if value is None:
            self._data.pop(name, None)
            if special:
                special._update(None)

        # Is the value a string type?
        elif isinstance(value, six.string_types):
            self._data[name] = value
            if special:
                special._update(value)

        # Is the value a type compatible with a special variable?
        elif (special and
              (isinstance(value, collections.Iterable) or
               isinstance(value, special._type))):
            # Let the special variable handle the change directly
            special._update(value)

            # Rebuild the value into a string
            special._rebuild()

        # Unable to handle the data type
        else:
            raise ValueError('invalid value %r' % value)

    def __delitem__(self, name):
        """
        Delete the value of a variable from the environment.

        :param name: The name of the variable to delete.
        """

        # Will raise the KeyError if the variable doesn't exist
        del self._data[name]

        # Notify special variables
        if name in self._special:
            self._special[name]._update(None)

    def copy(self):
        """
        Retrieve a copy of the Environment.  Note that this is a shallow
        copy.
        """

        return self.__class__(self._data.copy(), self._sensitive.copy(),
                              self._cwd)

    def _declare_special(self, name, sep, klass):
        """
        Declare an environment variable as a special variable.  This can
        be used even if the environment variable is not present.

        :param name: The name of the environment variable that should
                     be considered special.
        :param sep: The separator to be used.
        :param klass: The subclass of ``SpecialVariable`` used to
                      represent the variable.
        """

        # First, has it already been declared?
        if name in self._special:
            special = self._special[name]
            if not isinstance(special, klass) or sep != special._sep:
                raise ValueError('variable %s already declared as %s '
                                 'with separator "%s"' %
                                 (name, special.__class__.__name__,
                                  special._sep))

        # OK, it's new; declare it
        else:
            self._special[name] = klass(self, name, sep)

    def declare_list(self, name, sep=os.pathsep):
        """
        Declare an environment variable as a list-like special variable.
        This can be used even if the environment variable is not
        present.

        :param name: The name of the environment variable that should
                     be considered list-like.
        :param sep: The separator to be used.  Defaults to the value
                    of ``os.pathsep``.
        """

        self._declare_special(name, sep, ListVariable)

    def declare_set(self, name, sep=os.pathsep):
        """
        Declare an environment variable as a set-like special variable.
        This can be used even if the environment variable is not
        present.

        :param name: The name of the environment variable that should
                     be considered set-like.
        :param sep: The separator to be used.  Defaults to the value
                    of ``os.pathsep``.
        """

        self._declare_special(name, sep, SetVariable)

    def call(self, args, **kwargs):
        """
        A thin wrapper around ``subprocess.Popen``.  Takes the same
        options as ``subprocess.Popen``, with the exception of the
        ``cwd``, and ``env`` parameters, which come from the
        ``Environment`` instance.  Note that if the sole positional
        argument is a string, it will be converted into a sequence
        using the ``shlex.split()`` function.
        """

        # Convert string args into a sequence
        if isinstance(args, six.string_types):
            args = shlex.split(args)

        # Substitute cwd and env
        kwargs['cwd'] = self._cwd
        kwargs['env'] = self._data

        # Set a default for close_fds
        kwargs.setdefault('close_fds', True)

        return subprocess.Popen(args, **kwargs)

    @property
    def cwd(self):
        """
        Retrieve the current working directory processes will be executed
        in.
        """

        return self._cwd

    @cwd.setter
    def cwd(self, value):
        """
        Change the working directory that processes should be executed in.

        :param value: The new path to change to.  If relative, will be
                      interpreted relative to the current working
                      directory.
        """

        self._cwd = utils.canonicalize_path(self._cwd, value)

class EnvironmentAction(steps.SensitiveDictAction):
    """
    An action for updating and otherwise modifying the execution
    environment.  The base usage is::

        - env:
            set:
              ENV_VAR: value
            unset:
            - OTHER_VAR
            sensitive:
            - SENS_VAR
            files:
            - fname1.yaml
            - fname2.yaml

    This action would set the environment variable named "ENV_VAR" to
    the value "value", unset the environment variable "OTHER_VAR",
    mark the environment variable "SENS_VAR" as a sensitive variable,
    and read the files "fname1.yaml" and "fname2.yaml" from the
    directory containing the file specifying the action.

    Note that if an environment variable is present under both the
    "set" and "unset" keys, the "set" will take precedence.  Also note
    that variable file reading is performed before any other
    operations.
    """

    # Schema for validating the configuration; this contains tweaks
    # specific to environment variables
    schema = {
        'type': 'object',
        'properties': {
            'set': {
                'type': 'object',
                'additionalProperties': {'type': 'string'},
            },
            'unset': {
                'type': 'array',
                'items': {'type': 'string'},
            },
            'sensitive': {
                'type': 'array',
                'items': {'type': 'string'},
            },
            'files': {
                'type': 'array',
                'items': {'type': 'string'},
            },
        },
        'additionalProperties': False,
    }

    # The name of the context attribute affected by this action
    context_attr = 'environment'

class DirectoryAction(steps.Action):
    """
    An action for changing the working directory.  The base usage is::

        - chdir: path/to/change/to

    This action would cause a change to the indicated directory,
    relative to the current working directory.
    """

    # Schema for validating the configuration
    schema = {'type': 'string'}

    def __init__(self, ctxt, name, config, step_addr):
        """
        Initialize a ``DirectoryAction`` instance.

        :param ctxt: The context object.
        :param name: The name of the action.
        :param config: The configuration for the action.  This may be
                       a scalar value (e.g., "run: command"), a list,
                       or a dictionary.  If the configuration provided
                       is invalid for the action, a ``ConfigError``
                       should be raised.
        :param step_addr: The address of the step in the test
                          configuration.  Should be passed to the
                          ``ConfigError``.
        """

        # Perform superclass initialization
        super(DirectoryAction, self).__init__(ctxt, name, config, step_addr)

        # Save the directory to switch to
        self.target_dir = ctxt.template(config)

    def __call__(self, ctxt):
        """
        Invoke the action.  This changes to the directory specified by the
        configuration.

        :param ctxt: The context object.

        :returns: A ``StepResult`` object.
        """

        # Change to the target directory
        ctxt.environment.cwd = self.target_dir(ctxt)

        # We're done!
        return steps.StepResult(state=steps.SUCCESS)

class RunAction(steps.Action):
    """
    An action for invoking a shell command.  The base usage is::

        - run: ./command.py arg1 arg2 arg3

    Note that template substitution is performed *before* splitting
    the arguments up into a list.  (This splitting is performed using
    ``shlex.split()``, so shell quoting is honored.)  In particular,
    this means that any given template substitution that renders into
    a string containing spaces may result into multiple arguments
    passed to the command, unless shell quoting is used.  If more
    control over arguments is desired, use the more advanced list
    form::

        - run:
          - ./command.py
          - arg1
          - arg2
          - arg3

    In this form, no shell syntax quoting is honored.
    """

    # Schema for validating the configuration
    schema = {
        'oneOf': [
            {'type': 'string'},
            {
                'type': 'array',
                'items': {'type': 'string'},
            },
        ],
    }

    def __init__(self, ctxt, name, config, step_addr):
        """
        Initialize a ``RunAction`` instance.

        :param ctxt: The context object.
        :param name: The name of the action.
        :param config: The configuration for the action.  This may be
                       a scalar value (e.g., "run: command"), a list,
                       or a dictionary.  If the configuration provided
                       is invalid for the action, a ``ConfigError``
                       should be raised.
        :param step_addr: The address of the step in the test
                          configuration.  Should be passed to the
                          ``ConfigError``.
        """

        # Perform superclass initialization
        super(RunAction, self).__init__(ctxt, name, config, step_addr)

        # Build up the correct command
        if isinstance(config, six.string_types):
            self.command = ctxt.template(config)
        else:
            self.command = [ctxt.template(arg) for arg in config]

    def __call__(self, ctxt):
        """
        Invoke the action.  This executes the command specified in the
        configuration.

        :param ctxt: The context object.

        :returns: A ``StepResult`` object.
        """

        # First, do the correct splitting/rendering
        if isinstance(self.command, list):
            args = [arg(ctxt) for arg in self.command]
        else:
            args = shlex.split(self.command(ctxt))

        # Invoke the command
        subproc = ctxt.environment.call(args)

        # All done...
        return steps.StepResult(returncode=subproc.wait())


import json
import math
import random

import jsonschema

from pyrallelsa import Problem

# https://github.com/perrygeo/simanneal/blob/master/LICENSE.txt
# https://github.com/perrygeo/simanneal/blob/master/examples/salesman.py

def distance(a, b):
    """Calculates distance between two latitude-longitude coordinates."""
    R = 3963  # radius of Earth (miles)
    lat1, lon1 = math.radians(a[0]), math.radians(a[1])
    lat2, lon2 = math.radians(b[0]), math.radians(b[1])
    return math.acos(math.sin(lat1) * math.sin(lat2) +
                     math.cos(lat1) * math.cos(lat2) * math.cos(lon1 - lon2)) * R

def get_distance_matrix(cities):
    # create a distance matrix
    distance_matrix = {}
    for ka, va in cities.items():
        for kb, vb in cities.items():
            if kb == ka:
                distance_matrix["{},{}".format(ka, kb)] = 0.0
            else:
                distance_matrix["{},{}".format(ka, kb)] = distance(va, vb)
    return distance_matrix

class TSPProblem(Problem):
    """Traveling Salesman Problem Annealer

    :param State|None state: state of the current annealer process; if this
        is None, a new random state will be generated
    """

    # We have this because a distance matrix grows exponentially
    #  and your RAM will die
    MAX_CITIES_FOR_DISTANCE_MATRIX = 150

    def __init__(self, state=None, **problem_data):
        self._validate_problem_data(problem_data)

        self.cities = problem_data["cities"]
        self.locked_range = int(problem_data.get("locked_range", 2))

        start_city = problem_data.get("start_city")
        self.start_city = self.cities[0] if start_city is None else start_city
        assert self.start_city in self.cities

        if not problem_data.get("updates_enabled", False):
            self.update = lambda *args, **kwargs: None
        if len(self.cities) < self.MAX_CITIES_FOR_DISTANCE_MATRIX:
            self.distance_matrix = get_distance_matrix(self.cities)
            print("Using distance matrix")
        else:
            self.distance_matrix = None

        if state is None:
            state = self.cities.keys()
            random.shuffle(state)
        self.copy_strategy = "slice"
        super(TSPProblem, self).__init__(state)  # important!

    def _validate_problem_data(self, problem_data):
        schema = {
            "type": "object",
            "properties": {
                "cities": {"type": "object"},
                "start_city": {"type": "string"},
                "updates_enabled": {"type": "boolean"},
                "locked_range": {"type": "number"}
            },
            "required": ["cities"]
        }
        jsonschema.validate(problem_data, schema)

    def move(self, state=None):
        """Swaps two cities in the route.

        :type state: TSPState
        """
        state = self.state if state is None else state
        route = state
        a = random.randint(self.locked_range, len(route) - 1)
        b = random.randint(self.locked_range, len(route) - 1)
        route[a], route[b] = route[b], route[a]

    def energy(self, state=None):
        """Calculates the length of the route."""
        state = self.state if state is None else state
        route = state
        e = 0
        if self.distance_matrix:
            for i in range(len(route)):
                e += self.distance_matrix["{},{}".format(route[i-1], route[i])]
        else:
            for i in range(len(route)):
                e += distance(self.cities[route[i-1]], self.cities[route[i]])
        return e

    @classmethod
    def divide(cls, divisions, problem_data):
        """divide

        :type problem_data: dict
        """
        tspp = TSPProblem(**problem_data)

        def routes_for_subgroup(cs):
            for city in cs:
                if city == tspp.start_city:
                    continue
                cities = tspp.cities.keys()
                cities.remove(tspp.start_city)
                cities.remove(city)
                random.shuffle(cities)
                route = [tspp.start_city, city] + cities
                assert len(set(route)) == len(route)
                assert len(route) == len(tspp.cities)
                yield json.dumps(route)

        if divisions:
            chunk_size = int(math.ceil(len(tspp.cities) / divisions))
        else:
            chunk_size = 1
        for subgroup in chunks(tspp.cities.keys(), chunk_size):
            routes = list(routes_for_subgroup(subgroup))
            if routes:
                yield routes

def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i+n]


"""A TaskRecord backend using sqlite3

Authors:

* Min RK
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

import json
import os
import cPickle as pickle
from datetime import datetime

try:
    import sqlite3
except ImportError:
    sqlite3 = None

from zmq.eventloop import ioloop

from IPython.utils.traitlets import Unicode, Instance, List, Dict
from .dictdb import BaseDB
from IPython.utils.jsonutil import date_default, extract_dates, squash_dates

#-----------------------------------------------------------------------------
# SQLite operators, adapters, and converters
#-----------------------------------------------------------------------------

try:
    buffer
except NameError:
    # py3k
    buffer = memoryview

operators = {
 '$lt' : "<",
 '$gt' : ">",
 # null is handled weird with ==,!=
 '$eq' : "=",
 '$ne' : "!=",
 '$lte': "<=",
 '$gte': ">=",
 '$in' : ('=', ' OR '),
 '$nin': ('!=', ' AND '),
 # '$all': None,
 # '$mod': None,
 # '$exists' : None
}
null_operators = {
'=' : "IS NULL",
'!=' : "IS NOT NULL",
}

def _adapt_dict(d):
    return json.dumps(d, default=date_default)

def _convert_dict(ds):
    if ds is None:
        return ds
    else:
        if isinstance(ds, bytes):
            # If I understand the sqlite doc correctly, this will always be utf8
            ds = ds.decode('utf8')
        return extract_dates(json.loads(ds))

def _adapt_bufs(bufs):
    # this is *horrible*
    # copy buffers into single list and pickle it:
    if bufs and isinstance(bufs[0], (bytes, buffer)):
        return sqlite3.Binary(pickle.dumps(map(bytes, bufs),-1))
    elif bufs:
        return bufs
    else:
        return None

def _convert_bufs(bs):
    if bs is None:
        return []
    else:
        return pickle.loads(bytes(bs))

#-----------------------------------------------------------------------------
# SQLiteDB class
#-----------------------------------------------------------------------------

class SQLiteDB(BaseDB):
    """SQLite3 TaskRecord backend."""

    filename = Unicode('tasks.db', config=True,
        help="""The filename of the sqlite task database. [default: 'tasks.db']""")
    location = Unicode('', config=True,
        help="""The directory containing the sqlite task database.  The default
        is to use the cluster_dir location.""")
    table = Unicode("", config=True,
        help="""The SQLite Table to use for storing tasks for this session. If unspecified,
        a new table will be created with the Hub's IDENT.  Specifying the table will result
        in tasks from previous sessions being available via Clients' db_query and
        get_result methods.""")

    if sqlite3 is not None:
        _db = Instance('sqlite3.Connection')
    else:
        _db = None
    # the ordered list of column names
    _keys = List(['msg_id' ,
            'header' ,
            'content',
            'buffers',
            'submitted',
            'client_uuid' ,
            'engine_uuid' ,
            'started',
            'completed',
            'resubmitted',
            'received',
            'result_header' ,
            'result_content' ,
            'result_buffers' ,
            'queue' ,
            'pyin' ,
            'pyout',
            'pyerr',
            'stdout',
            'stderr',
        ])
    # sqlite datatypes for checking that db is current format
    _types = Dict({'msg_id' : 'text' ,
            'header' : 'dict text',
            'content' : 'dict text',
            'buffers' : 'bufs blob',
            'submitted' : 'timestamp',
            'client_uuid' : 'text',
            'engine_uuid' : 'text',
            'started' : 'timestamp',
            'completed' : 'timestamp',
            'resubmitted' : 'text',
            'received' : 'timestamp',
            'result_header' : 'dict text',
            'result_content' : 'dict text',
            'result_buffers' : 'bufs blob',
            'queue' : 'text',
            'pyin' : 'text',
            'pyout' : 'text',
            'pyerr' : 'text',
            'stdout' : 'text',
            'stderr' : 'text',
        })

    def __init__(self, **kwargs):
        super(SQLiteDB, self).__init__(**kwargs)
        if sqlite3 is None:
            raise ImportError("SQLiteDB requires sqlite3")
        if not self.table:
            # use session, and prefix _, since starting with # is illegal
            self.table = '_'+self.session.replace('-','_')
        if not self.location:
            # get current profile
            from IPython.core.application import BaseIPythonApplication
            if BaseIPythonApplication.initialized():
                app = BaseIPythonApplication.instance()
                if app.profile_dir is not None:
                    self.location = app.profile_dir.location
                else:
                    self.location = u'.'
            else:
                self.location = u'.'
        self._init_db()

        # register db commit as 2s periodic callback
        # to prevent clogging pipes
        # assumes we are being run in a zmq ioloop app
        loop = ioloop.IOLoop.instance()
        pc = ioloop.PeriodicCallback(self._db.commit, 2000, loop)
        pc.start()

    def _defaults(self, keys=None):
        """create an empty record"""
        d = {}
        keys = self._keys if keys is None else keys
        for key in keys:
            d[key] = None
        return d

    def _check_table(self):
        """Ensure that an incorrect table doesn't exist

        If a bad (old) table does exist, return False
        """
        cursor = self._db.execute("PRAGMA table_info(%s)"%self.table)
        lines = cursor.fetchall()
        if not lines:
            # table does not exist
            return True
        types = {}
        keys = []
        for line in lines:
            keys.append(line[1])
            types[line[1]] = line[2]
        if self._keys != keys:
            # key mismatch
            self.log.warn('keys mismatch')
            return False
        for key in self._keys:
            if types[key] != self._types[key]:
                self.log.warn(
                    'type mismatch: %s: %s != %s'%(key,types[key],self._types[key])
                )
                return False
        return True

    def _init_db(self):
        """Connect to the database and get new session number."""
        # register adapters
        sqlite3.register_adapter(dict, _adapt_dict)
        sqlite3.register_converter('dict', _convert_dict)
        sqlite3.register_adapter(list, _adapt_bufs)
        sqlite3.register_converter('bufs', _convert_bufs)
        # connect to the db
        dbfile = os.path.join(self.location, self.filename)
        self._db = sqlite3.connect(dbfile, detect_types=sqlite3.PARSE_DECLTYPES,
            # isolation_level = None)#,
             cached_statements=64)
        # print dir(self._db)
        first_table = previous_table = self.table
        i=0
        while not self._check_table():
            i+=1
            self.table = first_table+'_%i'%i
            self.log.warn(
                "Table %s exists and doesn't match db format, trying %s"%
                (previous_table, self.table)
            )
            previous_table = self.table

        self._db.execute("""CREATE TABLE IF NOT EXISTS %s
                (msg_id text PRIMARY KEY,
                header dict text,
                content dict text,
                buffers bufs blob,
                submitted timestamp,
                client_uuid text,
                engine_uuid text,
                started timestamp,
                completed timestamp,
                resubmitted text,
                received timestamp,
                result_header dict text,
                result_content dict text,
                result_buffers bufs blob,
                queue text,
                pyin text,
                pyout text,
                pyerr text,
                stdout text,
                stderr text)
                """%self.table)
        self._db.commit()

    def _dict_to_list(self, d):
        """turn a mongodb-style record dict into a list."""

        return [ d[key] for key in self._keys ]

    def _list_to_dict(self, line, keys=None):
        """Inverse of dict_to_list"""
        keys = self._keys if keys is None else keys
        d = self._defaults(keys)
        for key,value in zip(keys, line):
            d[key] = value

        return d

    def _render_expression(self, check):
        """Turn a mongodb-style search dict into an SQL query."""
        expressions = []
        args = []

        skeys = set(check.keys())
        skeys.difference_update(set(self._keys))
        skeys.difference_update(set(['buffers', 'result_buffers']))
        if skeys:
            raise KeyError("Illegal testing key(s): %s"%skeys)

        for name,sub_check in check.iteritems():
            if isinstance(sub_check, dict):
                for test,value in sub_check.iteritems():
                    try:
                        op = operators[test]
                    except KeyError:
                        raise KeyError("Unsupported operator: %r"%test)
                    if isinstance(op, tuple):
                        op, join = op

                    if value is None and op in null_operators:
                        expr = "%s %s" % (name, null_operators[op])
                    else:
                        expr = "%s %s ?"%(name, op)
                        if isinstance(value, (tuple,list)):
                            if op in null_operators and any([v is None for v in value]):
                                # equality tests don't work with NULL
                                raise ValueError("Cannot use %r test with NULL values on SQLite backend"%test)
                            expr = '( %s )'%( join.join([expr]*len(value)) )
                            args.extend(value)
                        else:
                            args.append(value)
                    expressions.append(expr)
            else:
                # it's an equality check
                if sub_check is None:
                    expressions.append("%s IS NULL" % name)
                else:
                    expressions.append("%s = ?"%name)
                    args.append(sub_check)

        expr = " AND ".join(expressions)
        return expr, args

    def add_record(self, msg_id, rec):
        """Add a new Task Record, by msg_id."""
        d = self._defaults()
        d.update(rec)
        d['msg_id'] = msg_id
        line = self._dict_to_list(d)
        tups = '(%s)'%(','.join(['?']*len(line)))
        self._db.execute("INSERT INTO %s VALUES %s"%(self.table, tups), line)
        # self._db.commit()

    def get_record(self, msg_id):
        """Get a specific Task Record, by msg_id."""
        cursor = self._db.execute("""SELECT * FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
        line = cursor.fetchone()
        if line is None:
            raise KeyError("No such msg: %r"%msg_id)
        return self._list_to_dict(line)

    def update_record(self, msg_id, rec):
        """Update the data in an existing record."""
        query = "UPDATE %s SET "%self.table
        sets = []
        keys = sorted(rec.keys())
        values = []
        for key in keys:
            sets.append('%s = ?'%key)
            values.append(rec[key])
        query += ', '.join(sets)
        query += ' WHERE msg_id == ?'
        values.append(msg_id)
        self._db.execute(query, values)
        # self._db.commit()

    def drop_record(self, msg_id):
        """Remove a record from the DB."""
        self._db.execute("""DELETE FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
        # self._db.commit()

    def drop_matching_records(self, check):
        """Remove a record from the DB."""
        expr,args = self._render_expression(check)
        query = "DELETE FROM %s WHERE %s"%(self.table, expr)
        self._db.execute(query,args)
        # self._db.commit()

    def find_records(self, check, keys=None):
        """Find records matching a query dict, optionally extracting subset of keys.

        Returns list of matching records.

        Parameters
        ----------

        check: dict
            mongodb-style query argument
        keys: list of strs [optional]
            if specified, the subset of keys to extract.  msg_id will *always* be
            included.
        """
        if keys:
            bad_keys = [ key for key in keys if key not in self._keys ]
            if bad_keys:
                raise KeyError("Bad record key(s): %s"%bad_keys)

        if keys:
            # ensure msg_id is present and first:
            if 'msg_id' in keys:
                keys.remove('msg_id')
            keys.insert(0, 'msg_id')
            req = ', '.join(keys)
        else:
            req = '*'
        expr,args = self._render_expression(check)
        query = """SELECT %s FROM %s WHERE %s"""%(req, self.table, expr)
        cursor = self._db.execute(query, args)
        matches = cursor.fetchall()
        records = []
        for line in matches:
            rec = self._list_to_dict(line, keys)
            records.append(rec)
        return records

    def get_history(self):
        """get all msg_ids, ordered by time submitted."""
        query = """SELECT msg_id FROM %s ORDER by submitted ASC"""%self.table
        cursor = self._db.execute(query)
        # will be a list of length 1 tuples
        return [ tup[0] for tup in cursor.fetchall()]

__all__ = ['SQLiteDB']


# encoding: utf-8
"""
Utilities for warnings.  Shoudn't we just use the built in warnings module.
"""

#-----------------------------------------------------------------------------
#  Copyright (C) 2008-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import sys

from IPython.utils import io

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------

def warn(msg,level=2,exit_val=1):
    """Standard warning printer. Gives formatting consistency.

    Output is sent to io.stderr (sys.stderr by default).

    Options:

    -level(2): allows finer control:
      0 -> Do nothing, dummy function.
      1 -> Print message.
      2 -> Print 'WARNING:' + message. (Default level).
      3 -> Print 'ERROR:' + message.
      4 -> Print 'FATAL ERROR:' + message and trigger a sys.exit(exit_val).

    -exit_val (1): exit value returned by sys.exit() for a level 4
    warning. Ignored for all other levels."""

    if level>0:
        header = ['','','WARNING: ','ERROR: ','FATAL ERROR: ']
        io.stderr.write('%s%s' % (header[level],msg))
        if level == 4:
            print >> io.stderr,'Exiting.\n'
            sys.exit(exit_val)

def info(msg):
    """Equivalent to warn(msg,level=1)."""

    warn(msg,level=1)

def error(msg):
    """Equivalent to warn(msg,level=3)."""

    warn(msg,level=3)

def fatal(msg,exit_val=1):
    """Equivalent to warn(msg,exit_val=exit_val,level=4)."""

    warn(msg,exit_val=exit_val,level=4)


import anyconfig
import os

from . import utils

class Reader():
    """Config Reader. Reads, validate and add default values to YAML"""

    def parse(self, config_file=None, specs=None, default_file=None):
        """Read a config_file, check the validity with a JSON Schema as specs
        and get default values from default_file if asked.

        All parameters are optionnal.

        If there is no config_file defined, read the venv base
        dir and try to get config/app.yml.

        If no specs, don't validate anything.

        If no default_file, don't merge with default values."""

        self._config_exists(config_file)
        self._specs_exists(specs)

        self.loaded_config = anyconfig.load(self.config_file, ac_parser='yaml')

        if default_file is not None:
            self._merge_default(default_file)

        if self.specs is None:
            return self.loaded_config

        self._validate()

        return self.loaded_config

    def _validate(self) -> None:
        (rc, err) = anyconfig.validate(self.loaded_config, anyconfig.load(self.specs, ac_parser='yaml'))

        if err != '':
            raise ValueError('Your config is not valid: {}'.format(err))

    def _config_exists(self, config_file: str) -> str:
        self.config_file = utils.get_venv_basedir() + '/config/app.yml'
        if config_file is not None:
            self.config_file = config_file

        if not os.path.isfile(self.config_file):
            raise IOError('Missing config file: "{}" does not exist'.format(self.config_file))

    def _merge_default(self, default_file: str) -> None:
        if not os.path.isfile(default_file):
            raise IOError("Your default ({}) does not exist".format(default_file))

        default_config = anyconfig.load(default_file)
        anyconfig.merge(default_config, self.loaded_config)

        self.loaded_config = default_config

    def _specs_exists(self, specs: str) -> None:
        self.specs = specs
        if self.specs is None:
            return

        # self.specs = os.path.abspath(os.path.dirname(__file__) + '/static/configspec.ini')
        if not os.path.isfile(self.specs):
            raise IOError("Your specs ({}) does not exist".format(self.specs))


from __future__ import unicode_literals

from decimal import Decimal

from django import template
from django.db.models import Model
from django.forms import ModelForm

from django.conf import settings
from django.core.urlresolvers import reverse

from django_baseline import html

register = template.Library()

@register.simple_tag
def table(rows):
    '''
    Output a simple table with several columns.
    '''

    output = '
'

    for row in rows:
        output += '
'
        for column in row:
            output += '{s}	'.format(s=column)
        output += '
'

    output += '
'

    return output

@register.simple_tag
def link(url, text='', classes='', target='', get="", **kwargs):
    '''
    Output a link tag.
    '''

    if not (url.startswith('http') or url.startswith('/')):
        # Handle additional reverse args.
        urlargs = {}

        for arg, val in kwargs.items():
            if arg[:4] == "url_":
                urlargs[arg[4:]] = val

        url = reverse(url, kwargs=urlargs)
        if get:
            url += '?' + get

    return html.tag('a', text or url, {
        'class': classes, 'target': target, 'href': url})

@register.simple_tag
def jsfile(url):
    '''
    Output a script tag to a js file.
    '''

    if not url.startswith('http://') and not url[:1] == '/':
        #add media_url for relative paths
        url = settings.STATIC_URL + url

    return ''.format(
        src=url)

@register.simple_tag
def cssfile(url):
    '''
    Output a link tag to a css stylesheet.
    '''

    if not url.startswith('http://') and not url[:1] == '/':
        #add media_url for relative paths
        url = settings.STATIC_URL + url

    return '<link href="{src}" rel="stylesheet">'.format(src=url)

@register.simple_tag
def img(url, alt='', classes='', style=''):
    '''
    Image tag helper.
    '''

    if not url.startswith('http://') and not url[:1] == '/':
        #add media_url for relative paths
        url = settings.STATIC_URL + url

    attr = {
        'class': classes,
        'alt': alt,
        'style': style,
        'src': url
    }

    return html.tag('img', '', attr)

def valid_numeric(arg):
    if isinstance(arg, (int, float, Decimal)):
        return arg
    try:
        return int(arg)
    except ValueError:
        return float(arg)

@register.filter
def sub(value, arg):
    """Subtract the arg from the value."""
    try:
        return valid_numeric(value) - valid_numeric(arg)
    except (ValueError, TypeError):
        try:
            return value - arg
        except Exception:
            return ''
sub.is_safe = False

@register.filter
def mul(value, arg):
    """Multiply the arg with the value."""
    try:
        return valid_numeric(value) * valid_numeric(arg)
    except (ValueError, TypeError):
        try:
            return value * arg
        except Exception:
            return ''
mul.is_safe = False

@register.filter
def div(value, arg):
    """Divide the arg by the value."""
    try:
        return valid_numeric(value) / valid_numeric(arg)
    except (ValueError, TypeError):
        try:
            return value / arg
        except Exception:
            return ''
div.is_safe = False

@register.filter(name='abs')
def absolute(value):
    """Return the absolute value."""
    try:
        return abs(valid_numeric(value))
    except (ValueError, TypeError):
        try:
            return abs(value)
        except Exception:
            return ''
absolute.is_safe = False

@register.filter
def mod(value, arg):
    """Return the modulo value."""
    try:
        return valid_numeric(value) % valid_numeric(arg)
    except (ValueError, TypeError):
        try:
            return value % arg
        except Exception:
            return ''
mod.is_safe = False

# Model related.

@register.filter
def model_verbose(obj, capitalize=True):
    """
    Return the verbose name of a model.
    The obj argument can be either a Model instance, or a ModelForm instance.
    This allows to retrieve the verbose name of the model of a ModelForm
    easily, without adding extra context vars.
    """

    if isinstance(obj, ModelForm):
        name = obj._meta.model._meta.verbose_name
    elif isinstance(obj, Model):
        name = obj._meta.verbose_name
    else:
        raise Exception('Unhandled type: ' + type(obj))

    return name.capitalize() if capitalize else name
model_verbose.is_safe = False

@register.filter
def user_can_edit(obj, user):
    """
    If a model implements the user_can_edit method,
    this filter returns the result of the method.
    """

    return obj.user_can_edit(user) if hasattr(obj, "user_can_edit") else None


# Python library imports.
#

# Django imports.
from django.db import models

def extendManager(mixinClass):
	'''
	Use as a class decorator to add extra methods to your model manager.
	Example usage:

		class Article(django.db.models.Model):
			published = models.DateTimeField()
			...

			@extendManager
			class objects(object):
				def getPublished(self):
					return self.filter(published__lte = django.utils.timezone.now()).order_by('-published')

		...

		publishedArticles = Article.objects.getPublished()
	'''

	class MixinManager(models.Manager, mixinClass):
		class MixinQuerySet(models.query.QuerySet, mixinClass):
			pass
		def get_queryset(self):
			return self.MixinQuerySet(self.model, using = self._db)

	return MixinManager()


"""
This module contains the main logic to run all the logic defined
"""

import commands
import xml.etree.ElementTree as ET
from optparse import OptionParser
from os import environ
from tabulate import tabulate
from subprocess import call
from storage import *
from configuration import *

TAGS_XML_DIR = environ["HOME"] + "/.mypy/myhelp/tags.xml"

tree = ET.parse(environ["HOME"] + "/.mypy/myhelp/tags.xml")

root = tree.getroot()
roottags = root.find('tags')
tags = roottags.iter('tag')
rootfiles = root.find('files')
files = rootfiles.iter('file')

def run():
    """Main method where all logic is defined"""

    config_option_help="'show' - displays configured options, 'set [section] [name] [value]' - sets config under a section,'set [name] [value]' - sets configuration globally"

    parser = OptionParser()

    parser.add_option("-a", "--add", action="store", type="string", dest="addfile", help="adds a notes")
    parser.add_option("-c", "--config", action="store", type="string", dest="config", help=config_option_help)
    parser.add_option("-e", "--edit", action="store", type="string", dest="editfile", help="edits a notes")
    parser.add_option("-o", "--open", action="store", type="string", dest="openfile", help="opens a notes")
    parser.add_option("-r", "--remove", action="store", type="string", dest="remove", help="removes a notes")

    options, args = parser.parse_args()

    if options.config:
        if options.config == "show":
            config_option_list=''
            config_sections = config.sections()
            for section in config_sections:
                config_option_list=config_option_list+section+"\n"
                section_items =config.items(section)
                for item in section_items:
                    config_option_list=config_option_list+"    "+item[0]+"    "+item[1]+"\n"
            print config_option_list
            quit()

    def add_notes(note_name,existing_tags):
       call([editor,environ["HOME"] + "/.mypy/myhelp/notes/"+note_name+".note"])
       definedtags = raw_input("Define Tags (separated by spaces): ").split(" ")
       definedtags.append(note_name)
       print definedtags
       print existing_tags
       definedtags=list(set(definedtags)-set(existing_tags))
       print definedtags
       if len(definedtags)>0:
           modify_tags_xml(note_name,definedtags,files,rootfiles,tags,roottags,tree,TAGS_XML_DIR)

    def get_tags_from_file(note_name):
        fil = get_file_from_files(note_name)
        filetags = fil.iter('tag')
        filetaglist=[]
        for tag in filetags:
            filetaglist.append(tag.text)
        return filetaglist

    if options.addfile:
       existing_tags=[]
       if isFile(options.addfile,files):
           existing_tags=get_tags_from_file(options.addfile)
           raw_input("Note exists with tags - "+" ".join(existing_tags)+"\nDo you want to edit the notes ? [Press enter to continue]\n")
       add_notes(options.addfile,existing_tags)
       quit()

    if options.editfile:
        if isFile(options.editfile,files):
            add_notes(note_name,[])
        else:
           raw_input("Note doesn't exist.\nDo you want add note ? [Press enter to continue]")
           add_notes(note_name)

    if options.remove:
        pass

    if len(args) != 1:
        print "Please use a search term\n example : myhelp <some tag word> "
        quit()
    _key_File = "Note"
    _key_Results = "    Results                                     "
    table={_key_Results:[]}
    for tag in tags:
        if tag.attrib["value"] == args[0]:
            fileelements = tag.iter("file")
            for fileelement in fileelements:
                f = open(
                    environ["HOME"] + "/.mypy/myhelp/notes/" + fileelement.text+".note", "r")
		table[_key_Results].append(f.read()+"\r\n\tfile: ~/.mypy/myhelp/notes/"+fileelement.text+".note")
                f.close()

    print tabulate(table,headers=[],tablefmt="rst")

if __name__ == "__main__":
    run()


# encoding: utf-8
"""
Simple utility for splitting user input. This is used by both inputsplitter and
prefilter.

Authors:

* Brian Granger
* Fernando Perez
"""

#-----------------------------------------------------------------------------
#  Copyright (C) 2008-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import re
import sys

from IPython.utils import py3compat
from IPython.utils.encoding import get_stream_enc

#-----------------------------------------------------------------------------
# Main function
#-----------------------------------------------------------------------------

# RegExp for splitting line contents into pre-char//first word-method//rest.
# For clarity, each group in on one line.

# WARNING: update the regexp if the escapes in interactiveshell are changed, as
# they are hardwired in.

# Although it's not solely driven by the regex, note that:
# ,;/% only trigger if they are the first character on the line
# ! and !! trigger if they are first char(s) *or* follow an indent
# ? triggers as first or last char.

line_split = re.compile("""
             ^(\s*)               # any leading space
             ([,;/%]|!!?|\?\??)?  # escape character or characters
             \s*(%{0,2}[\w\.\*]*)     # function/method, possibly with leading %
                                  # to correctly treat things like '?%magic'
             (.*?$|$)             # rest of line
             """, re.VERBOSE)

def split_user_input(line, pattern=None):
    """Split user input into initial whitespace, escape character, function part
    and the rest.
    """
    # We need to ensure that the rest of this routine deals only with unicode
    encoding = get_stream_enc(sys.stdin, 'utf-8')
    line = py3compat.cast_unicode(line, encoding)

    if pattern is None:
        pattern = line_split
    match = pattern.match(line)
    if not match:
        # print "match failed for line '%s'" % line
        try:
            ifun, the_rest = line.split(None,1)
        except ValueError:
            # print "split failed for line '%s'" % line
            ifun, the_rest = line, u''
        pre = re.match('^(\s*)(.*)',line).groups()[0]
        esc = ""
    else:
        pre, esc, ifun, the_rest = match.groups()

    #print 'line:<%s>' % line # dbg
    #print 'pre <%s> ifun <%s> rest <%s>' % (pre,ifun.strip(),the_rest) # dbg
    return pre, esc or '', ifun.strip(), the_rest.lstrip()

class LineInfo(object):
    """A single line of input and associated info.

    Includes the following as properties:

    line
      The original, raw line

    continue_prompt
      Is this line a continuation in a sequence of multiline input?

    pre
      Any leading whitespace.

    esc
      The escape character(s) in pre or the empty string if there isn't one.
      Note that '!!' and '??' are possible values for esc. Otherwise it will
      always be a single character.

    ifun
      The 'function part', which is basically the maximal initial sequence
      of valid python identifiers and the '.' character. This is what is
      checked for alias and magic transformations, used for auto-calling,
      etc. In contrast to Python identifiers, it may start with "%" and contain
      "*".

    the_rest
      Everything else on the line.
    """
    def __init__(self, line, continue_prompt=False):
        self.line            = line
        self.continue_prompt = continue_prompt
        self.pre, self.esc, self.ifun, self.the_rest = split_user_input(line)

        self.pre_char       = self.pre.strip()
        if self.pre_char:
            self.pre_whitespace = '' # No whitespace allowd before esc chars
        else:
            self.pre_whitespace = self.pre

    def ofind(self, ip):
        """Do a full, attribute-walking lookup of the ifun in the various
        namespaces for the given IPython InteractiveShell instance.

        Return a dict with keys: {found, obj, ospace, ismagic}

        Note: can cause state changes because of calling getattr, but should
        only be run if autocall is on and if the line hasn't matched any
        other, less dangerous handlers.

        Does cache the results of the call, so can be called multiple times
        without worrying about *further* damaging state.
        """
        return ip._ofind(self.ifun)

    def __str__(self):
        return "LineInfo [%s|%s|%s|%s]" %(self.pre, self.esc, self.ifun, self.the_rest)


from django import template
from django.core import urlresolvers

register = template.Library()

@register.simple_tag(takes_context = True)
def current(context, urlName, className = 'active', **kwargs):
	'''
	Return a class name (string) if the current URL matches the route name specified in ``className``.
	If any URL keyword arguments are provided, they must be matched as well.

	:param urlName: The route name that the current URL should match. Example: 'accounts:login'.
	:param className: The string that is returned if the current URL matches the specified one.
	:param kwargs: Any extra URL keyword arguments that must also be present in the current URL for it to match.
	:returns: ``className`` if the current URL matches, or an empty string otherwise.
	'''
	matches = pathMatches(context['request'].path, urlName, **kwargs)
	return className if matches else ''

def pathMatches(path, urlName, **kwargs):
	'''
	:param path: str
	:param urlName: str
	:returns: bool.
	'''
	resolved = urlresolvers.resolve(path)

	# Different URL name => the current URL cannot match.
	resolvedName = '{r.namespace}:{r.url_name}'.format(r = resolved) if resolved.namespace else resolved.url_name
	if urlName != resolvedName:
		return False

	# If any of the current keyword args is missing or different, the URL cannot match.
	for key, value in kwargs.items():
		currentArg = resolved.kwargs.get(key)
		if currentArg is None or str(value) != str(currentArg):
			return False

	return True


"""
Overview
========

The multiprocess plugin enables you to distribute your test run among a set of
worker processes that run tests in parallel. This can speed up CPU-bound test
runs (as long as the number of work processeses is around the number of
processors or cores available), but is mainly useful for IO-bound tests that
spend most of their time waiting for data to arrive from someplace else.

.. note ::

   See :doc:`../doc_tests/test_multiprocess/multiprocess` for
   additional documentation and examples. Use of this plugin on python
   2.5 or earlier requires the multiprocessing_ module, also available
   from PyPI.

.. _multiprocessing : http://code.google.com/p/python-multiprocessing/

How tests are distributed
=========================

The ideal case would be to dispatch each test to a worker process
separately. This ideal is not attainable in all cases, however, because many
test suites depend on context (class, module or package) fixtures.

The plugin can't know (unless you tell it -- see below!) if a context fixture
can be called many times concurrently (is re-entrant), or if it can be shared
among tests running in different processes. Therefore, if a context has
fixtures, the default behavior is to dispatch the entire suite to a worker as
a unit.

Controlling distribution
^^^^^^^^^^^^^^^^^^^^^^^^

There are two context-level variables that you can use to control this default
behavior.

If a context's fixtures are re-entrant, set ``_multiprocess_can_split_ = True``
in the context, and the plugin will dispatch tests in suites bound to that
context as if the context had no fixtures. This means that the fixtures will
execute concurrently and multiple times, typically once per test.

If a context's fixtures can be shared by tests running in different processes
-- such as a package-level fixture that starts an external http server or
initializes a shared database -- then set ``_multiprocess_shared_ = True`` in
the context. These fixtures will then execute in the primary nose process, and
tests in those contexts will be individually dispatched to run in parallel.

How results are collected and reported
======================================

As each test or suite executes in a worker process, results (failures, errors,
and specially handled exceptions like SkipTest) are collected in that
process. When the worker process finishes, it returns results to the main
nose process. There, any progress output is printed (dots!), and the
results from the test run are combined into a consolidated result
set. When results have been received for all dispatched tests, or all
workers have died, the result summary is output as normal.

Beware!
=======

Not all test suites will benefit from, or even operate correctly using, this
plugin. For example, CPU-bound tests will run more slowly if you don't have
multiple processors. There are also some differences in plugin
interactions and behaviors due to the way in which tests are dispatched and
loaded. In general, test loading under this plugin operates as if it were
always in directed mode instead of discovered mode. For instance, doctests
in test modules will always be found when using this plugin with the doctest
plugin.

But the biggest issue you will face is probably concurrency. Unless you
have kept your tests as religiously pure unit tests, with no side-effects, no
ordering issues, and no external dependencies, chances are you will experience
odd, intermittent and unexplainable failures and errors when using this
plugin. This doesn't necessarily mean the plugin is broken; it may mean that
your test suite is not safe for concurrency.

New Features in 1.1.0
=====================

* functions generated by test generators are now added to the worker queue
  making them multi-threaded.
* fixed timeout functionality, now functions will be terminated with a
  TimedOutException exception when they exceed their execution time. The
  worker processes are not terminated.
* added ``--process-restartworker`` option to restart workers once they are
  done, this helps control memory usage. Sometimes memory leaks can accumulate
  making long runs very difficult.
* added global _instantiate_plugins to configure which plugins are started
  on the worker processes.

"""

import logging
import os
import sys
import time
import traceback
import unittest
import pickle
import signal
import nose.case
from nose.core import TextTestRunner
from nose import failure
from nose import loader
from nose.plugins.base import Plugin
from nose.pyversion import bytes_
from nose.result import TextTestResult
from nose.suite import ContextSuite
from nose.util import test_address
try:
    # 2.7+
    from unittest.runner import _WritelnDecorator
except ImportError:
    from unittest import _WritelnDecorator
from Queue import Empty
from warnings import warn
try:
    from cStringIO import StringIO
except ImportError:
    import StringIO

# this is a list of plugin classes that will be checked for and created inside
# each worker process
_instantiate_plugins = None

log = logging.getLogger(__name__)

Process = Queue = Pool = Event = Value = Array = None

# have to inherit KeyboardInterrupt to it will interrupt process properly
class TimedOutException(KeyboardInterrupt):
    def __init__(self, value = "Timed Out"):
        self.value = value
    def __str__(self):
        return repr(self.value)

def _import_mp():
    global Process, Queue, Pool, Event, Value, Array
    try:
        from multiprocessing import Manager, Process
        #prevent the server process created in the manager which holds Python
        #objects and allows other processes to manipulate them using proxies
        #to interrupt on SIGINT (keyboardinterrupt) so that the communication
        #channel between subprocesses and main process is still usable after
        #ctrl+C is received in the main process.
        old=signal.signal(signal.SIGINT, signal.SIG_IGN)
        m = Manager()
        #reset it back so main process will receive a KeyboardInterrupt
        #exception on ctrl+c
        signal.signal(signal.SIGINT, old)
        Queue, Pool, Event, Value, Array = (
                m.Queue, m.Pool, m.Event, m.Value, m.Array
        )
    except ImportError:
        warn("multiprocessing module is not available, multiprocess plugin "
             "cannot be used", RuntimeWarning)

class TestLet:
    def __init__(self, case):
        try:
            self._id = case.id()
        except AttributeError:
            pass
        self._short_description = case.shortDescription()
        self._str = str(case)

    def id(self):
        return self._id

    def shortDescription(self):
        return self._short_description

    def __str__(self):
        return self._str

class MultiProcess(Plugin):
    """
    Run tests in multiple processes. Requires processing module.
    """
    score = 1000
    status = {}

    def options(self, parser, env):
        """
        Register command-line options.
        """
        parser.add_option("--processes", action="store",
                          default=env.get('NOSE_PROCESSES', 0),
                          dest="multiprocess_workers",
                          metavar="NUM",
                          help="Spread test run among this many processes. "
                          "Set a number equal to the number of processors "
                          "or cores in your machine for best results. "
                          "[NOSE_PROCESSES]")
        parser.add_option("--process-timeout", action="store",
                          default=env.get('NOSE_PROCESS_TIMEOUT', 10),
                          dest="multiprocess_timeout",
                          metavar="SECONDS",
                          help="Set timeout for return of results from each "
                          "test runner process. [NOSE_PROCESS_TIMEOUT]")
        parser.add_option("--process-restartworker", action="store_true",
                          default=env.get('NOSE_PROCESS_RESTARTWORKER', False),
                          dest="multiprocess_restartworker",
                          help="If set, will restart each worker process once"
                          " their tests are done, this helps control memory "
                          "leaks from killing the system. "
                          "[NOSE_PROCESS_RESTARTWORKER]")

    def configure(self, options, config):
        """
        Configure plugin.
        """
        try:
            self.status.pop('active')
        except KeyError:
            pass
        if not hasattr(options, 'multiprocess_workers'):
            self.enabled = False
            return
        # don't start inside of a worker process
        if config.worker:
            return
        self.config = config
        try:
            workers = int(options.multiprocess_workers)
        except (TypeError, ValueError):
            workers = 0
        if workers:
            _import_mp()
            if Process is None:
                self.enabled = False
                return
            self.enabled = True
            self.config.multiprocess_workers = workers
            t = float(options.multiprocess_timeout)
            self.config.multiprocess_timeout = t
            r = int(options.multiprocess_restartworker)
            self.config.multiprocess_restartworker = r
            self.status['active'] = True

    def prepareTestLoader(self, loader):
        """Remember loader class so MultiProcessTestRunner can instantiate
        the right loader.
        """
        self.loaderClass = loader.__class__

    def prepareTestRunner(self, runner):
        """Replace test runner with MultiProcessTestRunner.
        """
        # replace with our runner class
        return MultiProcessTestRunner(stream=runner.stream,
                                      verbosity=self.config.verbosity,
                                      config=self.config,
                                      loaderClass=self.loaderClass)

def signalhandler(sig, frame):
    raise TimedOutException()

class MultiProcessTestRunner(TextTestRunner):
    waitkilltime = 5.0 # max time to wait to terminate a process that does not
                       # respond to SIGILL
    def __init__(self, **kw):
        self.loaderClass = kw.pop('loaderClass', loader.defaultTestLoader)
        super(MultiProcessTestRunner, self).__init__(**kw)

    def collect(self, test, testQueue, tasks, to_teardown, result):
        # dispatch and collect results
        # put indexes only on queue because tests aren't picklable
        for case in self.nextBatch(test):
            log.debug("Next batch %s (%s)", case, type(case))
            if (isinstance(case, nose.case.Test) and
                isinstance(case.test, failure.Failure)):
                log.debug("Case is a Failure")
                case(result) # run here to capture the failure
                continue
            # handle shared fixtures
            if isinstance(case, ContextSuite) and case.context is failure.Failure:
                log.debug("Case is a Failure")
                case(result) # run here to capture the failure
                continue
            elif isinstance(case, ContextSuite) and self.sharedFixtures(case):
                log.debug("%s has shared fixtures", case)
                try:
                    case.setUp()
                except (KeyboardInterrupt, SystemExit):
                    raise
                except:
                    log.debug("%s setup failed", sys.exc_info())
                    result.addError(case, sys.exc_info())
                else:
                    to_teardown.append(case)
                    if case.factory:
                        ancestors=case.factory.context.get(case, [])
                        for an in ancestors[:2]:
                            #log.debug('reset ancestor %s', an)
                            if getattr(an, '_multiprocess_shared_', False):
                                an._multiprocess_can_split_=True
                            #an._multiprocess_shared_=False
                    self.collect(case, testQueue, tasks, to_teardown, result)

            else:
                test_addr = self.addtask(testQueue,tasks,case)
                log.debug("Queued test %s (%s) to %s",
                          len(tasks), test_addr, testQueue)

    def startProcess(self, iworker, testQueue, resultQueue, shouldStop, result):
        currentaddr = Value('c',bytes_(''))
        currentstart = Value('d',time.time())
        keyboardCaught = Event()
        p = Process(target=runner,
                   args=(iworker, testQueue,
                         resultQueue,
                         currentaddr,
                         currentstart,
                         keyboardCaught,
                         shouldStop,
                         self.loaderClass,
                         result.__class__,
                         pickle.dumps(self.config)))
        p.currentaddr = currentaddr
        p.currentstart = currentstart
        p.keyboardCaught = keyboardCaught
        old = signal.signal(signal.SIGILL, signalhandler)
        p.start()
        signal.signal(signal.SIGILL, old)
        return p

    def run(self, test):
        """
        Execute the test (which may be a test suite). If the test is a suite,
        distribute it out among as many processes as have been configured, at
        as fine a level as is possible given the context fixtures defined in
        the suite or any sub-suites.

        """
        log.debug("%s.run(%s) (%s)", self, test, os.getpid())
        wrapper = self.config.plugins.prepareTest(test)
        if wrapper is not None:
            test = wrapper

        # plugins can decorate or capture the output stream
        wrapped = self.config.plugins.setOutputStream(self.stream)
        if wrapped is not None:
            self.stream = wrapped

        testQueue = Queue()
        resultQueue = Queue()
        tasks = []
        completed = []
        workers = []
        to_teardown = []
        shouldStop = Event()

        result = self._makeResult()
        start = time.time()

        self.collect(test, testQueue, tasks, to_teardown, result)

        log.debug("Starting %s workers", self.config.multiprocess_workers)
        for i in range(self.config.multiprocess_workers):
            p = self.startProcess(i, testQueue, resultQueue, shouldStop, result)
            workers.append(p)
            log.debug("Started worker process %s", i+1)

        total_tasks = len(tasks)
        # need to keep track of the next time to check for timeouts in case
        # more than one process times out at the same time.
        nexttimeout=self.config.multiprocess_timeout
        thrownError = None

        try:
            while tasks:
                log.debug("Waiting for results (%s/%s tasks), next timeout=%.3fs",
                          len(completed), total_tasks,nexttimeout)
                try:
                    iworker, addr, newtask_addrs, batch_result = resultQueue.get(
                                                            timeout=nexttimeout)
                    log.debug('Results received for worker %d, %s, new tasks: %d',
                              iworker,addr,len(newtask_addrs))
                    try:
                        try:
                            tasks.remove(addr)
                        except ValueError:
                            log.warn('worker %s failed to remove from tasks: %s',
                                     iworker,addr)
                        total_tasks += len(newtask_addrs)
                        tasks.extend(newtask_addrs)
                    except KeyError:
                        log.debug("Got result for unknown task? %s", addr)
                        log.debug("current: %s",str(list(tasks)[0]))
                    else:
                        completed.append([addr,batch_result])
                    self.consolidate(result, batch_result)
                    if (self.config.stopOnError
                        and not result.wasSuccessful()):
                        # set the stop condition
                        shouldStop.set()
                        break
                    if self.config.multiprocess_restartworker:
                        log.debug('joining worker %s',iworker)
                        # wait for working, but not that important if worker
                        # cannot be joined in fact, for workers that add to
                        # testQueue, they will not terminate until all their
                        # items are read
                        workers[iworker].join(timeout=1)
                        if not shouldStop.is_set() and not testQueue.empty():
                            log.debug('starting new process on worker %s',iworker)
                            workers[iworker] = self.startProcess(iworker, testQueue, resultQueue, shouldStop, result)
                except Empty:
                    log.debug("Timed out with %s tasks pending "
                              "(empty testQueue=%r): %s",
                              len(tasks),testQueue.empty(),str(tasks))
                    any_alive = False
                    for iworker, w in enumerate(workers):
                        if w.is_alive():
                            worker_addr = bytes_(w.currentaddr.value,'ascii')
                            timeprocessing = time.time() - w.currentstart.value
                            if ( len(worker_addr) == 0
                                    and timeprocessing > self.config.multiprocess_timeout-0.1):
                                log.debug('worker %d has finished its work item, '
                                          'but is not exiting? do we wait for it?',
                                          iworker)
                            else:
                                any_alive = True
                            if (len(worker_addr) > 0
                                and timeprocessing > self.config.multiprocess_timeout-0.1):
                                log.debug('timed out worker %s: %s',
                                          iworker,worker_addr)
                                w.currentaddr.value = bytes_('')
                                # If the process is in C++ code, sending a SIGILL
                                # might not send a python KeybordInterrupt exception
                                # therefore, send multiple signals until an
                                # exception is caught. If this takes too long, then
                                # terminate the process
                                w.keyboardCaught.clear()
                                startkilltime = time.time()
                                while not w.keyboardCaught.is_set() and w.is_alive():
                                    if time.time()-startkilltime > self.waitkilltime:
                                        # have to terminate...
                                        log.error("terminating worker %s",iworker)
                                        w.terminate()
                                        # there is a small probability that the
                                        # terminated process might send a result,
                                        # which has to be specially handled or
                                        # else processes might get orphaned.
                                        workers[iworker] = w = self.startProcess(iworker, testQueue, resultQueue, shouldStop, result)
                                        break
                                    os.kill(w.pid, signal.SIGILL)
                                    time.sleep(0.1)
                    if not any_alive and testQueue.empty():
                        log.debug("All workers dead")
                        break
                nexttimeout=self.config.multiprocess_timeout
                for w in workers:
                    if w.is_alive() and len(w.currentaddr.value) > 0:
                        timeprocessing = time.time()-w.currentstart.value
                        if timeprocessing <= self.config.multiprocess_timeout:
                            nexttimeout = min(nexttimeout,
                                self.config.multiprocess_timeout-timeprocessing)
            log.debug("Completed %s tasks (%s remain)", len(completed), len(tasks))

        except (KeyboardInterrupt, SystemExit), e:
            log.info('parent received ctrl-c when waiting for test results')
            thrownError = e
            #resultQueue.get(False)

            result.addError(test, sys.exc_info())

        try:
            for case in to_teardown:
                log.debug("Tearing down shared fixtures for %s", case)
                try:
                    case.tearDown()
                except (KeyboardInterrupt, SystemExit):
                    raise
                except:
                    result.addError(case, sys.exc_info())

            stop = time.time()

            # first write since can freeze on shutting down processes
            result.printErrors()
            result.printSummary(start, stop)
            self.config.plugins.finalize(result)

            if thrownError is None:
                log.debug("Tell all workers to stop")
                for w in workers:
                    if w.is_alive():
                        testQueue.put('STOP', block=False)

            # wait for the workers to end
            for iworker,worker in enumerate(workers):
                if worker.is_alive():
                    log.debug('joining worker %s',iworker)
                    worker.join()
                    if worker.is_alive():
                        log.debug('failed to join worker %s',iworker)
        except (KeyboardInterrupt, SystemExit):
            log.info('parent received ctrl-c when shutting down: stop all processes')
            for worker in workers:
                if worker.is_alive():
                    worker.terminate()
            if thrownError: raise thrownError
            else: raise

        return result

    def addtask(testQueue,tasks,case):
        arg = None
        if isinstance(case,nose.case.Test) and hasattr(case.test,'arg'):
            # this removes the top level descriptor and allows real function
            # name to be returned
            case.test.descriptor = None
            arg = case.test.arg
        test_addr = MultiProcessTestRunner.address(case)
        testQueue.put((test_addr,arg), block=False)
        if arg is not None:
            test_addr += str(arg)
        if tasks is not None:
            tasks.append(test_addr)
        return test_addr
    addtask = staticmethod(addtask)

    def address(case):
        if hasattr(case, 'address'):
            file, mod, call = case.address()
        elif hasattr(case, 'context'):
            file, mod, call = test_address(case.context)
        else:
            raise Exception("Unable to convert %s to address" % case)
        parts = []
        if file is None:
            if mod is None:
                raise Exception("Unaddressable case %s" % case)
            else:
                parts.append(mod)
        else:
            # strip __init__.py(c) from end of file part
            # if present, having it there confuses loader
            dirname, basename = os.path.split(file)
            if basename.startswith('__init__'):
                file = dirname
            parts.append(file)
        if call is not None:
            parts.append(call)
        return ':'.join(map(str, parts))
    address = staticmethod(address)

    def nextBatch(self, test):
        # allows tests or suites to mark themselves as not safe
        # for multiprocess execution
        if hasattr(test, 'context'):
            if not getattr(test.context, '_multiprocess_', True):
                return

        if ((isinstance(test, ContextSuite)
             and test.hasFixtures(self.checkCanSplit))
            or not getattr(test, 'can_split', True)
            or not isinstance(test, unittest.TestSuite)):
            # regular test case, or a suite with context fixtures

            # special case: when run like nosetests path/to/module.py
            # the top-level suite has only one item, and it shares
            # the same context as that item. In that case, we want the
            # item, not the top-level suite
            if isinstance(test, ContextSuite):
                contained = list(test)
                if (len(contained) == 1
                    and getattr(contained[0],
                                'context', None) == test.context):
                    test = contained[0]
            yield test
        else:
            # Suite is without fixtures at this level; but it may have
            # fixtures at any deeper level, so we need to examine it all
            # the way down to the case level
            for case in test:
                for batch in self.nextBatch(case):
                    yield batch

    def checkCanSplit(context, fixt):
        """
        Callback that we use to check whether the fixtures found in a
        context or ancestor are ones we care about.

        Contexts can tell us that their fixtures are reentrant by setting
        _multiprocess_can_split_. So if we see that, we return False to
        disregard those fixtures.
        """
        if not fixt:
            return False
        if getattr(context, '_multiprocess_can_split_', False):
            return False
        return True
    checkCanSplit = staticmethod(checkCanSplit)

    def sharedFixtures(self, case):
        context = getattr(case, 'context', None)
        if not context:
            return False
        return getattr(context, '_multiprocess_shared_', False)

    def consolidate(self, result, batch_result):
        log.debug("batch result is %s" , batch_result)
        try:
            output, testsRun, failures, errors, errorClasses = batch_result
        except ValueError:
            log.debug("result in unexpected format %s", batch_result)
            failure.Failure(*sys.exc_info())(result)
            return
        self.stream.write(output)
        result.testsRun += testsRun
        result.failures.extend(failures)
        result.errors.extend(errors)
        for key, (storage, label, isfail) in errorClasses.items():
            if key not in result.errorClasses:
                # Ordinarily storage is result attribute
                # but it's only processed through the errorClasses
                # dict, so it's ok to fake it here
                result.errorClasses[key] = ([], label, isfail)
            mystorage, _junk, _junk = result.errorClasses[key]
            mystorage.extend(storage)
        log.debug("Ran %s tests (total: %s)", testsRun, result.testsRun)

def runner(ix, testQueue, resultQueue, currentaddr, currentstart,
           keyboardCaught, shouldStop, loaderClass, resultClass, config):
    try:
        try:
            return __runner(ix, testQueue, resultQueue, currentaddr, currentstart,
                    keyboardCaught, shouldStop, loaderClass, resultClass, config)
        except KeyboardInterrupt:
            log.debug('Worker %s keyboard interrupt, stopping',ix)
    except Empty:
        log.debug("Worker %s timed out waiting for tasks", ix)

def __runner(ix, testQueue, resultQueue, currentaddr, currentstart,
           keyboardCaught, shouldStop, loaderClass, resultClass, config):

    config = pickle.loads(config)
    dummy_parser = config.parserClass()
    if _instantiate_plugins is not None:
        for pluginclass in _instantiate_plugins:
            plugin = pluginclass()
            plugin.addOptions(dummy_parser,{})
            config.plugins.addPlugin(plugin)
    config.plugins.configure(config.options,config)
    config.plugins.begin()
    log.debug("Worker %s executing, pid=%d", ix,os.getpid())
    loader = loaderClass(config=config)
    loader.suiteClass.suiteClass = NoSharedFixtureContextSuite

    def get():
        return testQueue.get(timeout=config.multiprocess_timeout)

    def makeResult():
        stream = _WritelnDecorator(StringIO())
        result = resultClass(stream, descriptions=1,
                             verbosity=config.verbosity,
                             config=config)
        plug_result = config.plugins.prepareTestResult(result)
        if plug_result:
            return plug_result
        return result

    def batch(result):
        failures = [(TestLet(c), err) for c, err in result.failures]
        errors = [(TestLet(c), err) for c, err in result.errors]
        errorClasses = {}
        for key, (storage, label, isfail) in result.errorClasses.items():
            errorClasses[key] = ([(TestLet(c), err) for c, err in storage],
                                 label, isfail)
        return (
            result.stream.getvalue(),
            result.testsRun,
            failures,
            errors,
            errorClasses)
    for test_addr, arg in iter(get, 'STOP'):
        if shouldStop.is_set():
            log.exception('Worker %d STOPPED',ix)
            break
        result = makeResult()
        test = loader.loadTestsFromNames([test_addr])
        test.testQueue = testQueue
        test.tasks = []
        test.arg = arg
        log.debug("Worker %s Test is %s (%s)", ix, test_addr, test)
        try:
            if arg is not None:
                test_addr = test_addr + str(arg)
            currentaddr.value = bytes_(test_addr)
            currentstart.value = time.time()
            test(result)
            currentaddr.value = bytes_('')
            resultQueue.put((ix, test_addr, test.tasks, batch(result)))
        except KeyboardInterrupt, e: #TimedOutException:
            timeout = isinstance(e, TimedOutException)
            if timeout:
                keyboardCaught.set()
            if len(currentaddr.value):
                if timeout:
                    msg = 'Worker %s timed out, failing current test %s'
                else:
                    msg = 'Worker %s keyboard interrupt, failing current test %s'
                log.exception(msg,ix,test_addr)
                currentaddr.value = bytes_('')
                failure.Failure(*sys.exc_info())(result)
                resultQueue.put((ix, test_addr, test.tasks, batch(result)))
            else:
                if timeout:
                    msg = 'Worker %s test %s timed out'
                else:
                    msg = 'Worker %s test %s keyboard interrupt'
                log.debug(msg,ix,test_addr)
                resultQueue.put((ix, test_addr, test.tasks, batch(result)))
            if not timeout:
                raise
        except SystemExit:
            currentaddr.value = bytes_('')
            log.exception('Worker %s system exit',ix)
            raise
        except:
            currentaddr.value = bytes_('')
            log.exception("Worker %s error running test or returning "
                            "results",ix)
            failure.Failure(*sys.exc_info())(result)
            resultQueue.put((ix, test_addr, test.tasks, batch(result)))
        if config.multiprocess_restartworker:
            break
    log.debug("Worker %s ending", ix)

class NoSharedFixtureContextSuite(ContextSuite):
    """
    Context suite that never fires shared fixtures.

    When a context sets _multiprocess_shared_, fixtures in that context
    are executed by the main process. Using this suite class prevents them
    from executing in the runner process as well.

    """
    testQueue = None
    tasks = None
    arg = None
    def setupContext(self, context):
        if getattr(context, '_multiprocess_shared_', False):
            return
        super(NoSharedFixtureContextSuite, self).setupContext(context)

    def teardownContext(self, context):
        if getattr(context, '_multiprocess_shared_', False):
            return
        super(NoSharedFixtureContextSuite, self).teardownContext(context)
    def run(self, result):
        """Run tests in suite inside of suite fixtures.
        """
        # proxy the result for myself
        log.debug("suite %s (%s) run called, tests: %s",
                  id(self), self, self._tests)
        if self.resultProxy:
            result, orig = self.resultProxy(result, self), result
        else:
            result, orig = result, result
        try:
            #log.debug('setUp for %s', id(self));
            self.setUp()
        except KeyboardInterrupt:
            raise
        except:
            self.error_context = 'setup'
            result.addError(self, self._exc_info())
            return
        try:
            for test in self._tests:
                if (isinstance(test,nose.case.Test)
                    and self.arg is not None):
                    test.test.arg = self.arg
                else:
                    test.arg = self.arg
                test.testQueue = self.testQueue
                test.tasks = self.tasks
                if result.shouldStop:
                    log.debug("stopping")
                    break
                # each nose.case.Test will create its own result proxy
                # so the cases need the original result, to avoid proxy
                # chains
                #log.debug('running test %s in suite %s', test, self);
                try:
                    test(orig)
                except KeyboardInterrupt, e:
                    timeout = isinstance(e, TimedOutException)
                    if timeout:
                        msg = 'Timeout when running test %s in suite %s'
                    else:
                        msg = 'KeyboardInterrupt when running test %s in suite %s'
                    log.debug(msg, test, self)
                    err = (TimedOutException,TimedOutException(str(test)),
                           sys.exc_info()[2])
                    test.config.plugins.addError(test,err)
                    orig.addError(test,err)
                    if not timeout:
                        raise
        finally:
            self.has_run = True
            try:
                #log.debug('tearDown for %s', id(self));
                self.tearDown()
            except KeyboardInterrupt:
                raise
            except:
                self.error_context = 'teardown'
                result.addError(self, self._exc_info())


"""
A context manager for managing things injected into :mod:`__builtin__`.

Authors:

* Brian Granger
* Fernando Perez
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  The IPython Development Team.
#
#  Distributed under the terms of the BSD License.
#
#  Complete license in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import __builtin__

from IPython.config.configurable import Configurable

from IPython.utils.traitlets import Instance

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

class __BuiltinUndefined(object): pass
BuiltinUndefined = __BuiltinUndefined()

class __HideBuiltin(object): pass
HideBuiltin = __HideBuiltin()

class BuiltinTrap(Configurable):

    shell = Instance('IPython.core.interactiveshell.InteractiveShellABC')

    def __init__(self, shell=None):
        super(BuiltinTrap, self).__init__(shell=shell, config=None)
        self._orig_builtins = {}
        # We define this to track if a single BuiltinTrap is nested.
        # Only turn off the trap when the outermost call to __exit__ is made.
        self._nested_level = 0
        self.shell = shell
        # builtins we always add - if set to HideBuiltin, they will just
        # be removed instead of being replaced by something else
        self.auto_builtins = {'exit': HideBuiltin,
                              'quit': HideBuiltin,
                              'get_ipython': self.shell.get_ipython,
                              }
        # Recursive reload function
        try:
            from IPython.lib import deepreload
            if self.shell.deep_reload:
                self.auto_builtins['reload'] = deepreload.reload
            else:
                self.auto_builtins['dreload']= deepreload.reload
        except ImportError:
            pass

    def __enter__(self):
        if self._nested_level == 0:
            self.activate()
        self._nested_level += 1
        # I return self, so callers can use add_builtin in a with clause.
        return self

    def __exit__(self, type, value, traceback):
        if self._nested_level == 1:
            self.deactivate()
        self._nested_level -= 1
        # Returning False will cause exceptions to propagate
        return False

    def add_builtin(self, key, value):
        """Add a builtin and save the original."""
        bdict = __builtin__.__dict__
        orig = bdict.get(key, BuiltinUndefined)
        if value is HideBuiltin:
            if orig is not BuiltinUndefined: #same as 'key in bdict'
                self._orig_builtins[key] = orig
                del bdict[key]
        else:
            self._orig_builtins[key] = orig
            bdict[key] = value

    def remove_builtin(self, key, orig):
        """Remove an added builtin and re-set the original."""
        if orig is BuiltinUndefined:
            del __builtin__.__dict__[key]
        else:
            __builtin__.__dict__[key] = orig

    def activate(self):
        """Store ipython references in the __builtin__ namespace."""

        add_builtin = self.add_builtin
        for name, func in self.auto_builtins.iteritems():
            add_builtin(name, func)

    def deactivate(self):
        """Remove any builtins which might have been added by add_builtins, or
        restore overwritten ones to their previous values."""
        remove_builtin = self.remove_builtin
        for key, val in self._orig_builtins.iteritems():
            remove_builtin(key, val)
        self._orig_builtins.clear()
        self._builtins_added = False