<commit_before># -*- coding: utf-8 -*-
# Copyright (C) 2012-2015 MUJIN Inc

import threading

from . import zmq
from . import TimeoutError, GetMonotonicTime

import logging
log = logging.getLogger(__name__)

class ZmqSocketPool(object):

    _url = None  # url that sockets should connect to
    _ctx = None  # the context to use
    _ctxown = None  # the context owned exclusively
    _timeout = None  # timeout waiting for either send on receive while a socket is in the polling state
    _limit = None  # limit on number of socket alive at any time, None means unlimited

    _isok = False  # whether it is time to terminate
    _poller = None  # znq poller that tracks sockets in self._pollingsockets
    _sockets = None  # all sockets alive, a dictionary mapping from socket to True
    _pollingsockets = None  # sockets that are in the polling state, waiting for send or receive, a dictionary mapping from socket to timestamp when added to the poller
    _availablesockets = None  # list of sockets that are ready for use immediately

    _acquirecount = 0  # number of times a socket is acuired
    _releasecount = 0  # number of times a socket is released
    _opencount = 0  # number of times we opened a new socket
    _closecount = 0  # number of times we closed a socket

    def __init__(self, url, ctx=None, timeout=10.0, limit=None):
        """creates a socket pool, the pool can lease out socket for send and recv.
        if caller only does send but not recv, the socket will not be reused until internally the pool recv data and discard them.

        :param url: url for sockets to connect to
        :param ctx: optionally force socket to use provided zmq context instead of creating a new zmq context
        :param timeout: defaults to 10 seconds, specifies how long should we be polling the socket, until we consider it to be dead, note this is not a blocking timeout.
        :param limit: limit on number of socket alive at any time, None means unlimited
        """
        self._url = url
        self._ctx = ctx
        if self._ctx is None:
            self._ctxown = zmq.Context()
            self._ctx = self._ctxown
        self._timeout = timeout
        self._limit = limit

        self._isok = True
        self._poller = zmq.Poller()
        self._sockets = {}
        self._pollingsockets = {}
        self._availablesockets = []

        self._acquirecount = 0
        self._releasecount = 0
        self._opencount = 0
        self._closecount = 0

    def __del__(self):
        self.Destroy()

    def Destroy(self):
        self.SetDestroy()

        # at this point, all sockets should have been released
        assert(self._acquirecount == self._releasecount)
        sockets = list(self._pollingsockets.keys())
        while len(sockets) > 0:
            self._StopPollingSocket(sockets.pop())

        # close all socket ever created
        sockets = list(self._sockets.keys())
        while len(sockets) > 0:
            self._CloseSocket(sockets.pop())
        assert(self._opencount == self._closecount)

        # log.debug('sockets: opened = %d, closed = %d, acquired = %d, released = %d', self._opencount, self._closecount, self._acquirecount, self._releasecount)

        if self._ctxown is not None:
            try:
                self._ctxown.destroy()
            except Exception:
                log.exception('caught exception when destroying context')
            self._ctxown = None
        self._ctx = None

    def SetDestroy(self):
        # make sure no new socket can be created
        self._isok = False

        # make sure no one can acquire a socket now
        self._availablesockets = []

    def _OpenSocket(self):
        if not self._isok:
            return None

        socket = self._ctx.socket(zmq.REQ)
        socket.connect(self._url)
        assert(socket not in self._sockets)
        self._sockets[socket] = True
        self._opencount += 1

        # log.debug('opened a socket, url = %s opened = %d, closed = %d', self._url, self._opencount, self._closecount)
        return socket

    def _CloseSocket(self, socket):
        assert(socket in self._sockets)
        self._closecount += 1
        del self._sockets[socket]
        try:
            # make sure we do not linger when closing socket
            socket.close(linger=0)
        except Exception as e:
            log.exception('caught exception when closing socket: %s', e)

    def _StartPollingSocket(self, socket):
        assert(socket not in self._pollingsockets)
        self._pollingsockets[socket] = GetMonotonicTime()
        self._poller.register(socket, zmq.POLLIN | zmq.POLLOUT)

    def _StopPollingSocket(self, socket):
        assert(socket in self._pollingsockets)
        self._poller.unregister(socket)
        del self._pollingsockets[socket]

    def _Poll(self, timeout=0):
        """spin once and does internal polling of sockets
        """
        now = GetMonotonicTime()

        # poll for receives, non blocking if timeout is 0
        for socket, event in self._poller.poll(timeout):
            if (event & zmq.POLLIN) == zmq.POLLIN:
                # log.debug('a socket is ready for receive, url = %s, polling = %d, availble = %d', self._url, len(self._pollingsockets), len(self._availablesockets))

                # at least one message can be received without blocking
                try:
                    socket.recv(zmq.NOBLOCK)
                except Exception as e:
                    # when error occur, throw the socket away
                    log.exception('caught exception when recv: %s', e)
                    self._StopPollingSocket(socket)
                    self._CloseSocket(socket)
                    continue

                # reset the timestamp since we just called recv
                self._pollingsockets[socket] = now

        # poll again for send, non blocking if timeout is 0, some previously in recv state socket may now be available for sending
        for socket, event in self._poller.poll(timeout):
            if (event & zmq.POLLOUT) == zmq.POLLOUT:
                # log.debug('a socket is ready for send, url = %s, polling = %d, availble = %d', self._url, len(self._pollingsockets), len(self._availablesockets))

                # at least one message can be sent to socket without blocking
                self._StopPollingSocket(socket)
                self._availablesockets.append(socket)

        # check for socket that stayed in the polling state for too long
        timedoutsockets = []
        for socket, timestamp in self._pollingsockets.iteritems():
            if now - timestamp > self._timeout:
                timedoutsockets.append(socket)

        # close timed out sockets
        while len(timedoutsockets) > 0:
            socket = timedoutsockets.pop()
            self._StopPollingSocket(socket)
            self._CloseSocket(socket)

    def AcquireSocket(self, timeout=None, checkpreemptfn=None):
        """acquire a socket from the list of availble sockets for sending
        """
        # first we try a non blocking poll
        self._Poll(timeout=0)

        starttime = GetMonotonicTime()
        while self._isok:

            # if a socket is available, use it
            if len(self._availablesockets) > 0:
                socket = self._availablesockets.pop()
                self._acquirecount += 1
                return socket

            # if we don't have available socket but we have quota to create new socket, create a new one
            if self._limit is None or len(self._availablesockets) + len(self._pollingsockets) < self._limit:
                socket = self._OpenSocket()
                self._acquirecount += 1
                return socket

            # check for timeout
            elapsedtime = GetMonotonicTime() - starttime
            if timeout is not None and elapsedtime > timeout:
                raise TimeoutError(u'Timed out waiting for a socket to %s to become available after %f seconds' % (self._url, elapsedtime))

            if checkpreemptfn is not None:
                checkpreemptfn()

            # otherwise, wait by a small blocking poll
            self._Poll(timeout=50)

        # TODO: raise UserInterrupt
        return None

    def ReleaseSocket(self, socket, reuse=True):
        """release a socket after use, if caller did not call recv, the pool will take care of that
        """
        if socket is not None:
            self._releasecount += 1
            if self._isok and reuse:
                self._StartPollingSocket(socket)
            else:
                self._CloseSocket(socket)

class ZmqClient(object):

    _hostname = None
    _port = None
    _url = None
    _pool = None
    _socket = None
    _isok = False
    _callerthread = None  # last caller thread

    def __init__(self, hostname='', port=0, ctx=None, limit=100, url=None, checkpreemptfn=None, reusetimeout=10.0):
        """creates a new zmq client, uses zmq req socket over tcp

        :param hostname: hostname or ip to connect to
        :param port: port to connect to
        :param ctx: optionally specifies a zmq context to use, if not provided, a new zmq context will be created
        :param limit: defaults to 100, limit the number of underlying zmq socket to create, usually one socket will actually be used, but if server times out frequently and you use fireandforget, then this limit caps the number of sockets we can use
        :param url: allow passing of zmq socket url instead of hostname and port
        """

        self._hostname = hostname
        self._port = int(port)
        self._url = url
        if self._url is None:
            self._url = 'tcp://%s:%d' % (self._hostname, self._port)

        self._pool = ZmqSocketPool(self._url, ctx=ctx, limit=limit, timeout=reusetimeout)
        self._socket = None
        self._isok = True
        self._checkpreemptfn=checkpreemptfn

    def __del__(self):
        self.Destroy()

    def Destroy(self):
        self.SetDestroy()

        if self._pool is not None:
            self._ReleaseSocket()
            self._pool.Destroy()
            self._pool = None

    def SetDestroy(self):
        self._isok = False
        if self._pool is not None:
            self._pool.SetDestroy()

    def GetHostname(self):
        """returns the hostname given when constructing the client
        """
        return self._hostname

    def GetPort(self):
        """returns the port given when constructing the client
        """
        return self._port

    # TODO: this is for backward compatibility, remove once all callers are updated
    @property
    def hostname(self):
        return self._hostname

    # TODO: this is for backward compatibility, remove once all callers are updated
    @property
    def port(self):
        return self._port

    def _CheckCallerThread(self, context=None):
        """catch bad caller who use zmq client from multiple threads and causes random race conditions.
        """
        callerthread = repr(threading.current_thread())
        oldcallerthread = self._callerthread
        self._callerthread = callerthread
        if oldcallerthread is not None:
            # assert oldcallerthread == callerthread, 'zmqclient used from multiple threads: previously = %s, now = %s' % (oldcallerthread, callerthread)
            if oldcallerthread != callerthread:
                log.error('zmqclient used from multiple threads in %s, this is a bug in the caller: previously = %s, now = %s' % (context, oldcallerthread, callerthread))

    def _AcquireSocket(self, timeout=None, checkpreempt=True):
        # if we were holding on to a socket before, release it before acquiring one
        self._ReleaseSocket()
        self._socket = self._pool.AcquireSocket(timeout=timeout, checkpreemptfn=self._checkpreemptfn if checkpreempt else None)

    def _ReleaseSocket(self):
        if self._socket is not None:
            self._pool.ReleaseSocket(self._socket)
            self._socket = None

    def SetPreemptFn(self, checkpreemptfn):
        self._checkpreemptfn = checkpreemptfn

    def SendCommand(self, command, timeout=10.0, blockwait=True, fireandforget=False, sendjson=True, recvjson=True, checkpreempt=True):
        """sends command via established zmq socket

        :param command: command in json format
        :param timeout: if None, block. If >= 0, use as timeout
        :param blockwait: if True (default), will call receive also, otherwise, caller needs to call ReceiveCommand later
        :param fireandforget: if True, will send command and immediately return without trying to receive, blockwait will be set to False
        :param sendjson: if True (default), will send data as json
        :param recvjson: if True (default), will parse received data as json

        :return: returns the response from the zmq server in json format if blockwait is True
        """
        # log.debug('Sending command via ZMQ: %s', command)

        self._CheckCallerThread(command)

        if fireandforget:
            blockwait = False

        # acquire a socket for sending
        self._AcquireSocket(timeout=timeout, checkpreempt=checkpreempt)

        # we may be exiting, the pool refused to give us a socket
        if not self._isok or self._socket is None:
            # TODO: raise UserInterrupt
            return None

        releasesocket = True
        try:
            # send phase
            starttime = GetMonotonicTime()
            while self._isok:
                # timeout checking
                elapsedtime = GetMonotonicTime() - starttime
                if timeout is not None and elapsedtime > timeout:
                    raise TimeoutError(u'Timed out trying to send to %s after %f seconds' % (self._url, elapsedtime))

                if checkpreempt and self._checkpreemptfn is not None:
                    self._checkpreemptfn()

                # poll to see if we can send, if not, loop
                waitingevents = self._socket.poll(50, zmq.POLLOUT)
                if (waitingevents & zmq.POLLOUT) != zmq.POLLOUT:
                    continue

                if sendjson:
                    self._socket.send_json(command, zmq.NOBLOCK)
                else:
                    self._socket.send(command, zmq.NOBLOCK)

                # break when successfully sent
                break

            # for fire and forget, no need to receive
            if fireandforget:
                return None

            # keep the socket and let caller call receive
            if not blockwait:
                releasesocket = False
                return None

            # receive
            return self.ReceiveCommand(timeout=timeout, recvjson=recvjson, checkpreempt=checkpreempt)

        finally:
            # release socket
            if releasesocket:
                self._ReleaseSocket()

        # TODO: raise UserInterrupt
        return None

    def IsWaitingReply(self):
        return self._socket is not None

    def ReceiveCommand(self, timeout=10.0, recvjson=True, checkpreempt=True):
        """receive response to a previous SendCommand call, SendCommand must be called with blockwait=False and fireandforget=False

        :param timeout: if None, block. If >= 0, use as timeout
        :param recvjson: if True (default), will parse received data as json

        :return: returns the recv or recv_json response
        """
        self._CheckCallerThread('ReceiveCommand')

        # should have called SendCommand with blockwait=False first
        assert(self._socket is not None)

        try:
            # receive phase
            starttime = GetMonotonicTime()
            while self._isok:
                # timeout checking
                elapsedtime = GetMonotonicTime() - starttime
                if timeout is not None and elapsedtime > timeout:
                    raise TimeoutError(u'Timed out to get response from %s after %f seconds (timeout=%f)' % (self._url, elapsedtime, timeout))

                if checkpreempt and self._checkpreemptfn is not None:
                    self._checkpreemptfn()

                # poll to see if something has been received, if received nothing, loop
                startpolltime = GetMonotonicTime()
                waitingevents = self._socket.poll(50, zmq.POLLIN)
                endpolltime = GetMonotonicTime()
                if endpolltime - startpolltime > 0.2:  # due to python delays sometimes this can be 0.11s
                    log.critical('polling time took %fs!', endpolltime - startpolltime)
                if (waitingevents & zmq.POLLIN) != zmq.POLLIN:
                    continue

                if recvjson:
                    return self._socket.recv_json(zmq.NOBLOCK)
                else:
                    return self._socket.recv(zmq.NOBLOCK)

        finally:
            # release socket
            self._ReleaseSocket()

        # TODO: raise UserInterrupt
        return None
<commit_msg>fix whitespaces<commit_after># -*- coding: utf-8 -*-
# Copyright (C) 2012-2015 MUJIN Inc

import threading

from . import zmq
from . import TimeoutError, GetMonotonicTime

import logging
log = logging.getLogger(__name__)

class ZmqSocketPool(object):

    _url = None  # url that sockets should connect to
    _ctx = None  # the context to use
    _ctxown = None  # the context owned exclusively
    _timeout = None  # timeout waiting for either send on receive while a socket is in the polling state
    _limit = None  # limit on number of socket alive at any time, None means unlimited

    _isok = False  # whether it is time to terminate
    _poller = None  # znq poller that tracks sockets in self._pollingsockets
    _sockets = None  # all sockets alive, a dictionary mapping from socket to True
    _pollingsockets = None  # sockets that are in the polling state, waiting for send or receive, a dictionary mapping from socket to timestamp when added to the poller
    _availablesockets = None  # list of sockets that are ready for use immediately

    _acquirecount = 0  # number of times a socket is acuired
    _releasecount = 0  # number of times a socket is released
    _opencount = 0  # number of times we opened a new socket
    _closecount = 0  # number of times we closed a socket

    def __init__(self, url, ctx=None, timeout=10.0, limit=None):
        """creates a socket pool, the pool can lease out socket for send and recv.
        if caller only does send but not recv, the socket will not be reused until internally the pool recv data and discard them.

        :param url: url for sockets to connect to
        :param ctx: optionally force socket to use provided zmq context instead of creating a new zmq context
        :param timeout: defaults to 10 seconds, specifies how long should we be polling the socket, until we consider it to be dead, note this is not a blocking timeout.
        :param limit: limit on number of socket alive at any time, None means unlimited
        """
        self._url = url
        self._ctx = ctx
        if self._ctx is None:
            self._ctxown = zmq.Context()
            self._ctx = self._ctxown
        self._timeout = timeout
        self._limit = limit

        self._isok = True
        self._poller = zmq.Poller()
        self._sockets = {}
        self._pollingsockets = {}
        self._availablesockets = []

        self._acquirecount = 0
        self._releasecount = 0
        self._opencount = 0
        self._closecount = 0

    def __del__(self):
        self.Destroy()

    def Destroy(self):
        self.SetDestroy()

        # at this point, all sockets should have been released
        assert(self._acquirecount == self._releasecount)
        sockets = list(self._pollingsockets.keys())
        while len(sockets) > 0:
            self._StopPollingSocket(sockets.pop())

        # close all socket ever created
        sockets = list(self._sockets.keys())
        while len(sockets) > 0:
            self._CloseSocket(sockets.pop())
        assert(self._opencount == self._closecount)

        # log.debug('sockets: opened = %d, closed = %d, acquired = %d, released = %d', self._opencount, self._closecount, self._acquirecount, self._releasecount)

        if self._ctxown is not None:
            try:
                self._ctxown.destroy()
            except Exception:
                log.exception('caught exception when destroying context')
            self._ctxown = None
        self._ctx = None

    def SetDestroy(self):
        # make sure no new socket can be created
        self._isok = False

        # make sure no one can acquire a socket now
        self._availablesockets = []

    def _OpenSocket(self):
        if not self._isok:
            return None

        socket = self._ctx.socket(zmq.REQ)
        socket.connect(self._url)
        assert(socket not in self._sockets)
        self._sockets[socket] = True
        self._opencount += 1

        # log.debug('opened a socket, url = %s opened = %d, closed = %d', self._url, self._opencount, self._closecount)
        return socket

    def _CloseSocket(self, socket):
        assert(socket in self._sockets)
        self._closecount += 1
        del self._sockets[socket]
        try:
            # make sure we do not linger when closing socket
            socket.close(linger=0)
        except Exception as e:
            log.exception('caught exception when closing socket: %s', e)

    def _StartPollingSocket(self, socket):
        assert(socket not in self._pollingsockets)
        self._pollingsockets[socket] = GetMonotonicTime()
        self._poller.register(socket, zmq.POLLIN | zmq.POLLOUT)

    def _StopPollingSocket(self, socket):
        assert(socket in self._pollingsockets)
        self._poller.unregister(socket)
        del self._pollingsockets[socket]

    def _Poll(self, timeout=0):
        """spin once and does internal polling of sockets
        """
        now = GetMonotonicTime()

        # poll for receives, non blocking if timeout is 0
        for socket, event in self._poller.poll(timeout):
            if (event & zmq.POLLIN) == zmq.POLLIN:
                # log.debug('a socket is ready for receive, url = %s, polling = %d, availble = %d', self._url, len(self._pollingsockets), len(self._availablesockets))

                # at least one message can be received without blocking
                try:
                    socket.recv(zmq.NOBLOCK)
                except Exception as e:
                    # when error occur, throw the socket away
                    log.exception('caught exception when recv: %s', e)
                    self._StopPollingSocket(socket)
                    self._CloseSocket(socket)
                    continue

                # reset the timestamp since we just called recv
                self._pollingsockets[socket] = now

        # poll again for send, non blocking if timeout is 0, some previously in recv state socket may now be available for sending
        for socket, event in self._poller.poll(timeout):
            if (event & zmq.POLLOUT) == zmq.POLLOUT:
                # log.debug('a socket is ready for send, url = %s, polling = %d, availble = %d', self._url, len(self._pollingsockets), len(self._availablesockets))

                # at least one message can be sent to socket without blocking
                self._StopPollingSocket(socket)
                self._availablesockets.append(socket)

        # check for socket that stayed in the polling state for too long
        timedoutsockets = []
        for socket, timestamp in self._pollingsockets.iteritems():
            if now - timestamp > self._timeout:
                timedoutsockets.append(socket)

        # close timed out sockets
        while len(timedoutsockets) > 0:
            socket = timedoutsockets.pop()
            self._StopPollingSocket(socket)
            self._CloseSocket(socket)

    def AcquireSocket(self, timeout=None, checkpreemptfn=None):
        """acquire a socket from the list of availble sockets for sending
        """
        # first we try a non blocking poll
        self._Poll(timeout=0)

        starttime = GetMonotonicTime()
        while self._isok:

            # if a socket is available, use it
            if len(self._availablesockets) > 0:
                socket = self._availablesockets.pop()
                self._acquirecount += 1
                return socket

            # if we don't have available socket but we have quota to create new socket, create a new one
            if self._limit is None or len(self._availablesockets) + len(self._pollingsockets) < self._limit:
                socket = self._OpenSocket()
                self._acquirecount += 1
                return socket

            # check for timeout
            elapsedtime = GetMonotonicTime() - starttime
            if timeout is not None and elapsedtime > timeout:
                raise TimeoutError(u'Timed out waiting for a socket to %s to become available after %f seconds' % (self._url, elapsedtime))

            if checkpreemptfn is not None:
                checkpreemptfn()

            # otherwise, wait by a small blocking poll
            self._Poll(timeout=50)

        # TODO: raise UserInterrupt
        return None

    def ReleaseSocket(self, socket, reuse=True):
        """release a socket after use, if caller did not call recv, the pool will take care of that
        """
        if socket is not None:
            self._releasecount += 1
            if self._isok and reuse:
                self._StartPollingSocket(socket)
            else:
                self._CloseSocket(socket)

class ZmqClient(object):

    _hostname = None
    _port = None
    _url = None
    _pool = None
    _socket = None
    _isok = False
    _callerthread = None  # last caller thread

    def __init__(self, hostname='', port=0, ctx=None, limit=100, url=None, checkpreemptfn=None, reusetimeout=10.0):
        """creates a new zmq client, uses zmq req socket over tcp

        :param hostname: hostname or ip to connect to
        :param port: port to connect to
        :param ctx: optionally specifies a zmq context to use, if not provided, a new zmq context will be created
        :param limit: defaults to 100, limit the number of underlying zmq socket to create, usually one socket will actually be used, but if server times out frequently and you use fireandforget, then this limit caps the number of sockets we can use
        :param url: allow passing of zmq socket url instead of hostname and port
        """

        self._hostname = hostname
        self._port = int(port)
        self._url = url
        if self._url is None:
            self._url = 'tcp://%s:%d' % (self._hostname, self._port)

        self._pool = ZmqSocketPool(self._url, ctx=ctx, limit=limit, timeout=reusetimeout)
        self._socket = None
        self._isok = True
        self._checkpreemptfn = checkpreemptfn

    def __del__(self):
        self.Destroy()

    def Destroy(self):
        self.SetDestroy()

        if self._pool is not None:
            self._ReleaseSocket()
            self._pool.Destroy()
            self._pool = None

    def SetDestroy(self):
        self._isok = False
        if self._pool is not None:
            self._pool.SetDestroy()

    def GetHostname(self):
        """returns the hostname given when constructing the client
        """
        return self._hostname

    def GetPort(self):
        """returns the port given when constructing the client
        """
        return self._port

    # TODO: this is for backward compatibility, remove once all callers are updated
    @property
    def hostname(self):
        return self._hostname

    # TODO: this is for backward compatibility, remove once all callers are updated
    @property
    def port(self):
        return self._port

    def _CheckCallerThread(self, context=None):
        """catch bad caller who use zmq client from multiple threads and causes random race conditions.
        """
        callerthread = repr(threading.current_thread())
        oldcallerthread = self._callerthread
        self._callerthread = callerthread
        if oldcallerthread is not None:
            # assert oldcallerthread == callerthread, 'zmqclient used from multiple threads: previously = %s, now = %s' % (oldcallerthread, callerthread)
            if oldcallerthread != callerthread:
                log.error('zmqclient used from multiple threads in %s, this is a bug in the caller: previously = %s, now = %s' % (context, oldcallerthread, callerthread))

    def _AcquireSocket(self, timeout=None, checkpreempt=True):
        # if we were holding on to a socket before, release it before acquiring one
        self._ReleaseSocket()
        self._socket = self._pool.AcquireSocket(timeout=timeout, checkpreemptfn=self._checkpreemptfn if checkpreempt else None)

    def _ReleaseSocket(self):
        if self._socket is not None:
            self._pool.ReleaseSocket(self._socket)
            self._socket = None

    def SetPreemptFn(self, checkpreemptfn):
        self._checkpreemptfn = checkpreemptfn

    def SendCommand(self, command, timeout=10.0, blockwait=True, fireandforget=False, sendjson=True, recvjson=True, checkpreempt=True):
        """sends command via established zmq socket

        :param command: command in json format
        :param timeout: if None, block. If >= 0, use as timeout
        :param blockwait: if True (default), will call receive also, otherwise, caller needs to call ReceiveCommand later
        :param fireandforget: if True, will send command and immediately return without trying to receive, blockwait will be set to False
        :param sendjson: if True (default), will send data as json
        :param recvjson: if True (default), will parse received data as json

        :return: returns the response from the zmq server in json format if blockwait is True
        """
        # log.debug('Sending command via ZMQ: %s', command)

        self._CheckCallerThread(command)

        if fireandforget:
            blockwait = False

        # acquire a socket for sending
        self._AcquireSocket(timeout=timeout, checkpreempt=checkpreempt)

        # we may be exiting, the pool refused to give us a socket
        if not self._isok or self._socket is None:
            # TODO: raise UserInterrupt
            return None

        releasesocket = True
        try:
            # send phase
            starttime = GetMonotonicTime()
            while self._isok:
                # timeout checking
                elapsedtime = GetMonotonicTime() - starttime
                if timeout is not None and elapsedtime > timeout:
                    raise TimeoutError(u'Timed out trying to send to %s after %f seconds' % (self._url, elapsedtime))

                if checkpreempt and self._checkpreemptfn is not None:
                    self._checkpreemptfn()

                # poll to see if we can send, if not, loop
                waitingevents = self._socket.poll(50, zmq.POLLOUT)
                if (waitingevents & zmq.POLLOUT) != zmq.POLLOUT:
                    continue

                if sendjson:
                    self._socket.send_json(command, zmq.NOBLOCK)
                else:
                    self._socket.send(command, zmq.NOBLOCK)

                # break when successfully sent
                break

            # for fire and forget, no need to receive
            if fireandforget:
                return None

            # keep the socket and let caller call receive
            if not blockwait:
                releasesocket = False
                return None

            # receive
            return self.ReceiveCommand(timeout=timeout, recvjson=recvjson, checkpreempt=checkpreempt)

        finally:
            # release socket
            if releasesocket:
                self._ReleaseSocket()

        # TODO: raise UserInterrupt
        return None

    def IsWaitingReply(self):
        return self._socket is not None

    def ReceiveCommand(self, timeout=10.0, recvjson=True, checkpreempt=True):
        """receive response to a previous SendCommand call, SendCommand must be called with blockwait=False and fireandforget=False

        :param timeout: if None, block. If >= 0, use as timeout
        :param recvjson: if True (default), will parse received data as json

        :return: returns the recv or recv_json response
        """
        self._CheckCallerThread('ReceiveCommand')

        # should have called SendCommand with blockwait=False first
        assert(self._socket is not None)

        try:
            # receive phase
            starttime = GetMonotonicTime()
            while self._isok:
                # timeout checking
                elapsedtime = GetMonotonicTime() - starttime
                if timeout is not None and elapsedtime > timeout:
                    raise TimeoutError(u'Timed out to get response from %s after %f seconds (timeout=%f)' % (self._url, elapsedtime, timeout))

                if checkpreempt and self._checkpreemptfn is not None:
                    self._checkpreemptfn()

                # poll to see if something has been received, if received nothing, loop
                startpolltime = GetMonotonicTime()
                waitingevents = self._socket.poll(50, zmq.POLLIN)
                endpolltime = GetMonotonicTime()
                if endpolltime - startpolltime > 0.2:  # due to python delays sometimes this can be 0.11s
                    log.critical('polling time took %fs!', endpolltime - startpolltime)
                if (waitingevents & zmq.POLLIN) != zmq.POLLIN:
                    continue

                if recvjson:
                    return self._socket.recv_json(zmq.NOBLOCK)
                else:
                    return self._socket.recv(zmq.NOBLOCK)

        finally:
            # release socket
            self._ReleaseSocket()

        # TODO: raise UserInterrupt
        return None
<|endoftext|>


<commit_before>#!/usr/bin/python3
"""
Certificate watcher which reloads nginx or reconfigures it, depending on what
happens to externally supplied certificates. Only executed by start.py in case
of TLS_FLAVOR=[mail, cert]
"""

from os.path import exists, split as path_split
from os import system
import time
from watchdog.observers.polling import PollingObserver
from watchdog.events import FileSystemEventHandler, FileDeletedEvent, \
    FileCreatedEvent, FileModifiedEvent, FileMovedEvent

class ChangeHandler(FileSystemEventHandler):
    "watchdog-handler listening on any event, executing the correct configuration/reload steps"
    @staticmethod
    def reload_nginx():
        "merely reload nginx without re-configuring everything"
        if exists("/var/run/nginx.pid"):
            print("Reloading a running nginx")
            system("nginx -s reload")

    @staticmethod
    def reexec_config():
        "execute a reconfiguration of the system, which also reloads"
        print("Reconfiguring system")
        system("/config.py")

    def on_any_event(self, event):
        "event-listener checking if the affected files are the cert-files we're interested in"
        if event.is_directory:
            return

        filename = path_split(event.src_path)[-1]
        if isinstance(event, FileMovedEvent):
            filename = path_split(event.dest_path)[-1]

        if filename in ['cert.pem', 'key.pem']:
            # all cases except for FileModified need re-configure
            if isinstance(event, (FileCreatedEvent, FileMovedEvent, FileDeletedEvent)):
                ChangeHandler.reexec_config()
            # file modification needs only a nginx reload without config.py
            elif isinstance(event, FileModifiedEvent):
                ChangeHandler.reload_nginx()
        # cert files have been moved away, re-configure
        elif isinstance(event, FileMovedEvent) and path_split(event.src_path)[-1] in ['cert.pem', 'key.pem']:
            ChangeHandler.reexec_config()

if __name__ == '__main__':
    observer = PollingObserver()
    handler = ChangeHandler()
    observer.schedule(handler, "/certs", recursive=False)
    observer.start()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        observer.stop()

    observer.join()
<commit_msg>Use environment variables for cert paths/names in nginx certwatcher<commit_after>#!/usr/bin/python3
"""
Certificate watcher which reloads nginx or reconfigures it, depending on what
happens to externally supplied certificates. Only executed by start.py in case
of TLS_FLAVOR=[mail, cert]
"""

from os.path import exists, split as path_split, join as path_join
from os import system, getenv
import time
from watchdog.observers.polling import PollingObserver
from watchdog.events import FileSystemEventHandler, FileDeletedEvent, \
    FileCreatedEvent, FileModifiedEvent, FileMovedEvent

class ChangeHandler(FileSystemEventHandler):
    "watchdog-handler listening on any event, executing the correct configuration/reload steps"

    def __init__(self, cert_path, keypair_path):
        "Initialize a new changehandler"""
        super().__init__()
        self.cert_path = cert_path
        self.keypair_path = keypair_path

    @staticmethod
    def reload_nginx():
        "merely reload nginx without re-configuring everything"
        if exists("/var/run/nginx.pid"):
            print("Reloading a running nginx")
            system("nginx -s reload")

    @staticmethod
    def reexec_config():
        "execute a reconfiguration of the system, which also reloads"
        print("Reconfiguring system")
        system("/config.py")

    def on_any_event(self, event):
        "event-listener checking if the affected files are the cert-files we're interested in"
        if event.is_directory:
            return

        filename = event.src_path
        if isinstance(event, FileMovedEvent):
            filename = event.dest_path

        if filename in [self.cert_path, self.keypair_path]:
            # all cases except for FileModified need re-configure
            if isinstance(event, (FileCreatedEvent, FileMovedEvent, FileDeletedEvent)):
                ChangeHandler.reexec_config()
            # file modification needs only a nginx reload without config.py
            elif isinstance(event, FileModifiedEvent):
                ChangeHandler.reload_nginx()
        # cert files have been moved away, re-configure
        elif isinstance(event, FileMovedEvent) and event.src_path in [self.cert_path, self.keypair_path]:
            ChangeHandler.reexec_config()

if __name__ == '__main__':
    cert_path = path_join("/certs/", getenv("TLS_CERT_FILENAME", default="cert.pem"))
    cert_dir = path_split(cert_path)[0]
    keypair_path = path_join("/certs/", getenv("TLS_KEYPAIR_FILENAME", default="key.pem"))
    keypair_dir = path_split(keypair_path)[0]

    observer = PollingObserver()
    handler = ChangeHandler(cert_path, keypair_path)
    observer.schedule(handler, cert_dir, recursive=False)
    if keypair_dir != cert_dir:
       observer.schedule(handler, keypair_dir, recursive=False)
    observer.start()

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        observer.stop()

    observer.join()
<|endoftext|>


<commit_before>import contextlib
import dateutil.parser
import json
import logging
import os
import shutil
import tempfile
from zipfile import ZipFile
from django.db import transaction

from rdrf.models import Registry
from .catalogue import DataGroupImporterCatalogue, ModelImporterCatalogue, MongoCollectionImporterCatalogue
from .importers import get_meta_value, allow_if_forced
import definitions
from .utils import DelegateMixin, IndentedLogger, app_schema_version

logger = logging.getLogger(__name__)

@contextlib.contextmanager
def zipfile_contents(zipfile):
    tmpdir = tempfile.mkdtemp(suffix='import')
    archive = ZipFile(zipfile)
    archive.extractall(tmpdir)

    yield tmpdir

    shutil.rmtree(tmpdir)

class ZipFileImporter(object):
    def __init__(self, zipfile, catalogue=None):
        if catalogue is None:
            catalogue = definitions.Catalogue(
                    DataGroupImporterCatalogue(),
                    ModelImporterCatalogue(),
                    MongoCollectionImporterCatalogue())
        self.catalogue = catalogue
        self.zipfile = zipfile
        self.workdir = None
        self.meta = None
        self.logger = logging.getLogger(__name__)

    def find_workdir(self, startdir):
        """Return the workdir and the meta file location in the archive.

        The zip file can contain the META file at top level or in a nested
        directory. We find the META file and the directory containing it will
        be the working directory."""

        meta = [(path, os.path.join(path, 'META')) for path, _, files in os.walk(startdir) if 'META' in files]
        if len(meta) == 0:
            raise Exception("Invalid export file '%s'.'META' file is missing." % self.zipfile)

        return meta[0]

    def extract_meta_info(self, meta_file):
        with open(meta_file) as f:
            self.meta = json.load(f)

    def create_importer(self, export_type, requested_import_type=None):
        self.file_export_type = definitions.EXPORT_TYPES.from_name(export_type)
        if requested_import_type is None:
            requested_import_type = self.file_export_type.code

        self.requested_type = definitions.EXPORT_TYPES.from_code(requested_import_type)
        if not (self.requested_type is self.file_export_type or self.requested_type in self.file_export_type.includes):
            raise Exception("Invalid import type '%s' requested for file '%s' with type '%s'." %
                    (requested_import_type, self.zipfile, self.file_export_type.code))

        if self.requested_type in definitions.EXPORT_TYPES.registry_types:
            return RegistryImporter(self)
        if self.requested_type is definitions.EXPORT_TYPES.CDES:
            return GenericImporter(self)
        if self.requested_type is definitions.EXPORT_TYPES.REFDATA:
            return GenericImporter(self)
        raise Exception("Unrecognized export type '%s'." % (self.requested_type.code))

    def do_import(self, import_type=None, verbose=False, indented_logs=True, simulate=False, force=False):
        if verbose:
            self.logger.setLevel(logging.DEBUG)
        self.child_logger = logger
        if indented_logs:
            self.child_logger = IndentedLogger(self.logger)
        self.simulate = simulate
        self.force = force

        with zipfile_contents(self.zipfile) as tmpdir:
            self.workdir, meta_file = self.find_workdir(tmpdir)
            self.extract_meta_info(meta_file)

            importer = self.create_importer(get_meta_value(self.meta, 'type'), requested_import_type=import_type)
            self.logger.debug('  Import type: %s (%s)', self.requested_type.name, self.requested_type.code)
            importer.output_import_info()
            importer.do_import()

    def inspect(self):
        self.logger.setLevel(logging.DEBUG)
        with zipfile_contents(self.zipfile) as tmpdir:
            self.workdir, meta_file = self.find_workdir(tmpdir)
            self.extract_meta_info(meta_file)
            importer = self.create_importer(get_meta_value(self.meta, 'type'))
            importer.output_import_info()

class RegistryLevelChecks(DelegateMixin):
    def __init__(self, importer):
        DelegateMixin.__init__(self, delegate_to=importer)

    def check_registry_export_type_in_meta(self):
        export_type = get_meta_value(self.meta, 'type')
        if export_type not in definitions.EXPORT_TYPES.registry_types_names:
            raise Exception("Invalid export type '%s' for registry import. Should be one of '%s'."
                % (export_type, ', '.join(definitions.EXPORT_TYPES.registry_types_names)))

    @allow_if_forced
    def check_registry_does_not_exist(self):
        if self.force:
            return
        if Registry.objects.filter(code=self.registry_code).exists():
            raise Exception("Registry '%s' already exists." % self.registry_code)

class BaseImporter(DelegateMixin):
    def __init__(self, zipfile_importer):
        DelegateMixin.__init__(self, delegate_to=zipfile_importer)

    def maybe_filter_meta(self, meta):
        if self.file_export_type is self.requested_type:
            return meta
        else:
            return filter(definitions.META_FILTERS[self.requested_type], meta)

    @allow_if_forced
    def check_app_schema_versions_match(self):
        app_schema_version_different = self.diff_app_versions()
        if len(app_schema_version_different) > 0:
            raise Exception('Schema difference detected between your registry and the export file.'
                    ' App(s) with different schema: %s' % ', '.join(app_schema_version_different))

    def import_datagroups(self, meta):
        meta = self.maybe_filter_meta(meta)
        self.check_app_schema_versions_match()
        self.logger.debug('Importing %d datagroups', len(meta))
        datagroup_importers = self.catalogue.datagroups
        for data_group_meta in meta:
            importer = datagroup_importers.get(data_group_meta['name'])(self.catalogue)
            options = {
                'logger': self.child_logger,
                'simulate': self.simulate,
                'force': self.force,
            }
            if hasattr(self, 'registry_code'):
                options['registry_code'] = self.registry_code
            importer.do_import(data_group_meta, self.workdir, **options)

    def output_import_info(self):
        logger = IndentedLogger(self.logger, indent_level=2)
        zipfile_type = definitions.EXPORT_TYPES.from_name(get_meta_value(self.meta, 'type'))
        logger.debug('Zipfile type: %s (%s)', zipfile_type.name, zipfile_type.code)
        also_includes = zipfile_type.includes
        if also_includes:
            logger.debug('(also includes import types: %s)' %
                    ', '.join(map(lambda t: "'%s' (%s)" % (t.name, t.code), also_includes)))

        app_schema_version_different = self.diff_app_versions()
        if len(app_schema_version_different) > 0:
            logger.warn('WARNING: Schema difference detected between your registry and the export file.'
                    'App(s): %s', ', '.join(app_schema_version_different))

        return logger

    def diff_app_versions(self):
        meta = self.maybe_filter_meta(get_meta_value(self.meta, 'data_groups'))
        app_versions = dict(reduce(lambda d, x: d + x.get('app_versions', {}).items(), meta, []))
        return filter(lambda app: app_versions[app] != app_schema_version(app), app_versions)

class RegistryImporter(BaseImporter):
    def __init__(self, zipfile_importer):
        BaseImporter.__init__(self, zipfile_importer)
        self.checks = RegistryLevelChecks(self)
        self.registry_code = None

    def do_import(self):
        self.checks.check_registry_export_type_in_meta()
        self.registry_code = get_meta_value(self.meta, 'registry.code')
        self.checks.check_registry_does_not_exist()

        # TODO We should probably do all of this in a transaction as below.
        # However, this would work only for the RDBMS and Mongo data changes
        # won't be rolled back. Therefore we could end up with some Mongo data
        # being inserted, but the relational data being rolled back.
        # Given that we can't avoid ending up with inconsistent data, I'm
        # leaving this off for now.
        #
        # with transaction.atomic():
        #    self.import_datagroups(self.registry_code, get_meta_value(self.meta, 'data_groups'))

        # TODO implement fake transactions by keeping a record of mongo documents inserted
        # and deleting them all on exception
        self.import_datagroups(get_meta_value(self.meta, 'data_groups'))

    def output_import_info(self):
        logger = BaseImporter.output_import_info(self)
        registry = get_meta_value(self.meta, 'registry')
        logger.debug('Registry: %s (%s v%s)', registry.get('name'), registry.get('code'), registry.get('version'))
        logger.debug(' ' * len('Registry: ') + '%s', registry.get('description'))
        logger.debug('Exported at: %s', dateutil.parser.parse(get_meta_value(self.meta, 'exported_at')))

META_FILTERS = {
        definitions.EXPORT_TYPES.CDES: lambda m: m['name'] == 'CDEs',
        definitions.EXPORT_TYPES.REFDATA: lambda m: m['name'] == 'Reference Data',
        definitions.EXPORT_TYPES.REGISTRY_DEF: lambda m: m['name'] != 'Registry Data',
}

class GenericImporter(BaseImporter):
    def do_import(self):
        with transaction.atomic():
            self.import_datagroups(get_meta_value(self.meta, 'data_groups'))
<commit_msg>RDR-1420: Reset DB sequences after import<commit_after>import contextlib
import dateutil.parser
import json
import logging
import os
import shutil
from StringIO import StringIO
import tempfile
from zipfile import ZipFile

from django.core.management import call_command
from django.db import transaction, connection

from rdrf.models import Registry
from .catalogue import DataGroupImporterCatalogue, ModelImporterCatalogue, MongoCollectionImporterCatalogue
from .importers import get_meta_value, allow_if_forced
import definitions
from .utils import DelegateMixin, IndentedLogger, app_schema_version

logger = logging.getLogger(__name__)

@contextlib.contextmanager
def zipfile_contents(zipfile):
    tmpdir = tempfile.mkdtemp(suffix='import')
    archive = ZipFile(zipfile)
    archive.extractall(tmpdir)

    yield tmpdir

    shutil.rmtree(tmpdir)

class ZipFileImporter(object):
    def __init__(self, zipfile, catalogue=None):
        if catalogue is None:
            catalogue = definitions.Catalogue(
                    DataGroupImporterCatalogue(),
                    ModelImporterCatalogue(),
                    MongoCollectionImporterCatalogue())
        self.catalogue = catalogue
        self.zipfile = zipfile
        self.workdir = None
        self.meta = None
        self.logger = logging.getLogger(__name__)

    def find_workdir(self, startdir):
        """Return the workdir and the meta file location in the archive.

        The zip file can contain the META file at top level or in a nested
        directory. We find the META file and the directory containing it will
        be the working directory."""

        meta = [(path, os.path.join(path, 'META')) for path, _, files in os.walk(startdir) if 'META' in files]
        if len(meta) == 0:
            raise Exception("Invalid export file '%s'.'META' file is missing." % self.zipfile)

        return meta[0]

    def extract_meta_info(self, meta_file):
        with open(meta_file) as f:
            self.meta = json.load(f)

    def create_importer(self, export_type, requested_import_type=None):
        self.file_export_type = definitions.EXPORT_TYPES.from_name(export_type)
        if requested_import_type is None:
            requested_import_type = self.file_export_type.code

        self.requested_type = definitions.EXPORT_TYPES.from_code(requested_import_type)
        if not (self.requested_type is self.file_export_type or self.requested_type in self.file_export_type.includes):
            raise Exception("Invalid import type '%s' requested for file '%s' with type '%s'." %
                    (requested_import_type, self.zipfile, self.file_export_type.code))

        if self.requested_type in definitions.EXPORT_TYPES.registry_types:
            return RegistryImporter(self)
        if self.requested_type is definitions.EXPORT_TYPES.CDES:
            return GenericImporter(self)
        if self.requested_type is definitions.EXPORT_TYPES.REFDATA:
            return GenericImporter(self)
        raise Exception("Unrecognized export type '%s'." % (self.requested_type.code))

    def do_import(self, import_type=None, verbose=False, indented_logs=True, simulate=False, force=False):
        if verbose:
            self.logger.setLevel(logging.DEBUG)
        self.child_logger = logger
        if indented_logs:
            self.child_logger = IndentedLogger(self.logger)
        self.simulate = simulate
        self.force = force

        with zipfile_contents(self.zipfile) as tmpdir:
            self.workdir, meta_file = self.find_workdir(tmpdir)
            self.extract_meta_info(meta_file)

            importer = self.create_importer(get_meta_value(self.meta, 'type'), requested_import_type=import_type)
            self.logger.debug('  Import type: %s (%s)', self.requested_type.name, self.requested_type.code)
            importer.output_import_info()
            importer.do_import()

    def inspect(self):
        self.logger.setLevel(logging.DEBUG)
        with zipfile_contents(self.zipfile) as tmpdir:
            self.workdir, meta_file = self.find_workdir(tmpdir)
            self.extract_meta_info(meta_file)
            importer = self.create_importer(get_meta_value(self.meta, 'type'))
            importer.output_import_info()

class RegistryLevelChecks(DelegateMixin):
    def __init__(self, importer):
        DelegateMixin.__init__(self, delegate_to=importer)

    def check_registry_export_type_in_meta(self):
        export_type = get_meta_value(self.meta, 'type')
        if export_type not in definitions.EXPORT_TYPES.registry_types_names:
            raise Exception("Invalid export type '%s' for registry import. Should be one of '%s'."
                % (export_type, ', '.join(definitions.EXPORT_TYPES.registry_types_names)))

    @allow_if_forced
    def check_registry_does_not_exist(self):
        if self.force:
            return
        if Registry.objects.filter(code=self.registry_code).exists():
            raise Exception("Registry '%s' already exists." % self.registry_code)

class BaseImporter(DelegateMixin):
    def __init__(self, zipfile_importer):
        DelegateMixin.__init__(self, delegate_to=zipfile_importer)

    def maybe_filter_meta(self, meta):
        if self.file_export_type is self.requested_type:
            return meta
        else:
            return filter(definitions.META_FILTERS[self.requested_type], meta)

    @allow_if_forced
    def check_app_schema_versions_match(self):
        app_schema_version_different = self.diff_app_versions()
        if len(app_schema_version_different) > 0:
            raise Exception('Schema difference detected between your registry and the export file.'
                    ' App(s) with different schema: %s' % ', '.join(app_schema_version_different))

    def reset_sql_sequences(self):
        meta = self.maybe_filter_meta(get_meta_value(self.meta, 'data_groups'))
        apps = set(reduce(lambda d, x: d + x.get('app_versions', {}).keys(), meta, []))
        self.logger.debug('Apps we are resetting: %s', apps)

        os.environ['DJANGO_COLORS'] = 'nocolor'
        commands = StringIO()

        for app in apps:
            call_command('sqlsequencereset', app, stdout=commands)

        cursor = connection.cursor()
        s = commands.getvalue()
        self.logger.debug('SQL Reset commands: %s', s)
        # cursor.execute(commands.getvalue())
        cursor.execute(s)

    def import_datagroups(self, meta):
        meta = self.maybe_filter_meta(meta)
        self.check_app_schema_versions_match()
        self.logger.debug('Importing %d datagroups', len(meta))
        datagroup_importers = self.catalogue.datagroups
        for data_group_meta in meta:
            importer = datagroup_importers.get(data_group_meta['name'])(self.catalogue)
            options = {
                'logger': self.child_logger,
                'simulate': self.simulate,
                'force': self.force,
            }
            if hasattr(self, 'registry_code'):
                options['registry_code'] = self.registry_code
            importer.do_import(data_group_meta, self.workdir, **options)
        self.reset_sql_sequences()

    def output_import_info(self):
        logger = IndentedLogger(self.logger, indent_level=2)
        zipfile_type = definitions.EXPORT_TYPES.from_name(get_meta_value(self.meta, 'type'))
        logger.debug('Zipfile type: %s (%s)', zipfile_type.name, zipfile_type.code)
        also_includes = zipfile_type.includes
        if also_includes:
            logger.debug('(also includes import types: %s)' %
                    ', '.join(map(lambda t: "'%s' (%s)" % (t.name, t.code), also_includes)))

        app_schema_version_different = self.diff_app_versions()
        if len(app_schema_version_different) > 0:
            logger.warn('WARNING: Schema difference detected between your registry and the export file.'
                    'App(s): %s', ', '.join(app_schema_version_different))

        return logger

    def diff_app_versions(self):
        meta = self.maybe_filter_meta(get_meta_value(self.meta, 'data_groups'))
        app_versions = dict(reduce(lambda d, x: d + x.get('app_versions', {}).items(), meta, []))
        return filter(lambda app: app_versions[app] != app_schema_version(app), app_versions)

class RegistryImporter(BaseImporter):
    def __init__(self, zipfile_importer):
        BaseImporter.__init__(self, zipfile_importer)
        self.checks = RegistryLevelChecks(self)
        self.registry_code = None

    def do_import(self):
        self.checks.check_registry_export_type_in_meta()
        self.registry_code = get_meta_value(self.meta, 'registry.code')
        self.checks.check_registry_does_not_exist()

        # TODO We should probably do all of this in a transaction as below.
        # However, this would work only for the RDBMS and Mongo data changes
        # won't be rolled back. Therefore we could end up with some Mongo data
        # being inserted, but the relational data being rolled back.
        # Given that we can't avoid ending up with inconsistent data, I'm
        # leaving this off for now.
        #
        # with transaction.atomic():
        #    self.import_datagroups(self.registry_code, get_meta_value(self.meta, 'data_groups'))

        # TODO implement fake transactions by keeping a record of mongo documents inserted
        # and deleting them all on exception
        self.import_datagroups(get_meta_value(self.meta, 'data_groups'))

    def output_import_info(self):
        logger = BaseImporter.output_import_info(self)
        registry = get_meta_value(self.meta, 'registry')
        logger.debug('Registry: %s (%s v%s)', registry.get('name'), registry.get('code'), registry.get('version'))
        logger.debug(' ' * len('Registry: ') + '%s', registry.get('description'))
        logger.debug('Exported at: %s', dateutil.parser.parse(get_meta_value(self.meta, 'exported_at')))

META_FILTERS = {
        definitions.EXPORT_TYPES.CDES: lambda m: m['name'] == 'CDEs',
        definitions.EXPORT_TYPES.REFDATA: lambda m: m['name'] == 'Reference Data',
        definitions.EXPORT_TYPES.REGISTRY_DEF: lambda m: m['name'] != 'Registry Data',
}

class GenericImporter(BaseImporter):
    def do_import(self):
        with transaction.atomic():
            self.import_datagroups(get_meta_value(self.meta, 'data_groups'))
<|endoftext|>


<commit_before># coding=utf-8
#-------------------------------------------------------------------------------

#Author: Alexandre Manhaes Savio <alexsavio@gmail.com>
#Grupo de Inteligencia Computational <www.ehu.es/ccwintco>
#Universidad del Pais Vasco UPV/EHU
#
#2013, Alexandre Manhaes Savio
#Use this at your own risk!
#-------------------------------------------------------------------------------

import nibabel as nib
from nipy.io.nifti_ref import nifti2nipy
from nipy.core.reference.array_coords import ArrayCoordMap

def voxcoord_to_mm(cm, i, j, k):
    '''
    Parameters
    ----------
    cm: nipy.core.reference.coordinate_map.CoordinateMap

    i, j, k: floats
        Voxel coordinates

    Returns
    -------
    Triplet with real 3D world coordinates

    '''
    try:
        mm = cm([i, j, k])
    except Exception as exc:
        raise Exception('Error on converting coordinates.') from exc
    else:
        return mm

def mm_to_voxcoord(cm, x, y, z):
    '''
    Parameters
    ----------
    cm: nipy.core.reference.coordinate_map.CoordinateMap

    x, y, z: floats
        Physical coordinates

    Returns
    -------
    Triplet with 3D voxel coordinates
    '''
    try:
        vox = cm.inverse()([x, y, z])
    except Exception as exc:
        raise Exception('Error on converting coordinates') from exc
    else:
        return vox

def get_3D_coordmap(img):
    '''
    Gets a 3D CoordinateMap from img.

    Parameters
    ----------
    img: nib.Nifti1Image or nipy Image

    Returns
    -------
    nipy.core.reference.coordinate_map.CoordinateMap
    '''
    if isinstance(img, nib.Nifti1Image):
        img = nifti2nipy(img)

    if img.ndim == 4:
        from nipy.core.reference.coordinate_map import drop_io_dim
        cm = drop_io_dim(img.coordmap, 3)
    else:
        cm = img.coordmap

    return cm

def get_coordmap_array(coordmap, shape):
    '''
    See: http://nipy.org/nipy/stable/api/generated/nipy.core.reference.array_coords.html?highlight=grid#nipy.core.reference.array_coords.Grid
    '''
    return ArrayCoordMap(coordmap, shape)
<commit_msg>add voxspace_to_mmspace function<commit_after># coding=utf-8
#-------------------------------------------------------------------------------

#Author: Alexandre Manhaes Savio <alexsavio@gmail.com>
#Grupo de Inteligencia Computational <www.ehu.es/ccwintco>
#Universidad del Pais Vasco UPV/EHU
#
#2016, Alexandre Manhaes Savio
#Use this at your own risk!
#-------------------------------------------------------------------------------

import nibabel as nib
import numpy as np
from nipy.io.nifti_ref import nifti2nipy
from nipy.core.reference.array_coords import ArrayCoordMap

def voxspace_to_mmspace(img):
    """ Return a grid with coordinates in 3D physical space for `img`."""
    shape, affine = img.shape[:3], img.affine
    coords = np.array(np.meshgrid(*(range(i) for i in shape), indexing='ij'))
    coords = np.rollaxis(coords, 0, len(shape) + 1)
    mm_coords = nib.affines.apply_affine(affine, coords)

    return mm_coords

def voxcoord_to_mm(cm, i, j, k):
    '''
    Parameters
    ----------
    cm: nipy.core.reference.coordinate_map.CoordinateMap

    i, j, k: floats
        Voxel coordinates

    Returns
    -------
    Triplet with real 3D world coordinates

    '''
    try:
        mm = cm([i, j, k])
    except Exception as exc:
        raise Exception('Error on converting coordinates.') from exc
    else:
        return mm

def mm_to_voxcoord(cm, x, y, z):
    '''
    Parameters
    ----------
    cm: nipy.core.reference.coordinate_map.CoordinateMap

    x, y, z: floats
        Physical coordinates

    Returns
    -------
    Triplet with 3D voxel coordinates
    '''
    try:
        vox = cm.inverse()([x, y, z])
    except Exception as exc:
        raise Exception('Error on converting coordinates') from exc
    else:
        return vox

def get_3D_coordmap(img):
    '''
    Gets a 3D CoordinateMap from img.

    Parameters
    ----------
    img: nib.Nifti1Image or nipy Image

    Returns
    -------
    nipy.core.reference.coordinate_map.CoordinateMap
    '''
    if isinstance(img, nib.Nifti1Image):
        img = nifti2nipy(img)

    if img.ndim == 4:
        from nipy.core.reference.coordinate_map import drop_io_dim
        cm = drop_io_dim(img.coordmap, 3)
    else:
        cm = img.coordmap

    return cm

def get_coordmap_array(coordmap, shape):
    '''
    See: http://nipy.org/nipy/stable/api/generated/nipy.core.reference.array_coords.html?highlight=grid#nipy.core.reference.array_coords.Grid
    '''
    return ArrayCoordMap(coordmap, shape)
<|endoftext|>


<commit_before>#!env/python3
# coding: utf-8
try:
    import ipdb
except ImportError:
    pass

import os
import json
import aiohttp

import aiohttp_jinja2
import tarfile
import datetime
import time
import uuid

from aiohttp import web
from urllib.parse import parse_qsl

from config import *
from core.framework.common import *
from core.framework.tus import *
from core.model import *
from core.core import core
from api_rest.rest import *

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Customization of the TUS protocol for the download of pirus files
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File TUS wrapper
class FileWrapper (TusFileWrapper):
    def __init__(self, id):
        self.file = File.from_id(id)
        if self.file is not None:
            self.id = id
            self.name = self.file.name
            self.upload_offset = self.file.upload_offset
            self.path = self.file.path
            self.size = self.file.size
            self.upload_url = "file/upload/" + str(id)
        else:
            raise RegovarException("TUS File wrapper init error : Unknow id: {}".format(id))

    async def save(self):
        from core.core import core
        try:
            f = File.from_id(self.id)
            f.upload_offset=self.upload_offset
            f.save()
            await core.notify_all_co(data={"action": "file_upload", "data" : f.to_json(["id", "size", "upload_offset", "status", "name"])})
        except Exception as ex:
            return TusManager.build_response(code=500, body="Unexpected error occured: {}".format(ex))

    async def complete(self, checksum=None, checksum_type="md5"):
        try:
            log ('Upload of the file (id={0}) is complete.'.format(self.id))
            core.files.upload_finish(self.id, checksum, checksum_type)
            f = File.from_id(self.id)
            await core.notify_all_co(data={"action": "file_upload", "data" : f.to_json()})
        except Exception as ex:
            return TusManager.build_response(code=500, body="Unexpected error occured: {}".format(ex))

    @staticmethod
    def new_upload(request, filename, file_size):
        """
            Create and return the wrapper to manipulate the uploading file
        """
        pfile = core.files.upload_init(filename, file_size)
        return FileWrapper(pfile.id)

# set mapping
tus_manager.route_maping["/file/upload"] = FileWrapper

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# REST FILE API HANDLER
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

class FileHandler:

    #@staticmethod
    #def format_file_json(file_json):
        #if not isinstance(file_json, dict): return file_json
        #if "path" in file_json.keys():
            #path = file_json["path"]
            #if path and path.startswith(FILES_DIR):
                #file_json["path"] = "http://{}/dl/file{}".format(HOST_P, path[len(JOBS_DIR):])
            #elif path and path.startswith(JOBS_DIR):
                #file_json["path"] = "http://{}/dl/job{}".format(HOST_P, path[len(JOBS_DIR):])
        #return file_json

    @user_role('Authenticated')
    def list(self, request):
        # Generic processing of the get query
        fields, query, order, offset, limit = process_generic_get(request.query_string, File.public_fields)
        depth = request.query["depth"] if "depth" in request.query else 0
        # Get range meta data
        range_data = {
            "range_offset" : offset,
            "range_limit"  : limit,
            "range_total"  : File.count(),
            "range_max"    : RANGE_MAX,
        }
        # Return result of the query for PirusFile
        files = core.files.get(fields, query, order, offset, limit, depth)
        return rest_success(check_local_path([f.to_json() for f in files]), range_data)

    @user_role('Authenticated')
    async def edit(self, request):
        # TODO : implement PUT to edit file metadata (and remove the obsolete  "simple post" replaced by TUS upload )
        file_id = request.match_info.get('file_id', "")
        params = await request.json()
        return rest_error("Not yet implemented")

    @user_role('Authenticated')
    def delete(self, request):
        file_id = request.match_info.get('file_id', "")
        try:
            return rest_success(check_local_path(core.files.delete(file_id).to_json()))
        except Exception as ex:
            return rest_error("Error occured : " + str(ex))

    @user_role('Authenticated')
    def get(self, request):
        file_id = request.match_info.get('file_id', -1)
        file = File.from_id(file_id, 2)
        if not file:
            return rest_error("Unable to find the file (id={})".format(file_id))
        return rest_success(check_local_path(file.to_json(File.public_fields)))

    # Resumable download implement the TUS.IO protocol.
    def tus_config(self, request):
        return tus_manager.options(request)

    def tus_upload_init(self, request):
        return tus_manager.creation(request)

    def tus_upload_resume(self, request):
        return tus_manager.resume(request)

    async def tus_upload_chunk(self, request):
        result = await tus_manager.patch(request)
        return result

    def tus_upload_delete(self, request):
        return tus_manager.delete_file(request)

<commit_msg>fix file update<commit_after>#!env/python3
# coding: utf-8
try:
    import ipdb
except ImportError:
    pass

import os
import json
import aiohttp

import aiohttp_jinja2
import tarfile
import datetime
import time
import uuid

from aiohttp import web
from urllib.parse import parse_qsl

from config import *
from core.framework.common import *
from core.framework.tus import *
from core.model import *
from core.core import core
from api_rest.rest import *

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Customization of the TUS protocol for the download of files
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File TUS wrapper
class FileWrapper (TusFileWrapper):
    def __init__(self, id):
        self.file = File.from_id(id)
        if self.file is not None:
            self.id = id
            self.name = self.file.name
            self.upload_offset = self.file.upload_offset
            self.path = self.file.path
            self.size = self.file.size
            self.upload_url = "file/upload/" + str(id)
        else:
            raise RegovarException("TUS File wrapper init error : Unknow id: {}".format(id))

    async def save(self):
        from core.core import core
        try:
            f = File.from_id(self.id)
            f.upload_offset=self.upload_offset
            f.save()
            await core.notify_all_co(data={"action": "file_upload", "data" : f.to_json(["id", "size", "upload_offset", "status", "name"])})
        except Exception as ex:
            return TusManager.build_response(code=500, body="Unexpected error occured: {}".format(ex))

    async def complete(self, checksum=None, checksum_type="md5"):
        try:
            log ('Upload of the file (id={0}) is complete.'.format(self.id))
            core.files.upload_finish(self.id, checksum, checksum_type)
            f = File.from_id(self.id)
            await core.notify_all_co(data={"action": "file_upload", "data" : f.to_json()})
        except Exception as ex:
            return TusManager.build_response(code=500, body="Unexpected error occured: {}".format(ex))

    @staticmethod
    def new_upload(request, filename, file_size):
        """
            Create and return the wrapper to manipulate the uploading file
        """
        pfile = core.files.upload_init(filename, file_size)
        return FileWrapper(pfile.id)

# set mapping
tus_manager.route_maping["/file/upload"] = FileWrapper

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# REST FILE API HANDLER
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

class FileHandler:

    @user_role('Authenticated')
    def list(self, request):
        # Generic processing of the get query
        fields, query, order, offset, limit = process_generic_get(request.query_string, File.public_fields)
        depth = request.query["depth"] if "depth" in request.query else 0
        # Get range meta data
        range_data = {
            "range_offset" : offset,
            "range_limit"  : limit,
            "range_total"  : File.count(),
            "range_max"    : RANGE_MAX,
        }
        # Return result of the query for PirusFile
        files = core.files.get(fields, query, order, offset, limit, depth)
        return rest_success(check_local_path([f.to_json() for f in files]), range_data)

    @user_role('Authenticated')
    async def edit(self, request):
        file_id = request.match_info.get('file_id', "")
        data = await request.json()
        file = File.from_id(file_id, 2)
        if not file:
            return rest_error("Unable to find the file (id={})".format(file_id))
        try:
            file.load(data)
        except Exception as ex:
            return rest_error("Unable to update file with provided information.", exception=ex)
        return  rest_success(check_local_path(file.to_json(File.public_fields)))

    @user_role('Authenticated')
    def delete(self, request):
        file_id = request.match_info.get('file_id', "")
        try:
            return rest_success(check_local_path(core.files.delete(file_id).to_json()))
        except Exception as ex:
            return rest_error("Unable to delete file {}".format(file_id), exception=ex)

    @user_role('Authenticated')
    def get(self, request):
        file_id = request.match_info.get('file_id', -1)
        file = File.from_id(file_id, 2)
        if not file:
            return rest_error("Unable to find the file (id={})".format(file_id))
        return rest_success(check_local_path(file.to_json(File.public_fields)))

    # Resumable download implement the TUS.IO protocol.
    def tus_config(self, request):
        return tus_manager.options(request)

    def tus_upload_init(self, request):
        return tus_manager.creation(request)

    def tus_upload_resume(self, request):
        return tus_manager.resume(request)

    async def tus_upload_chunk(self, request):
        result = await tus_manager.patch(request)
        return result

    def tus_upload_delete(self, request):
        return tus_manager.delete_file(request)

<|endoftext|>


<commit_before>import urllib
import glob
import pandas as pd
from bquery import ctable
# do not forget to install numexpr
workdir = '/home/carst/Desktop/taxi/'

# wget https://storage.googleapis.com/tlc-trip-data/2015/yellow_tripdata_2015-01.csv
for year in [2014, 2015]:
    for month in range(1, 13):
        filename = 'yellow_tripdata_' + str(year) + '-' + ('0' + str(month))[-2:] + '.csv'
        url = 'https://storage.googleapis.com/tlc-trip-data/' + str(year) + '/' + filename
        print(url)
        urllib.urlretrieve(url, workdir + '/' + filename)

file_list = sorted(glob.glob(workdir + 'yellow_tripdata_*.csv'))
if not file_list:
    raise ValueError('No Files Found')

expected_len = len(file_list) * 20*10**6

for i, filename in enumerate(file_list):

    print(filename)

    import_df = pd.read_csv(filename)

    # lower columns because of input inconsistencies
    import_df.columns = [x.lower() for x in import_df.columns]
    import_df.columns = [x.strip() for x in import_df.columns]
    import_df.columns = [x.replace('tpep_', '') for x in import_df.columns]

    import_df['nr_rides'] = 1

    import_df['pickup_date'] = import_df['pickup_datetime'].str[0:10]
    import_df['pickup_date'] = import_df['pickup_date'].str.replace('-', '')
    import_df['pickup_year'] = import_df['pickup_date'].str[0:4].astype(int)
    import_df['pickup_yearmonth'] = import_df['pickup_date'].str[0:6].astype(int)
    import_df['pickup_month'] = import_df['pickup_date'].str[4:6].astype(int)
    import_df['pickup_date'] = import_df['pickup_date'].astype(int)
    import_df['pickup_time'] = import_df['pickup_datetime'].str[11:]
    import_df['pickup_time'] = import_df['pickup_time'].str.replace(':', '')
    import_df['pickup_hour'] = import_df['pickup_time'].str[0:2].astype(int)
    import_df['pickup_time'] = import_df['pickup_time'].astype(int)
    del import_df['pickup_datetime']

    import_df['dropoff_date'] = import_df['dropoff_datetime'].str[0:10]
    import_df['dropoff_date'] = import_df['dropoff_date'].str.replace('-', '')
    import_df['dropoff_year'] = import_df['dropoff_date'].str[0:4].astype(int)
    import_df['dropoff_yearmonth'] = import_df['dropoff_date'].str[0:6].astype(int)
    import_df['dropoff_month'] = import_df['dropoff_date'].str[4:6].astype(int)
    import_df['dropoff_date'] = import_df['dropoff_date'].astype(int)
    import_df['dropoff_time'] = import_df['dropoff_datetime'].str[11:]
    import_df['dropoff_time'] = import_df['dropoff_time'].str.replace(':', '')
    import_df['dropoff_hour'] = import_df['dropoff_time'].str[0:2].astype(int)
    import_df['dropoff_time'] = import_df['dropoff_time'].astype(int)
    del import_df['dropoff_datetime']

    if i == 0:
        import_ct = ctable.fromdataframe(import_df, rootdir=workdir+'taxi', expectedlen=expected_len, mode='w')
        del import_df
    else:
        temp_ct = ctable.fromdataframe(import_df)
        import_ct.append(temp_ct)
        del temp_ct
        del import_df

import_ct.flush()<commit_msg>Comparison branch  for the taxi data set<commit_after>import urllib
import glob
import pandas as pd
from bquery import ctable
import bquery
import bcolz

# do not forget to install numexpr
workdir = '/home/carst/Desktop/taxi/'

# wget https://storage.googleapis.com/tlc-trip-data/2015/yellow_tripdata_2015-01.csv
for year in [2015]:
    for month in range(1, 13):
        filename = 'yellow_tripdata_' + str(year) + '-' + ('0' + str(month))[-2:] + '.csv'
        url = 'https://storage.googleapis.com/tlc-trip-data/' + str(year) + '/' + filename
        print(url)
        urllib.urlretrieve(url, workdir + '/' + filename)

file_list = sorted(glob.glob(workdir + 'yellow_tripdata_*.csv'))
if not file_list:
    raise ValueError('No Files Found')

expected_len = len(file_list) * 50000000

for i, filename in enumerate(file_list):

    print(filename)

    import_df = pd.read_csv(filename)

    # lower columns because of input inconsistencies
    import_df.columns = [x.lower() for x in import_df.columns]
    import_df.columns = [x.strip() for x in import_df.columns]
    import_df.columns = [x.replace('tpep_', '') for x in import_df.columns]

    import_df['nr_rides'] = 1

    import_df['pickup_date'] = import_df['pickup_datetime'].str[0:10]
    import_df['pickup_date'] = import_df['pickup_date'].str.replace('-', '')
    import_df['pickup_year'] = import_df['pickup_date'].str[0:4].astype(int)
    import_df['pickup_yearmonth'] = import_df['pickup_date'].str[0:6].astype(int)
    import_df['pickup_month'] = import_df['pickup_date'].str[4:6].astype(int)
    import_df['pickup_date'] = import_df['pickup_date'].astype(int)
    import_df['pickup_time'] = import_df['pickup_datetime'].str[11:]
    import_df['pickup_time'] = import_df['pickup_time'].str.replace(':', '')
    import_df['pickup_hour'] = import_df['pickup_time'].str[0:2].astype(int)
    import_df['pickup_time'] = import_df['pickup_time'].astype(int)
    del import_df['pickup_datetime']

    import_df['dropoff_date'] = import_df['dropoff_datetime'].str[0:10]
    import_df['dropoff_date'] = import_df['dropoff_date'].str.replace('-', '')
    import_df['dropoff_year'] = import_df['dropoff_date'].str[0:4].astype(int)
    import_df['dropoff_yearmonth'] = import_df['dropoff_date'].str[0:6].astype(int)
    import_df['dropoff_month'] = import_df['dropoff_date'].str[4:6].astype(int)
    import_df['dropoff_date'] = import_df['dropoff_date'].astype(int)
    import_df['dropoff_time'] = import_df['dropoff_datetime'].str[11:]
    import_df['dropoff_time'] = import_df['dropoff_time'].str.replace(':', '')
    import_df['dropoff_hour'] = import_df['dropoff_time'].str[0:2].astype(int)
    import_df['dropoff_time'] = import_df['dropoff_time'].astype(int)
    del import_df['dropoff_datetime']

    if i == 0:
        import_ct = ctable.fromdataframe(import_df, rootdir=workdir + 'taxi', expectedlen=expected_len, mode='w')
        del import_df
    else:
        temp_ct = ctable.fromdataframe(import_df)
        import_ct.append(temp_ct)
        del temp_ct
        del import_import_ct = ctable.fromdataframe(import_df, rootdir=workdir + 'taxi', expectedlen=expected_len,
                                                    mode='w')
        df

import_ct.flush()

import_ct = ctable(rootdir=workdir + 'taxi', mode='a')

import_ct.cache_factor([
    'dropoff_date',
    'dropoff_hour',
    'dropoff_latitude',
    'dropoff_longitude',
    'dropoff_month',
    'dropoff_time',
    'dropoff_year',
    'dropoff_yearmonth',
    'payment_type',
    'pickup_date',
    'pickup_hour',
    'pickup_latitude',
    'pickup_longitude',
    'pickup_month',
    'pickup_time',
    'pickup_year',
    'pickup_yearmonth',
    'ratecodeid',
    'store_and_fwd_flag',
    'vendorid'])

measure_list = ['extra',
                'fare_amount',
                'improvement_surcharge',
                'mta_tax',
                'nr_rides',
                'passenger_count',
                'tip_amount',
                'tolls_amount',
                'total_amount',
                'trip_distance']

bquery.set_nthreads(bquery.ncores)

import_ct.groupby(['pickup_yearmonth'], ['nr_rides'])
%timeit import_ct.groupby(['pickup_yearmonth'], ['nr_rides'])

# In [22]: %timeit import_ct.groupby(['pickup_yearmonth'], ['nr_rides'])
# 1 loop, best of 3: 6.91 s per loop

 %time nyc2015.payment_type.value_counts().compute()

import_ct.groupby(['payment_type'], ['nr_rides'])
%time import_ct.groupby(['payment_type'], ['nr_rides'])

In [26]: %time import_ct.groupby(['payment_type'], ['nr_rides'])
CPU times: user 11.3 s, sys: 284 ms, total: 11.5 s
Wall time: 7.41 s

compared to dask winine m3.2xlarge nodes on EC2. These have eight cores and 30GB of RAM each
>>> %time nyc2015.payment_type.value_counts().compute()
CPU times: user 132 ms, sys: 0 ns, total: 132 ms
Wall time: 558 ms

13.8 times slower but without parallelization and only 84mb resident + 400mb virtual memory usage
<|endoftext|>


<commit_before>import numpy as np
from brew import Ensemble

class EPIC(object):

    def __init__(self):
        pass

    def fit(self, ensemble, X, y, classes):
        classifiers = ensemble.classifiers
        self.classes_ = classes
        V = np.zeros((y.shape[0], len(self.classes_)))
        C = classifiers

        index_table = {}
        i = 0
        for classe in self.classes_:
            index_table[classe] = i
            i = i + 1

        for c_i in C:
            for i in range(y.shape[0]):
                prediction = c_i.predict(X[i])[0]
                index = index_table[prediction]
                V[i][index] = V[i][index] + 1

        OL = []
        for c_i in C:
            IC = 0.0
            for j in range(y.shape[0]):
                pred = c_i.predict(X[j])[0]
                index_pred = index_table[pred]

                alpha, beta, gamma = 0, 0, 1
                if index_pred == index_table[y[j]]:
                    alpha = 1 if index_pred == np.argmin(V[j]) else 0
                    beta = 1 if index_pred == np.argmax(V[j]) else 0
                    gamma = 0

                v_max = np.argmax(V[j])
                v_sec = sorted(np.array(V[j]))[len(V[j]) - 2]
                v_cor = V[j][index_table[y[j]]]

                IC = IC + alpha * (2*v_max - V[j][index_pred]) + beta * v_sec + gamma * (v_cor - V[j][index_pred] - v_max)

            OL = OL + [[c_i, IC]]

        OL = sorted(OL, key=lambda e: e[1], reverse=True)
        self.ensemble = list(zip(*OL)[0])
        return self

    def get(self, p=0.1):
        return self.ensemble[:int(p*len(self.ensemble))]

<commit_msg>epic updated<commit_after># -*- coding: utf-8 -*-

import numpy as np

from brew.base import Ensemble
from .base import Prunner

class EPIC(Prunner):

    def __init__(self):
        pass

    def fit(self, ensemble, X, y):
        classifiers = ensemble.classifiers
        self.classes_ = list(set(y))
        V = np.zeros((y.shape[0], len(self.classes_)))
        C = classifiers

        index_table = {}
        i = 0
        for classe in self.classes_:
            index_table[classe] = i
            i = i + 1

        for c_i in C:
            for i in range(y.shape[0]):
                prediction = c_i.predict(X[i])[0]
                index = index_table[prediction]
                V[i][index] = V[i][index] + 1

        OL = []
        for c_i in C:
            IC = 0.0
            for j in range(y.shape[0]):
                pred = c_i.predict(X[j])[0]
                index_pred = index_table[pred]

                alpha, beta, gamma = 0, 0, 1
                if index_pred == index_table[y[j]]:
                    alpha = 1 if index_pred == np.argmin(V[j]) else 0
                    beta = 1 if index_pred == np.argmax(V[j]) else 0
                    gamma = 0

                v_max = np.argmax(V[j])
                v_sec = sorted(np.array(V[j]))[len(V[j]) - 2]
                v_cor = V[j][index_table[y[j]]]

                IC = IC + alpha * (2*v_max - V[j][index_pred]) + beta * v_sec + gamma * (v_cor - V[j][index_pred] - v_max)

            OL = OL + [[c_i, IC]]

        OL = sorted(OL, key=lambda e: e[1], reverse=True)
        self.classifiers = list(zip(*OL)[0])
        return self

    def get(self, p=0.1):
        return self.classifiers[:int(p*len(self.classifiers))]

<|endoftext|>