<commit_before># -*- encoding: utf-8 -*-

#
# Copyright P. Christeas <p_christ@hol.gr> 2008,2009
#
#
# WARNING: This program as such is intended to be used by professional
# programmers who take the whole responsability of assessing all potential
# consequences resulting from its eventual inadequacies and bugs
# End users who are looking for a ready-to-use solution with commercial
# garantees and support are strongly adviced to contract a Free Software
# Service Company
#
# This program is Free Software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
###############################################################################

""" This file contains instance of the http server.

"""
from websrv_lib import *
import netsvc
import threading
import tools
import os
import socket
import xmlrpclib

from SimpleXMLRPCServer import SimpleXMLRPCDispatcher

try:
    import fcntl
except ImportError:
    fcntl = None

try:
    from ssl import SSLError
except ImportError:
    class SSLError(Exception): pass

class ThreadedHTTPServer(ConnThreadingMixIn, SimpleXMLRPCDispatcher, HTTPServer):
    """ A threaded httpd server, with all the necessary functionality for us.

        It also inherits the xml-rpc dispatcher, so that some xml-rpc functions
        will be available to the request handler
    """
    encoding = None
    allow_none = False
    allow_reuse_address = 1
    _send_traceback_header = False
    i = 0

    def __init__(self, addr, requestHandler,
                 logRequests=True, allow_none=False, encoding=None, bind_and_activate=True):
        self.logRequests = logRequests

        SimpleXMLRPCDispatcher.__init__(self, allow_none, encoding)
        HTTPServer.__init__(self, addr, requestHandler)

        # [Bug #1222790] If possible, set close-on-exec flag; if a
        # method spawns a subprocess, the subprocess shouldn't have
        # the listening socket open.
        if fcntl is not None and hasattr(fcntl, 'FD_CLOEXEC'):
            flags = fcntl.fcntl(self.fileno(), fcntl.F_GETFD)
            flags |= fcntl.FD_CLOEXEC
            fcntl.fcntl(self.fileno(), fcntl.F_SETFD, flags)

    def handle_error(self, request, client_address):
        """ Override the error handler
        """
        import traceback
        netsvc.Logger().notifyChannel("init", netsvc.LOG_ERROR,"Server error in request from %s:\n%s" %
            (client_address,traceback.format_exc()))

class MultiHandler2(MultiHTTPHandler):
    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_DEBUG,format % args)

    def log_error(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_ERROR,format % args)

class SecureMultiHandler2(SecureMultiHTTPHandler):
    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('https',netsvc.LOG_DEBUG,format % args)

    def getcert_fnames(self):
        tc = tools.config
        fcert = tc.get_misc('httpsd','sslcert', 'ssl/server.cert')
        fkey = tc.get_misc('httpsd','sslkey', 'ssl/server.key')
        return (fcert,fkey)

    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_DEBUG,format % args)

    def log_error(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_ERROR,format % args)

class HttpDaemon(threading.Thread, netsvc.Server):
    def __init__(self, interface, port):
        threading.Thread.__init__(self)
        netsvc.Server.__init__(self)
        self.__port = port
        self.__interface = interface

        try:
            self.server = ThreadedHTTPServer((interface, port), MultiHandler2)
            self.server.vdirs = []
            self.server.logRequests = True
            netsvc.Logger().notifyChannel("web-services", netsvc.LOG_INFO,
                "starting HTTP service at %s port %d" % (interface or '0.0.0.0', port,))
        except Exception, e:
            netsvc.Logger().notifyChannel('httpd', netsvc.LOG_CRITICAL, "Error occur when starting the server daemon: %s" % (e,))
            raise

    def attach(self, path, gw):
        pass

    def stop(self):
        self.running = False
        if os.name != 'nt':
            self.server.socket.shutdown( hasattr(socket, 'SHUT_RDWR') and socket.SHUT_RDWR or 2 )
        self.server.socket.close()

    def run(self):
        #self.server.register_introspection_functions()

        self.running = True
        while self.running:
            self.server.handle_request()
        return True

class HttpSDaemon(threading.Thread, netsvc.Server):
    def __init__(self, interface, port):
        threading.Thread.__init__(self)
        netsvc.Server.__init__(self)
        self.__port = port
        self.__interface = interface

        try:
            self.server = ThreadedHTTPServer((interface, port), SecureMultiHandler2)
            self.server.vdirs = []
            self.server.logRequests = True
            netsvc.Logger().notifyChannel("web-services", netsvc.LOG_INFO,
                "starting HTTPS service at %s port %d" % (interface or '0.0.0.0', port,))
        except SSLError, e:
            netsvc.Logger().notifyChannel('httpd-ssl', netsvc.LOG_CRITICAL, "Can not load the certificate and/or the private key files")
            raise
        except Exception, e:
            netsvc.Logger().notifyChannel('httpd-ssl', netsvc.LOG_CRITICAL, "Error occur when starting the server daemon: %s" % (e,))
            raise

    def attach(self, path, gw):
        pass

    def stop(self):
        self.running = False
        if os.name != 'nt':
            self.server.socket.shutdown( hasattr(socket, 'SHUT_RDWR') and socket.SHUT_RDWR or 2 )
        self.server.socket.close()

    def run(self):
        #self.server.register_introspection_functions()

        self.running = True
        while self.running:
            self.server.handle_request()
        return True

httpd = None
httpsd = None

def init_servers():
    global httpd, httpsd
    if tools.config.get_misc('httpd','enable', True):
        httpd = HttpDaemon(tools.config.get_misc('httpd','interface', ''), \
            int(tools.config.get_misc('httpd','port', tools.config.get('port',8069))))

    if tools.config.get_misc('httpsd','enable', False):
        httpsd = HttpSDaemon(tools.config.get_misc('httpsd','interface', ''), \
            int(tools.config.get_misc('httpsd','port', 8071)))

def reg_http_service(hts, secure_only = False):
    """ Register some handler to httpd.
        hts must be an HTTPDir
    """
    global httpd, httpsd
    if not isinstance(hts, HTTPDir):
        raise Exception("Wrong class for http service")

    if httpd and not secure_only:
        httpd.server.vdirs.append(hts)

    if httpsd:
        httpsd.server.vdirs.append(hts)

    if (not httpd) and (not httpsd):
        netsvc.Logger().notifyChannel('httpd',netsvc.LOG_WARNING,"No httpd available to register service %s" % hts.path)
    return

import SimpleXMLRPCServer
class XMLRPCRequestHandler(netsvc.OpenERPDispatcher,FixSendError,SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
    rpc_paths = []
    protocol_version = 'HTTP/1.1'
    def _dispatch(self, method, params):
        try:
            service_name = self.path.split("/")[-1]
            return self.dispatch(service_name, method, params)
        except netsvc.OpenERPDispatcherException, e:
            raise xmlrpclib.Fault(tools.exception_to_unicode(e.exception), e.traceback)

    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('xmlrpc',netsvc.LOG_DEBUG_RPC,format % args)

    def handle(self):
        pass

    def finish(self):
        pass

    def setup(self):
        self.connection = dummyconn()
        if not len(XMLRPCRequestHandler.rpc_paths):
            XMLRPCRequestHandler.rpc_paths = map(lambda s: '/%s' % s, netsvc.ExportService._services.keys())
        pass

def init_xmlrpc():
    if not tools.config.get_misc('xmlrpc','enable', True):
        return
    reg_http_service(HTTPDir('/xmlrpc/',XMLRPCRequestHandler))
    # Example of http file serving:
    # reg_http_service(HTTPDir('/test/',HTTPHandler))
    netsvc.Logger().notifyChannel("web-services", netsvc.LOG_INFO,
            "Registered XML-RPC over HTTP")

class OerpAuthProxy(AuthProxy):
    """ Require basic authentication..

        This is a copy of the BasicAuthProxy, which however checks/caches the db
        as well.
    """
    def __init__(self,provider):
        AuthProxy.__init__(self,provider)
        self.auth_creds = {}
        self.auth_tries = 0
        self.last_auth = None

    def checkRequest(self,handler,path = '/'):
        if self.auth_creds:
            return True
        auth_str = handler.headers.get('Authorization',False)
        try:
            db = handler.get_db_from_path(path)
            print "Got db:",db
        except:
            if path.startswith('/'):
                path = path[1:]
            psp= path.split('/')
            if len(psp)>1:
                db = psp[0]
            else:
                #FIXME!
                self.provider.log("Wrong path: %s, failing auth" %path)
                raise AuthRejectedExc("Authorization failed. Wrong sub-path.")

        if auth_str and auth_str.startswith('Basic '):
            auth_str=auth_str[len('Basic '):]
            (user,passwd) = base64.decodestring(auth_str).split(':')
            self.provider.log("Found user=\"%s\", passwd=\"***\" for db=\"%s\"" %(user,db))
            acd = self.provider.authenticate(db,user,passwd,handler.client_address)
            if acd != False:
                self.auth_creds[db] = acd
                self.last_auth=db
                return True
        if self.auth_tries > 5:
            self.provider.log("Failing authorization after 5 requests w/o password")
            raise AuthRejectedExc("Authorization failed.")
        self.auth_tries += 1
        raise AuthRequiredExc(atype = 'Basic', realm=self.provider.realm)

import security
class OpenERPAuthProvider(AuthProvider):
    def __init__(self,realm = 'OpenERP User'):
        self.realm = realm

    def setupAuth(self, multi, handler):
        if not multi.sec_realms.has_key(self.realm):
            multi.sec_realms[self.realm] = OerpAuthProxy(self)
        handler.auth_proxy = multi.sec_realms[self.realm]

    def authenticate(self, db, user, passwd, client_address):
        try:
            uid = security.login(db,user,passwd)
            if uid is False:
                return False
            return (user, passwd, db, uid)
        except Exception,e:
            netsvc.Logger().notifyChannel("auth",netsvc.LOG_DEBUG,"Fail auth:"+ str(e))
            return False

    def log(self, msg):
        netsvc.Logger().notifyChannel("auth",netsvc.LOG_INFO,msg)

#eof
<commit_msg>[fix] actually fix hang in OSX when trying to shutdown server via SIGINT lost during moves/merges from 5.0 to trunk<commit_after># -*- encoding: utf-8 -*-

#
# Copyright P. Christeas <p_christ@hol.gr> 2008,2009
#
#
# WARNING: This program as such is intended to be used by professional
# programmers who take the whole responsability of assessing all potential
# consequences resulting from its eventual inadequacies and bugs
# End users who are looking for a ready-to-use solution with commercial
# garantees and support are strongly adviced to contract a Free Software
# Service Company
#
# This program is Free Software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
###############################################################################

""" This file contains instance of the http server.

"""
from websrv_lib import *
import netsvc
import threading
import tools
import os
import socket
import xmlrpclib

from SimpleXMLRPCServer import SimpleXMLRPCDispatcher

try:
    import fcntl
except ImportError:
    fcntl = None

try:
    from ssl import SSLError
except ImportError:
    class SSLError(Exception): pass

class ThreadedHTTPServer(ConnThreadingMixIn, SimpleXMLRPCDispatcher, HTTPServer):
    """ A threaded httpd server, with all the necessary functionality for us.

        It also inherits the xml-rpc dispatcher, so that some xml-rpc functions
        will be available to the request handler
    """
    encoding = None
    allow_none = False
    allow_reuse_address = 1
    _send_traceback_header = False
    i = 0

    def __init__(self, addr, requestHandler,
                 logRequests=True, allow_none=False, encoding=None, bind_and_activate=True):
        self.logRequests = logRequests

        SimpleXMLRPCDispatcher.__init__(self, allow_none, encoding)
        HTTPServer.__init__(self, addr, requestHandler)

        # [Bug #1222790] If possible, set close-on-exec flag; if a
        # method spawns a subprocess, the subprocess shouldn't have
        # the listening socket open.
        if fcntl is not None and hasattr(fcntl, 'FD_CLOEXEC'):
            flags = fcntl.fcntl(self.fileno(), fcntl.F_GETFD)
            flags |= fcntl.FD_CLOEXEC
            fcntl.fcntl(self.fileno(), fcntl.F_SETFD, flags)

    def handle_error(self, request, client_address):
        """ Override the error handler
        """
        import traceback
        netsvc.Logger().notifyChannel("init", netsvc.LOG_ERROR,"Server error in request from %s:\n%s" %
            (client_address,traceback.format_exc()))

class MultiHandler2(MultiHTTPHandler):
    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_DEBUG,format % args)

    def log_error(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_ERROR,format % args)

class SecureMultiHandler2(SecureMultiHTTPHandler):
    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('https',netsvc.LOG_DEBUG,format % args)

    def getcert_fnames(self):
        tc = tools.config
        fcert = tc.get_misc('httpsd','sslcert', 'ssl/server.cert')
        fkey = tc.get_misc('httpsd','sslkey', 'ssl/server.key')
        return (fcert,fkey)

    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_DEBUG,format % args)

    def log_error(self, format, *args):
        netsvc.Logger().notifyChannel('http',netsvc.LOG_ERROR,format % args)

class BaseHttpDaemon(threading.Thread, netsvc.Server):
    def __init__(self, interface, port, handler):
        threading.Thread.__init__(self)
        netsvc.Server.__init__(self)
        self.__port = port
        self.__interface = interface

        try:
            self.server = ThreadedHTTPServer((interface, port), handler)
            self.server.vdirs = []
            self.server.logRequests = True
            netsvc.Logger().notifyChannel(
                "web-services", netsvc.LOG_INFO,
                "starting HTTPS service at %s port %d" %
                (interface or '0.0.0.0', port,))
        except Exception, e:
            netsvc.Logger().notifyChannel(
                'httpd', netsvc.LOG_CRITICAL,
                "Error occur when starting the server daemon: %s" % (e,))
            raise

    def attach(self, path, gw):
        pass

    def stop(self):
        self.running = False
        if os.name != 'nt':
            try:
                self.server.socket.shutdown(
                    hasattr(socket, 'SHUT_RDWR') and socket.SHUT_RDWR or 2)
            except socket.error, e:
                if e.errno != 57: raise
                # OSX, socket shutdowns both sides if any side closes it
                # causing an error 57 'Socket is not connected' on shutdown
                # of the other side (or something), see
                # http://bugs.python.org/issue4397
                netsvc.Logger().notifyChannel(
                    'server', netsvc.LOG_DEBUG,
                    '"%s" when shutting down server socket, '
                    'this is normal under OS X'%e)
        self.server.socket.close()

    def run(self):
        self.running = True
        while self.running:
            self.server.handle_request()
        return True

class HttpDaemon(BaseHttpDaemon):
    def __init__(self, interface, port):
        super(HttpDaemon, self).__init__(interface, port,
                                         handler=MultiHandler2)

class HttpSDaemon(threading.Thread, netsvc.Server):
    def __init__(self, interface, port):
        try:
            super(HttpDaemon, self).__init__(interface, port,
                                             handler=SecureMultiHandler2)
        except SSLError, e:
            netsvc.Logger().notifyChannel(
                'httpd-ssl', netsvc.LOG_CRITICAL,
                "Can not load the certificate and/or the private key files")
            raise

httpd = None
httpsd = None

def init_servers():
    global httpd, httpsd
    if tools.config.get_misc('httpd','enable', True):
        httpd = HttpDaemon(tools.config.get_misc('httpd','interface', ''), \
            int(tools.config.get_misc('httpd','port', tools.config.get('port',8069))))

    if tools.config.get_misc('httpsd','enable', False):
        httpsd = HttpSDaemon(tools.config.get_misc('httpsd','interface', ''), \
            int(tools.config.get_misc('httpsd','port', 8071)))

def reg_http_service(hts, secure_only = False):
    """ Register some handler to httpd.
        hts must be an HTTPDir
    """
    global httpd, httpsd
    if not isinstance(hts, HTTPDir):
        raise Exception("Wrong class for http service")

    if httpd and not secure_only:
        httpd.server.vdirs.append(hts)

    if httpsd:
        httpsd.server.vdirs.append(hts)

    if (not httpd) and (not httpsd):
        netsvc.Logger().notifyChannel('httpd',netsvc.LOG_WARNING,"No httpd available to register service %s" % hts.path)
    return

import SimpleXMLRPCServer
class XMLRPCRequestHandler(netsvc.OpenERPDispatcher,FixSendError,SimpleXMLRPCServer.SimpleXMLRPCRequestHandler):
    rpc_paths = []
    protocol_version = 'HTTP/1.1'
    def _dispatch(self, method, params):
        try:
            service_name = self.path.split("/")[-1]
            return self.dispatch(service_name, method, params)
        except netsvc.OpenERPDispatcherException, e:
            raise xmlrpclib.Fault(tools.exception_to_unicode(e.exception), e.traceback)

    def log_message(self, format, *args):
        netsvc.Logger().notifyChannel('xmlrpc',netsvc.LOG_DEBUG_RPC,format % args)

    def handle(self):
        pass

    def finish(self):
        pass

    def setup(self):
        self.connection = dummyconn()
        if not len(XMLRPCRequestHandler.rpc_paths):
            XMLRPCRequestHandler.rpc_paths = map(lambda s: '/%s' % s, netsvc.ExportService._services.keys())
        pass

def init_xmlrpc():
    if not tools.config.get_misc('xmlrpc','enable', True):
        return
    reg_http_service(HTTPDir('/xmlrpc/',XMLRPCRequestHandler))
    # Example of http file serving:
    # reg_http_service(HTTPDir('/test/',HTTPHandler))
    netsvc.Logger().notifyChannel("web-services", netsvc.LOG_INFO,
            "Registered XML-RPC over HTTP")

class OerpAuthProxy(AuthProxy):
    """ Require basic authentication..

        This is a copy of the BasicAuthProxy, which however checks/caches the db
        as well.
    """
    def __init__(self,provider):
        AuthProxy.__init__(self,provider)
        self.auth_creds = {}
        self.auth_tries = 0
        self.last_auth = None

    def checkRequest(self,handler,path = '/'):
        if self.auth_creds:
            return True
        auth_str = handler.headers.get('Authorization',False)
        try:
            db = handler.get_db_from_path(path)
            print "Got db:",db
        except:
            if path.startswith('/'):
                path = path[1:]
            psp= path.split('/')
            if len(psp)>1:
                db = psp[0]
            else:
                #FIXME!
                self.provider.log("Wrong path: %s, failing auth" %path)
                raise AuthRejectedExc("Authorization failed. Wrong sub-path.")

        if auth_str and auth_str.startswith('Basic '):
            auth_str=auth_str[len('Basic '):]
            (user,passwd) = base64.decodestring(auth_str).split(':')
            self.provider.log("Found user=\"%s\", passwd=\"***\" for db=\"%s\"" %(user,db))
            acd = self.provider.authenticate(db,user,passwd,handler.client_address)
            if acd != False:
                self.auth_creds[db] = acd
                self.last_auth=db
                return True
        if self.auth_tries > 5:
            self.provider.log("Failing authorization after 5 requests w/o password")
            raise AuthRejectedExc("Authorization failed.")
        self.auth_tries += 1
        raise AuthRequiredExc(atype = 'Basic', realm=self.provider.realm)

import security
class OpenERPAuthProvider(AuthProvider):
    def __init__(self,realm = 'OpenERP User'):
        self.realm = realm

    def setupAuth(self, multi, handler):
        if not multi.sec_realms.has_key(self.realm):
            multi.sec_realms[self.realm] = OerpAuthProxy(self)
        handler.auth_proxy = multi.sec_realms[self.realm]

    def authenticate(self, db, user, passwd, client_address):
        try:
            uid = security.login(db,user,passwd)
            if uid is False:
                return False
            return (user, passwd, db, uid)
        except Exception,e:
            netsvc.Logger().notifyChannel("auth",netsvc.LOG_DEBUG,"Fail auth:"+ str(e))
            return False

    def log(self, msg):
        netsvc.Logger().notifyChannel("auth",netsvc.LOG_INFO,msg)

#eof
<|endoftext|>


<commit_before>import copy
import functools

import elasticsearch
from elasticsearch.helpers import bulk

from inbox.api.kellogs import encode
from inbox.config import config
from inbox.log import get_logger
from inbox.search.query import DSLQueryEngine, MessageQuery, ThreadQuery
from inbox.search.mappings import NAMESPACE_INDEX_MAPPING

log = get_logger()

# Uncomment to enable logging of exactly which queries are made against the
# elasticsearch server, which you can paste directly into curl... also logs
# results, which is VERY verbose.
# import logging
# es_tracer = logging.getLogger('elasticsearch.trace')
# es_tracer.propagate = True

# Uncomment to disable verbose elasticsearch module logging.
# import logging
# es_logger = logging.getLogger('elasticsearch')
# es_logger.propagate = False

INDEX_SETTINGS = {
    "number_of_shards": 1
}

class Serializer(elasticsearch.serializer.JSONSerializer):
    """Override the Elasticsearch client's default serializer so that we
    serialize datetimes as Unix timestamps."""
    def default(self, data):
        serialized = encode(data)
        if serialized is not None:
            return serialized
        raise TypeError

class SearchEngineError(Exception):
    """ Raised when connecting to the Elasticsearch server fails. """
    pass

def new_connection():
    """
    Get a new connection to the Elasticsearch server defined in the config.

    """
    elasticsearch_hosts = config.get('ELASTICSEARCH_HOSTS')
    if not elasticsearch_hosts:
        raise SearchEngineError('No search hosts configured')
    return elasticsearch.Elasticsearch(hosts=elasticsearch_hosts,
                                       serializer=Serializer())

def wrap_es_errors(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except elasticsearch.TransportError as e:
            raise SearchEngineError(e)
    return wrapper

class NamespaceSearchEngine(object):
    """
    Interface to create and interact with the Elasticsearch datastore
    (i.e. index) for a namespace, identified by the namespace public id.

    """
    MAPPINGS = NAMESPACE_INDEX_MAPPING

    def __init__(self, namespace_public_id, create_index=False):
        self.index_id = namespace_public_id

        # TODO(emfree): probably want to try to keep persistent connections
        # around, instead of creating a new one each time.
        self._connection = new_connection()
        self.log = log.new(component='search', index=namespace_public_id)

        if create_index:
            self.create_index()

        self.messages = MessageSearchAdaptor(index_id=namespace_public_id,
                                             log=self.log)
        self.threads = ThreadSearchAdaptor(index_id=namespace_public_id,
                                           log=self.log)

    @wrap_es_errors
    def create_index(self):
        """
        Create an index for the namespace. If it already exists,
        re-configure it.

        """
        try:
            self.log.info('create_index')
            self._connection.indices.create(
                index=self.index_id,
                body={'mappings': NAMESPACE_INDEX_MAPPING,
                      'settings': INDEX_SETTINGS})

        except elasticsearch.exceptions.RequestError:
            self.log.warning('create_index error, will re-configure.')
            # If the index already exists, ensure the right mappings are still
            # configured. Only works if action.auto_create_index = False.
            return self.configure_index()

    @wrap_es_errors
    def configure_index(self):
        """
        Configure the index for the namespace. In case of mapping conflicts,
        delete and re-create it.

        """
        try:
            self.log.info('configure_index')
            for doc_type, mapping in self.MAPPINGS.items():
                self._connection.indices.put_mapping(
                    index=self.index_id, doc_type=doc_type, body=mapping)
        except elasticsearch.exceptions.RequestError:
            self.log.warning('configure_index error, will delete + create.')
            self.delete_index()
            self.create_index()

    @wrap_es_errors
    def delete_index(self):
        """
        Delete the index for the namespace.
        Use with care.

        Raises SearchEngineError if deletion fails for any error except
        'IndexNotFound' (ignores).

        """
        self.log.info('delete_index')

        # TODO[k]: Rigorous error handling?
        self._connection.indices.delete(index=[self.index_id], ignore=[404])

    @wrap_es_errors
    def refresh_index(self):
        """
        Manually refresh the index for the namespace (happens periodically by
        default). Makes all operations performed since the last refresh
        available for search.

        """
        self.log.info('refresh_index')
        self._connection.indices.refresh(index=[self.index_id])

class BaseSearchAdaptor(object):
    """
    Base adaptor between the Nilas API and Elasticsearch for a single index and
    document type. Subclasses implement the document type specific logic.

    """
    def __init__(self, index_id, doc_type, query_class, log):
        self.index_id = index_id
        self.doc_type = doc_type
        self.query_engine = DSLQueryEngine(query_class)

        self.log = log

        # TODO(emfree): probably want to try to keep persistent connections
        # around, instead of creating a new one each time.
        self._connection = new_connection()

    @wrap_es_errors
    def _index_document(self, object_repr, **kwargs):
        """
        (Re)index a document for the object with Nilas API representation
        `object_repr`.

        """
        assert self.index_id == object_repr['namespace_id']

        index_args = dict(
            index=self.index_id,
            doc_type=self.doc_type,
            id=object_repr['id'],
            body=object_repr)
        index_args.update(**kwargs)
        try:
            self._connection.index(**index_args)
        except elasticsearch.exceptions.TransportError as e:
            self.log.error('Index failure', error=e.error,
                           doc_type=self.doc_type, object_id=index_args['_id'])
            raise

    @wrap_es_errors
    def _bulk(self, objects, parent=None):
        """
        Perform a batch of index operations rather than a single one.

        Arguments
        ---------
        objects:
            list of (op_type, object) tuples.

            op_type defines the index operation to perform
            ('index' for creates, updates and 'delete' for deletes)

            object is a dict of document attributes required for the operation.

        Returns
        -------
        Count of index operations on success, raises SearchEngineError on
        failure.

        """
        index_args = []

        def raise_error(failure):
            for op_type, info in failure.iteritems():
                if info.get('status') not in [None, 404]:
                    return True
            return False

        for op, object_repr in objects:
            args = dict(_op_type=op,
                        _index=self.index_id,
                        _type=self.doc_type,
                        _id=object_repr['id'])

            if op != 'delete':
                args.update(dict(_source=object_repr))

                if parent is not None:
                    args.update(dict(_parent=object_repr[parent]))

            index_args.append(args)

        try:
            count, failures = bulk(self._connection, index_args)
        except elasticsearch.exceptions.TransportError as e:
            self.log.error('Bulk index failure', error=e.error,
                           doc_type=self.doc_type,
                           object_ids=[i['_id'] for i in index_args])
            raise SearchEngineError('Bulk index failure!')
        if count != len(objects):
            self.log.error('Bulk index failure',
                           error='Not all indices created',
                           doc_type=self.doc_type,
                           object_ids=[i['_id'] for i in index_args],
                           failures=failures)

            if any(raise_error(f) for f in failures):
                raise SearchEngineError('Bulk index failure!')

        return count

    @wrap_es_errors
    def count(self, query=None):
        search_kwargs = dict(
            index=self.index_id,
            doc_type=self.doc_type,
        )
        if query:
            search_kwargs.update(body=self.search_engine.generate_query(query))

        raw_result = self._connection.count(**search_kwargs)
        return raw_result.get("count", 0)

    @wrap_es_errors
    def search(self, query, sort, max_results=100, offset=0, explain=True):
        """ Perform a search and return the results. """
        dsl_query = self.query_engine.generate_query(query)

        self.log.debug('search query', query=query, dsl_query=dsl_query)

        search_kwargs = dict(index=self.index_id,
                             doc_type=self.doc_type,
                             body=dsl_query,
                             size=max_results,
                             from_=offset,
                             explain=explain)

        # Split this out to a Sort class with subclasses for
        # MessageSort/ThreadSort if we expand sorting to be more flexible.
        if sort != 'relevance':
            if self.doc_type == 'message':
                timestamp_field = 'date'
            if self.doc_type == 'thread':
                timestamp_field = 'last_message_timestamp'
            search_kwargs['sort'] = '{}:desc'.format(timestamp_field)
            # Uncomment this to include calculated scores in the results'
            # _score field (a little slower):
            # search_kwargs['body']['track_scores'] = True

        raw_results = self._connection.search(**search_kwargs)

        self._log_query(query, raw_results)

        total, api_results = self.query_engine.process_results(raw_results)
        return dict(total=total, results=api_results)

    @wrap_es_errors
    def get_mapping(self):
        return self._connection.indices.get_mapping(index=self.index_id,
                                                    doc_type=self.doc_type)

    def _log_query(self, query, raw_results):
        """
        Log query and result info, stripping out actual result bodies but
        keeping ids and metadata.

        """
        log_results = copy.deepcopy(raw_results)
        for hit in log_results['hits']['hits']:
            del hit['_source']
        self.log.debug('search query results', query=query,
                       results=log_results)

class MessageSearchAdaptor(BaseSearchAdaptor):
    """ Adaptor for the 'message' document type. """
    def __init__(self, index_id, log):
        BaseSearchAdaptor.__init__(self, index_id=index_id, doc_type='message',
                                   query_class=MessageQuery, log=log)

    def index(self, object_repr):
        self._index_document(object_repr, parent=object_repr['thread_id'])

    def bulk_index(self, objects):
        return self._bulk(objects, parent='thread_id')

class ThreadSearchAdaptor(BaseSearchAdaptor):
    """ Adaptor for the 'thread' document type. """
    def __init__(self, index_id, log):
        BaseSearchAdaptor.__init__(self, index_id=index_id, doc_type='thread',
                                   query_class=ThreadQuery, log=log)

    def index(self, object_repr):
        self._index_document(object_repr)

    def bulk_index(self, objects):
        return self._bulk(objects)
<commit_msg>Increasing elasticsearch default timeout to 60 seconds<commit_after>import copy
import functools

import elasticsearch
from elasticsearch.helpers import bulk

from inbox.api.kellogs import encode
from inbox.config import config
from inbox.log import get_logger
from inbox.search.query import DSLQueryEngine, MessageQuery, ThreadQuery
from inbox.search.mappings import NAMESPACE_INDEX_MAPPING

log = get_logger()

# Uncomment to enable logging of exactly which queries are made against the
# elasticsearch server, which you can paste directly into curl... also logs
# results, which is VERY verbose.
# import logging
# es_tracer = logging.getLogger('elasticsearch.trace')
# es_tracer.propagate = True

# Uncomment to disable verbose elasticsearch module logging.
# import logging
# es_logger = logging.getLogger('elasticsearch')
# es_logger.propagate = False

INDEX_SETTINGS = {
    "number_of_shards": 1
}
ES_DEFAULT_TIMEOUT = 60

class Serializer(elasticsearch.serializer.JSONSerializer):
    """Override the Elasticsearch client's default serializer so that we
    serialize datetimes as Unix timestamps."""
    def default(self, data):
        serialized = encode(data)
        if serialized is not None:
            return serialized
        raise TypeError

class SearchEngineError(Exception):
    """ Raised when connecting to the Elasticsearch server fails. """
    pass

def new_connection():
    """
    Get a new connection to the Elasticsearch server defined in the config.

    """
    elasticsearch_hosts = config.get('ELASTICSEARCH_HOSTS')
    if not elasticsearch_hosts:
        raise SearchEngineError('No search hosts configured')
    return elasticsearch.Elasticsearch(hosts=elasticsearch_hosts,
                                       serializer=Serializer(),
                                       timeout=ES_DEFAULT_TIMEOUT)

def wrap_es_errors(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except elasticsearch.TransportError as e:
            raise SearchEngineError(e)
    return wrapper

class NamespaceSearchEngine(object):
    """
    Interface to create and interact with the Elasticsearch datastore
    (i.e. index) for a namespace, identified by the namespace public id.

    """
    MAPPINGS = NAMESPACE_INDEX_MAPPING

    def __init__(self, namespace_public_id, create_index=False):
        self.index_id = namespace_public_id

        # TODO(emfree): probably want to try to keep persistent connections
        # around, instead of creating a new one each time.
        self._connection = new_connection()
        self.log = log.new(component='search', index=namespace_public_id)

        if create_index:
            self.create_index()

        self.messages = MessageSearchAdaptor(index_id=namespace_public_id,
                                             log=self.log)
        self.threads = ThreadSearchAdaptor(index_id=namespace_public_id,
                                           log=self.log)

    @wrap_es_errors
    def create_index(self):
        """
        Create an index for the namespace. If it already exists,
        re-configure it.

        """
        try:
            self.log.info('create_index')
            self._connection.indices.create(
                index=self.index_id,
                body={'mappings': NAMESPACE_INDEX_MAPPING,
                      'settings': INDEX_SETTINGS})

        except elasticsearch.exceptions.RequestError:
            self.log.warning('create_index error, will re-configure.')
            # If the index already exists, ensure the right mappings are still
            # configured. Only works if action.auto_create_index = False.
            return self.configure_index()

    @wrap_es_errors
    def configure_index(self):
        """
        Configure the index for the namespace. In case of mapping conflicts,
        delete and re-create it.

        """
        try:
            self.log.info('configure_index')
            for doc_type, mapping in self.MAPPINGS.items():
                self._connection.indices.put_mapping(
                    index=self.index_id, doc_type=doc_type, body=mapping)
        except elasticsearch.exceptions.RequestError:
            self.log.warning('configure_index error, will delete + create.')
            self.delete_index()
            self.create_index()

    @wrap_es_errors
    def delete_index(self):
        """
        Delete the index for the namespace.
        Use with care.

        Raises SearchEngineError if deletion fails for any error except
        'IndexNotFound' (ignores).

        """
        self.log.info('delete_index')

        # TODO[k]: Rigorous error handling?
        self._connection.indices.delete(index=[self.index_id], ignore=[404])

    @wrap_es_errors
    def refresh_index(self):
        """
        Manually refresh the index for the namespace (happens periodically by
        default). Makes all operations performed since the last refresh
        available for search.

        """
        self.log.info('refresh_index')
        self._connection.indices.refresh(index=[self.index_id])

class BaseSearchAdaptor(object):
    """
    Base adaptor between the Nilas API and Elasticsearch for a single index and
    document type. Subclasses implement the document type specific logic.

    """
    def __init__(self, index_id, doc_type, query_class, log):
        self.index_id = index_id
        self.doc_type = doc_type
        self.query_engine = DSLQueryEngine(query_class)

        self.log = log

        # TODO(emfree): probably want to try to keep persistent connections
        # around, instead of creating a new one each time.
        self._connection = new_connection()

    @wrap_es_errors
    def _index_document(self, object_repr, **kwargs):
        """
        (Re)index a document for the object with Nilas API representation
        `object_repr`.

        """
        assert self.index_id == object_repr['namespace_id']

        index_args = dict(
            index=self.index_id,
            doc_type=self.doc_type,
            id=object_repr['id'],
            body=object_repr)
        index_args.update(**kwargs)
        try:
            self._connection.index(**index_args)
        except elasticsearch.exceptions.TransportError as e:
            self.log.error('Index failure', error=e.error,
                           doc_type=self.doc_type, object_id=index_args['_id'])
            raise

    @wrap_es_errors
    def _bulk(self, objects, parent=None):
        """
        Perform a batch of index operations rather than a single one.

        Arguments
        ---------
        objects:
            list of (op_type, object) tuples.

            op_type defines the index operation to perform
            ('index' for creates, updates and 'delete' for deletes)

            object is a dict of document attributes required for the operation.

        Returns
        -------
        Count of index operations on success, raises SearchEngineError on
        failure.

        """
        index_args = []

        def raise_error(failure):
            for op_type, info in failure.iteritems():
                if info.get('status') not in [None, 404]:
                    return True
            return False

        for op, object_repr in objects:
            args = dict(_op_type=op,
                        _index=self.index_id,
                        _type=self.doc_type,
                        _id=object_repr['id'])

            if op != 'delete':
                args.update(dict(_source=object_repr))

                if parent is not None:
                    args.update(dict(_parent=object_repr[parent]))

            index_args.append(args)

        try:
            count, failures = bulk(self._connection, index_args)
        except elasticsearch.exceptions.TransportError as e:
            self.log.error('Bulk index failure', error=e.error,
                           doc_type=self.doc_type,
                           object_ids=[i['_id'] for i in index_args])
            raise SearchEngineError('Bulk index failure!')
        if count != len(objects):
            self.log.error('Bulk index failure',
                           error='Not all indices created',
                           doc_type=self.doc_type,
                           object_ids=[i['_id'] for i in index_args],
                           failures=failures)

            if any(raise_error(f) for f in failures):
                raise SearchEngineError('Bulk index failure!')

        return count

    @wrap_es_errors
    def count(self, query=None):
        search_kwargs = dict(
            index=self.index_id,
            doc_type=self.doc_type,
        )
        if query:
            search_kwargs.update(body=self.search_engine.generate_query(query))

        raw_result = self._connection.count(**search_kwargs)
        return raw_result.get("count", 0)

    @wrap_es_errors
    def search(self, query, sort, max_results=100, offset=0, explain=True):
        """ Perform a search and return the results. """
        dsl_query = self.query_engine.generate_query(query)

        self.log.debug('search query', query=query, dsl_query=dsl_query)

        search_kwargs = dict(index=self.index_id,
                             doc_type=self.doc_type,
                             body=dsl_query,
                             size=max_results,
                             from_=offset,
                             explain=explain)

        # Split this out to a Sort class with subclasses for
        # MessageSort/ThreadSort if we expand sorting to be more flexible.
        if sort != 'relevance':
            if self.doc_type == 'message':
                timestamp_field = 'date'
            if self.doc_type == 'thread':
                timestamp_field = 'last_message_timestamp'
            search_kwargs['sort'] = '{}:desc'.format(timestamp_field)
            # Uncomment this to include calculated scores in the results'
            # _score field (a little slower):
            # search_kwargs['body']['track_scores'] = True

        raw_results = self._connection.search(**search_kwargs)

        self._log_query(query, raw_results)

        total, api_results = self.query_engine.process_results(raw_results)
        return dict(total=total, results=api_results)

    @wrap_es_errors
    def get_mapping(self):
        return self._connection.indices.get_mapping(index=self.index_id,
                                                    doc_type=self.doc_type)

    def _log_query(self, query, raw_results):
        """
        Log query and result info, stripping out actual result bodies but
        keeping ids and metadata.

        """
        log_results = copy.deepcopy(raw_results)
        for hit in log_results['hits']['hits']:
            del hit['_source']
        self.log.debug('search query results', query=query,
                       results=log_results)

class MessageSearchAdaptor(BaseSearchAdaptor):
    """ Adaptor for the 'message' document type. """
    def __init__(self, index_id, log):
        BaseSearchAdaptor.__init__(self, index_id=index_id, doc_type='message',
                                   query_class=MessageQuery, log=log)

    def index(self, object_repr):
        self._index_document(object_repr, parent=object_repr['thread_id'])

    def bulk_index(self, objects):
        return self._bulk(objects, parent='thread_id')

class ThreadSearchAdaptor(BaseSearchAdaptor):
    """ Adaptor for the 'thread' document type. """
    def __init__(self, index_id, log):
        BaseSearchAdaptor.__init__(self, index_id=index_id, doc_type='thread',
                                   query_class=ThreadQuery, log=log)

    def index(self, object_repr):
        self._index_document(object_repr)

    def bulk_index(self, objects):
        return self._bulk(objects)
<|endoftext|>


<commit_before>from __future__ import unicode_literals

import io
import json
import traceback
import hashlib
import os
import subprocess
import sys
from zipimport import zipimporter

from .utils import encode_compat_str

from .version import __version__

def rsa_verify(message, signature, key):
    from hashlib import sha256
    assert isinstance(message, bytes)
    byte_size = (len(bin(key[0])) - 2 + 8 - 1) // 8
    signature = ('%x' % pow(int(signature, 16), key[1], key[0])).encode()
    signature = (byte_size * 2 - len(signature)) * b'0' + signature
    asn1 = b'3031300d060960864801650304020105000420'
    asn1 += sha256(message).hexdigest().encode()
    if byte_size < len(asn1) // 2 + 11:
        return False
    expected = b'0001' + (byte_size - len(asn1) // 2 - 3) * b'ff' + b'00' + asn1
    return expected == signature

def update_self(to_screen, verbose, opener):
    """Update the program file with the latest version from the repository"""

    UPDATE_URL = 'https://ytdl-org.github.io/youtube-dl/update/'
    VERSION_URL = UPDATE_URL + 'LATEST_VERSION'
    JSON_URL = UPDATE_URL + 'versions.json'
    UPDATES_RSA_KEY = (0x9d60ee4d8f805312fdb15a62f87b95bd66177b91df176765d13514a0f1754bcd2057295c5b6f1d35daa6742c3ffc9a82d3e118861c207995a8031e151d863c9927e304576bc80692bc8e094896fcf11b66f3e29e04e3a71e9a11558558acea1840aec37fc396fb6b65dc81a1c4144e03bd1c011de62e3f1357b327d08426fe93, 65537)

    if not isinstance(globals().get('__loader__'), zipimporter) and not hasattr(sys, 'frozen'):
        to_screen('It looks like you installed youtube-dl with a package manager, pip, setup.py or a tarball. Please use that to update.')
        return

    # Check if there is a new version
    try:
        newversion = opener.open(VERSION_URL).read().decode('utf-8').strip()
    except Exception:
        if verbose:
            to_screen(encode_compat_str(traceback.format_exc()))
        to_screen('ERROR: can\'t find the current version. Please try again later.')
        return
    if newversion == __version__:
        to_screen('youtube-dl is up-to-date (' + __version__ + ')')
        return

    # Download and check versions info
    try:
        versions_info = opener.open(JSON_URL).read().decode('utf-8')
        versions_info = json.loads(versions_info)
    except Exception:
        if verbose:
            to_screen(encode_compat_str(traceback.format_exc()))
        to_screen('ERROR: can\'t obtain versions info. Please try again later.')
        return
    if 'signature' not in versions_info:
        to_screen('ERROR: the versions file is not signed or corrupted. Aborting.')
        return
    signature = versions_info['signature']
    del versions_info['signature']
    if not rsa_verify(json.dumps(versions_info, sort_keys=True).encode('utf-8'), signature, UPDATES_RSA_KEY):
        to_screen('ERROR: the versions file signature is invalid. Aborting.')
        return

    version_id = versions_info['latest']

    def version_tuple(version_str):
        return tuple(map(int, version_str.split('.')))
    if version_tuple(__version__) >= version_tuple(version_id):
        to_screen('youtube-dl is up to date (%s)' % __version__)
        return

    to_screen('Updating to version ' + version_id + ' ...')
    version = versions_info['versions'][version_id]

    print_notes(to_screen, versions_info['versions'])

    # sys.executable is set to the full pathname of the exe-file for py2exe
    filename = sys.executable if hasattr(sys, 'frozen') else sys.argv[0]

    if not os.access(filename, os.W_OK):
        to_screen('ERROR: no write permissions on %s' % filename)
        return

    # Py2EXE
    if hasattr(sys, 'frozen'):
        exe = filename
        directory = os.path.dirname(exe)
        if not os.access(directory, os.W_OK):
            to_screen('ERROR: no write permissions on %s' % directory)
            return

        try:
            urlh = opener.open(version['exe'][0])
            newcontent = urlh.read()
            urlh.close()
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to download latest version')
            return

        newcontent_hash = hashlib.sha256(newcontent).hexdigest()
        if newcontent_hash != version['exe'][1]:
            to_screen('ERROR: the downloaded file hash does not match. Aborting.')
            return

        try:
            with open(exe + '.new', 'wb') as outf:
                outf.write(newcontent)
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to write the new version')
            return

        try:
            bat = os.path.join(directory, 'youtube-dl-updater.bat')
            with io.open(bat, 'w') as batfile:
                batfile.write('''
@echo off
echo Waiting for file handle to be closed ...
ping 127.0.0.1 -n 5 -w 1000 > NUL
move /Y "%s.new" "%s" > NUL
echo Updated youtube-dl to version %s.
start /b "" cmd /c del "%%~f0"&exit /b"
                \n''' % (exe, exe, version_id))

            subprocess.Popen([bat])  # Continues to run in the background
            return  # Do not show premature success messages
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to overwrite current version')
            return

    # Zip unix package
    elif isinstance(globals().get('__loader__'), zipimporter):
        try:
            urlh = opener.open(version['bin'][0])
            newcontent = urlh.read()
            urlh.close()
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to download latest version')
            return

        newcontent_hash = hashlib.sha256(newcontent).hexdigest()
        if newcontent_hash != version['bin'][1]:
            to_screen('ERROR: the downloaded file hash does not match. Aborting.')
            return

        try:
            with open(filename, 'wb') as outf:
                outf.write(newcontent)
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to overwrite current version')
            return

    to_screen('Updated youtube-dl. Restart youtube-dl to use the new version.')

def get_notes(versions, fromVersion):
    notes = []
    for v, vdata in sorted(versions.items()):
        if v > fromVersion:
            notes.extend(vdata.get('notes', []))
    return notes

def print_notes(to_screen, versions, fromVersion=__version__):
    notes = get_notes(versions, fromVersion)
    if notes:
        to_screen('PLEASE NOTE:')
        for note in notes:
            to_screen(note)
<commit_msg>[update] Hide update URLs behind redirect<commit_after>from __future__ import unicode_literals

import io
import json
import traceback
import hashlib
import os
import subprocess
import sys
from zipimport import zipimporter

from .utils import encode_compat_str

from .version import __version__

def rsa_verify(message, signature, key):
    from hashlib import sha256
    assert isinstance(message, bytes)
    byte_size = (len(bin(key[0])) - 2 + 8 - 1) // 8
    signature = ('%x' % pow(int(signature, 16), key[1], key[0])).encode()
    signature = (byte_size * 2 - len(signature)) * b'0' + signature
    asn1 = b'3031300d060960864801650304020105000420'
    asn1 += sha256(message).hexdigest().encode()
    if byte_size < len(asn1) // 2 + 11:
        return False
    expected = b'0001' + (byte_size - len(asn1) // 2 - 3) * b'ff' + b'00' + asn1
    return expected == signature

def update_self(to_screen, verbose, opener):
    """Update the program file with the latest version from the repository"""

    UPDATE_URL = 'https://yt-dl.org/update/'
    VERSION_URL = UPDATE_URL + 'LATEST_VERSION'
    JSON_URL = UPDATE_URL + 'versions.json'
    UPDATES_RSA_KEY = (0x9d60ee4d8f805312fdb15a62f87b95bd66177b91df176765d13514a0f1754bcd2057295c5b6f1d35daa6742c3ffc9a82d3e118861c207995a8031e151d863c9927e304576bc80692bc8e094896fcf11b66f3e29e04e3a71e9a11558558acea1840aec37fc396fb6b65dc81a1c4144e03bd1c011de62e3f1357b327d08426fe93, 65537)

    if not isinstance(globals().get('__loader__'), zipimporter) and not hasattr(sys, 'frozen'):
        to_screen('It looks like you installed youtube-dl with a package manager, pip, setup.py or a tarball. Please use that to update.')
        return

    # Check if there is a new version
    try:
        newversion = opener.open(VERSION_URL).read().decode('utf-8').strip()
    except Exception:
        if verbose:
            to_screen(encode_compat_str(traceback.format_exc()))
        to_screen('ERROR: can\'t find the current version. Please try again later.')
        return
    if newversion == __version__:
        to_screen('youtube-dl is up-to-date (' + __version__ + ')')
        return

    # Download and check versions info
    try:
        versions_info = opener.open(JSON_URL).read().decode('utf-8')
        versions_info = json.loads(versions_info)
    except Exception:
        if verbose:
            to_screen(encode_compat_str(traceback.format_exc()))
        to_screen('ERROR: can\'t obtain versions info. Please try again later.')
        return
    if 'signature' not in versions_info:
        to_screen('ERROR: the versions file is not signed or corrupted. Aborting.')
        return
    signature = versions_info['signature']
    del versions_info['signature']
    if not rsa_verify(json.dumps(versions_info, sort_keys=True).encode('utf-8'), signature, UPDATES_RSA_KEY):
        to_screen('ERROR: the versions file signature is invalid. Aborting.')
        return

    version_id = versions_info['latest']

    def version_tuple(version_str):
        return tuple(map(int, version_str.split('.')))
    if version_tuple(__version__) >= version_tuple(version_id):
        to_screen('youtube-dl is up to date (%s)' % __version__)
        return

    to_screen('Updating to version ' + version_id + ' ...')
    version = versions_info['versions'][version_id]

    print_notes(to_screen, versions_info['versions'])

    # sys.executable is set to the full pathname of the exe-file for py2exe
    filename = sys.executable if hasattr(sys, 'frozen') else sys.argv[0]

    if not os.access(filename, os.W_OK):
        to_screen('ERROR: no write permissions on %s' % filename)
        return

    # Py2EXE
    if hasattr(sys, 'frozen'):
        exe = filename
        directory = os.path.dirname(exe)
        if not os.access(directory, os.W_OK):
            to_screen('ERROR: no write permissions on %s' % directory)
            return

        try:
            urlh = opener.open(version['exe'][0])
            newcontent = urlh.read()
            urlh.close()
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to download latest version')
            return

        newcontent_hash = hashlib.sha256(newcontent).hexdigest()
        if newcontent_hash != version['exe'][1]:
            to_screen('ERROR: the downloaded file hash does not match. Aborting.')
            return

        try:
            with open(exe + '.new', 'wb') as outf:
                outf.write(newcontent)
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to write the new version')
            return

        try:
            bat = os.path.join(directory, 'youtube-dl-updater.bat')
            with io.open(bat, 'w') as batfile:
                batfile.write('''
@echo off
echo Waiting for file handle to be closed ...
ping 127.0.0.1 -n 5 -w 1000 > NUL
move /Y "%s.new" "%s" > NUL
echo Updated youtube-dl to version %s.
start /b "" cmd /c del "%%~f0"&exit /b"
                \n''' % (exe, exe, version_id))

            subprocess.Popen([bat])  # Continues to run in the background
            return  # Do not show premature success messages
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to overwrite current version')
            return

    # Zip unix package
    elif isinstance(globals().get('__loader__'), zipimporter):
        try:
            urlh = opener.open(version['bin'][0])
            newcontent = urlh.read()
            urlh.close()
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to download latest version')
            return

        newcontent_hash = hashlib.sha256(newcontent).hexdigest()
        if newcontent_hash != version['bin'][1]:
            to_screen('ERROR: the downloaded file hash does not match. Aborting.')
            return

        try:
            with open(filename, 'wb') as outf:
                outf.write(newcontent)
        except (IOError, OSError):
            if verbose:
                to_screen(encode_compat_str(traceback.format_exc()))
            to_screen('ERROR: unable to overwrite current version')
            return

    to_screen('Updated youtube-dl. Restart youtube-dl to use the new version.')

def get_notes(versions, fromVersion):
    notes = []
    for v, vdata in sorted(versions.items()):
        if v > fromVersion:
            notes.extend(vdata.get('notes', []))
    return notes

def print_notes(to_screen, versions, fromVersion=__version__):
    notes = get_notes(versions, fromVersion)
    if notes:
        to_screen('PLEASE NOTE:')
        for note in notes:
            to_screen(note)
<|endoftext|>


<commit_before>#!/usr/bin/python -tt
#
# yum-rocket
# Fast, threaded downloads for YUM
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import os
import time
import threading
import Queue
import urllib
from urlparse import urlparse, urljoin

from yum.plugins import TYPE_CORE, PluginYumExit

requires_api_version = '2.5'
plugin_type = (TYPE_CORE,)

spanmirrors = 3
threadcount = 5
repo_list   = dict()

def init_hook(conduit):
    global threadcount, spanmirrors
    threadcount = conduit.confInt('main', 'threadcount', default=5)
    spanmirrors = conduit.confInt('main', 'spanmirrors', default=3)
    if hasattr(conduit, 'registerPackageName'):
        conduit.registerPackageName('yum-rocket')

def predownload_hook(conduit):
    global threadcount, spanmirrors, repo_list

    def prioritize_po_repos(po):
        """ Organize a repository list based on mirror speed.

        This function will (somewhat) intelligently choose from a number of
        mirrors to download from. It will choose only the fastest set of
        mirrors, and keep track of how many threads are already downloading
        from each of them to help distribute the load over multiple fast
        mirror servers.
        """
        repo_url = None
        load = 1
        urlcount = 0
        for url in po.repo._urls:
            if urlcount < spanmirrors:
                if not po.repo.id in repo_list.keys():
                    repo_list[po.repo.id] = dict()
                if not url in repo_list[po.repo.id].keys():
                    repo_list[po.repo.id][url] = 0
                    urlcount += 1

        while not repo_url:
            for r in repo_list[po.repo.id].keys():
                if repo_list[po.repo.id][r] < load:
                    repo_list[po.repo.id][r] += 1
                    repo_url = r
                    break
            load += 1

        po.repo._urls = repo_list[po.repo.id].keys()
        po.repo._urls.remove(repo_url)
        po.repo._urls = [repo_url] + po.repo._urls

        return repo_url

    def mirror_load(id, url, load):
        """ Update the load value of a given mirror """
        if id in repo_list.keys() and url in repo_list[id].keys():
            repo_list[id][url] += load

    def getPackage(po, thread_id):
        """ Handle downloading a package from a repository. """
        repo_url = prioritize_po_repos(po)
        repo_host = urlparse(repo_url).netloc
        conduit.verbose_logger.info('[%s] start: %s [%s]' %
                                      (thread_id, po, repo_host))
        mirror_load(po.repo.id, repo_url, 1)
        url = urljoin(repo_url, po.remote_path)
        urllib.urlretrieve(url, po.localPkg())
        mirror_load(po.repo.id, repo_url, -1)
        conduit.verbose_logger.info('[%s] done: %s' %
                                      (thread_id, po))

    def wait_on_threads(threads):
        """ Wait for a list of threads to finish working. """
        while len(threads) > 0:
            for thread in threads:
                if not thread.is_alive():
                    threads.remove(thread)

    class PkgDownloadThread(threading.Thread):
        def __init__(self, q, run_event):
            threading.Thread.__init__(self)
            self.q = q
            self.run_event = run_event

        def run(self):
            while self.run_event.is_set() and not self.q.empty():
                po = self.q.get()
                getPackage(po, self.name)
                self.q.task_done()
            if not self.run_event.is_set():
                conduit.logger.warn('Stopping %s...' % self.name)

    download_po = []
    for po in conduit.getDownloadPackages():
        local = po.localPkg()
        if os.path.exists(local):
            conduit.verbose_logger.debug('using local copy of %s' % po)
        else:
            download_po.append(po)

        dirstat = os.statvfs(po.repo.pkgdir)
        if (dirstat.f_bavail * dirstat.f_bsize) <= long(po.size):
            raise PluginYumExit, ('Insufficient disk space in directory %s' %
                                 po.repo.pkgdir)

    # Let's thread this bitch!
    if len(download_po) < threadcount:
        threadcount = len(download_po)
    if (len(download_po) > 0):
        conduit.verbose_logger.info('Spawning %d download threads' %
                                   threadcount)

    q = Queue.Queue()
    for po in download_po:
        q.put(po)

    run_event = threading.Event()
    run_event.set()

    threads = []
    for i in range(0, threadcount):
        thread = PkgDownloadThread(q, run_event)
        thread.start()
        threads.append(thread)

    try:
        wait_on_threads(threads)
    except:
        conduit.logger.warn('\n\nCaught exit signal\nHang tight, '
                         'waiting for threads to exit...\n')
        run_event.clear()
        wait_on_threads(threads)
        raise PluginYumExit, 'Threads terminated'
<commit_msg>Formatting adjustments in rocket.py<commit_after>#!/usr/bin/python -tt
#
# yum-rocket
# Fast, threaded downloads for YUM
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import os
import time
import threading
import Queue
import urllib
from urlparse import urlparse, urljoin

from yum.plugins import TYPE_CORE, PluginYumExit

requires_api_version = '2.5'
plugin_type = (TYPE_CORE,)

spanmirrors = 3
threadcount = 5
repo_list   = dict()

def init_hook(conduit):
    global threadcount, spanmirrors
    threadcount = conduit.confInt('main', 'threadcount', default=5)
    spanmirrors = conduit.confInt('main', 'spanmirrors', default=3)
    if hasattr(conduit, 'registerPackageName'):
        conduit.registerPackageName('yum-rocket')

def predownload_hook(conduit):
    global threadcount, spanmirrors, repo_list

    def prioritize_po_repos(po):
        """ Organize a repository list based on mirror speed.

        This function will (somewhat) intelligently choose from a number of
        mirrors to download from. It will choose only the fastest set of
        mirrors, and keep track of how many threads are already downloading
        from each of them to help distribute the load over multiple fast
        mirror servers.
        """
        repo_url = None
        load = 1
        urlcount = 0
        for url in po.repo._urls:
            if urlcount < spanmirrors:
                if not po.repo.id in repo_list.keys():
                    repo_list[po.repo.id] = dict()
                if not url in repo_list[po.repo.id].keys():
                    repo_list[po.repo.id][url] = 0
                    urlcount += 1

        while not repo_url:
            for r in repo_list[po.repo.id].keys():
                if repo_list[po.repo.id][r] < load:
                    repo_list[po.repo.id][r] += 1
                    repo_url = r
                    break
            load += 1

        po.repo._urls = repo_list[po.repo.id].keys()
        po.repo._urls.remove(repo_url)
        po.repo._urls = [repo_url] + po.repo._urls

        return repo_url

    def mirror_load(id, url, load):
        """ Update the load value of a given mirror """
        if id in repo_list.keys() and url in repo_list[id].keys():
            repo_list[id][url] += load

    def getPackage(po, thread_id):
        """ Handle downloading a package from a repository. """
        repo_url = prioritize_po_repos(po)
        repo_host = urlparse(repo_url).netloc
        conduit.verbose_logger.info('[%s] start: %s [%s]' %
                                    (thread_id, po, repo_host))
        mirror_load(po.repo.id, repo_url, 1)
        url = urljoin(repo_url, po.remote_path)
        urllib.urlretrieve(url, po.localPkg())
        mirror_load(po.repo.id, repo_url, -1)
        conduit.verbose_logger.info('[%s] done: %s' %
                                    (thread_id, po))

    def wait_on_threads(threads):
        """ Wait for a list of threads to finish working. """
        while len(threads) > 0:
            for thread in threads:
                if not thread.is_alive():
                    threads.remove(thread)

    class PkgDownloadThread(threading.Thread):
        def __init__(self, q, run_event):
            threading.Thread.__init__(self)
            self.q = q
            self.run_event = run_event

        def run(self):
            while self.run_event.is_set() and not self.q.empty():
                po = self.q.get()
                getPackage(po, self.name)
                self.q.task_done()
            if not self.run_event.is_set():
                conduit.logger.warn('Stopping %s...' % self.name)

    download_po = []
    for po in conduit.getDownloadPackages():
        local = po.localPkg()
        if os.path.exists(local):
            conduit.verbose_logger.debug('using local copy of %s' % po)
        else:
            download_po.append(po)

        dirstat = os.statvfs(po.repo.pkgdir)
        if (dirstat.f_bavail * dirstat.f_bsize) <= long(po.size):
            raise PluginYumExit, ('Insufficient disk space in directory %s' %
                                  po.repo.pkgdir)

    # Let's thread this bitch!
    if len(download_po) < threadcount:
        threadcount = len(download_po)
    if (len(download_po) > 0):
        conduit.verbose_logger.info('Spawning %d download threads' %
                                    threadcount)

    q = Queue.Queue()
    for po in download_po:
        q.put(po)

    run_event = threading.Event()
    run_event.set()

    threads = []
    for i in range(0, threadcount):
        thread = PkgDownloadThread(q, run_event)
        thread.start()
        threads.append(thread)

    try:
        wait_on_threads(threads)
    except:
        conduit.logger.warn('\n\nCaught exit signal\nHang tight, '
                            'waiting for threads to exit...\n')
        run_event.clear()
        wait_on_threads(threads)
        raise PluginYumExit, 'Threads terminated'
<|endoftext|>


<commit_before>from .bases import CellLoop
from .compatibility import one_dimension, two_dimensions, activity
from .utils import offset_pos

from itertools import product, izip
import numpy as np

class OneDimCellLoop(CellLoop):
    """The OneDimCellLoop iterates over all cells in order from 0 to sizeX."""

    requires_features = [one_dimension]

    def get_pos(self):
        return "loop_x"

    def visit(self):
        super(OneDimCellLoop, self).visit()
        self.code.add_code("loop_begin",
                """for(int loop_x=0; loop_x < sizeX; loop_x++) {""")
        self.code.add_code("loop_end",
                """}""")

    def get_iter(self):
        return iter(izip(xrange(0, self.code.acc.get_size_of())))

    def build_name(self, parts):
        parts.insert(0, "1d")

class TwoDimCellLoop(CellLoop):
    """The TwoDimCellLoop iterates over all cells from left to right, then from
    top to bottom."""

    requires_features = [two_dimensions]

    def get_pos(self):
        return "loop_x", "loop_y"

    def visit(self):
        super(TwoDimCellLoop, self).visit()
        size_names = self.code.acc.size_names
        self.code.add_code("loop_begin",
            """for(int loop_x=0; loop_x < %s; loop_x++) {
                for(int loop_y=0; loop_y < %s; loop_y++) {""" % (size_names))
        self.code.add_code("loop_end",
                """}
                }""")

    def get_iter(self):
        return iter(product(xrange(0, self.code.acc.get_size_of(0)),
                            xrange(0, self.code.acc.get_size_of(1))))

    def build_name(self, parts):
        parts.insert(0, "2d")

class SparseCellLoop(CellLoop):
    """The SparseCellLoopBase offers common code for loops that only calculate
    those fields, where the neighbours have changed in the last step.

    This is based on a list of positions called `sparse_list` as well as a mask
    of booleans called `sparse_mask`, that only internally gets used to make
    sure, that no fields are enlisted more than once.

    The `sparse_list` is duplicated into the property `prev_sparse_list`, from
    which reads are performed. Additionally, the number of active cells is
    taken from an "activity" stats before the loop to see how many valid
    entries exist in the array.

    For the pure-py version, a normal python set is used."""

    requires_features = [activity]

    def set_target(self, target):
        """Adds the activity mask and position list to the target attributes."""
        super(SparseCellLoop, self).set_target(target)

        if len(target.size) == 1:
            self.requires_features.append(one_dimension)
        elif len(target.size) == 2:
            self.requires_features.append(two_dimensions)
        else:
            raise NotImplementedError("SparseCellLoop has not been modified to"
                    " work with more than 2 dimensions.")

        size = self.calculate_size()
        target.sparse_mask = np.zeros(size, dtype=np.bool)
        target.sparse_list = np.zeros(size, dtype=np.int)
        target.sparse_set = set()

    def get_pos(self):
        return self.position_names

    def calculate_size(self):
        """Calculate how big the mask and list have to be.

        The current strategy is to just allocate one field for each field of the
        configuration."""
        return reduce(lambda a, b: a * b, self.target.size)

    def mark_cell_py(self, pos):
        positions = [offset_pos(pos, offs) for offs in
                     self.code.neigh.affected_cells()]
        positions = [pos
                         if self.code.border.is_position_valid(pos)
                         else self.code.border.correct_position(pos)
                     for pos in positions]
        self.target.sparse_set.update(positions)

    def new_config(self):
        size = self.calculate_size()
        print size
        self.target.sparse_mask = np.ones(size, dtype=np.bool)
        self.target.sparse_list = np.array(range(size), dtype=np.int)
        self.target.prev_sparse_list = self.target.sparse_list.copy()
        print self.target.prev_sparse_list
        self.target.sparse_set = set(product(*[range(siz) for siz in self.target.size]))

    def get_iter(self):
        # iterate over a copy of the set, so that it can be modified while running
        print len(self.target.sparse_set)
        the_list = list(self.target.sparse_set)
        self.target.sparse_set.clear()
        return iter(the_list)

    def visit(self):
        super(SparseCellLoop, self).visit()

        self.code.attrs.append("sparse_mask")
        self.code.attrs.append("sparse_list")
        self.code.attrs.append("prev_sparse_list")

        self.code.add_py_hook("loop_end",
            """if was_active: self.loop.mark_cell_py(pos)""")

        self.code.add_code("localvars",
            """int valid_positions = activity(0);
               if(valid_positions == -1) {
                   valid_positions = %(size_a)s * %(size_b)s;
               }
               int sparse_cell_write_idx = 0;""" %
                   dict(size_a=self.code.acc.size_names[0],
                        size_b=self.code.acc.size_names[1] if len(self.code.acc.size_names) == 2 else "1"))
        self.code.add_code("loop_begin",
            """for(int cell_idx=0; cell_idx < valid_positions; cell_idx++) {""")
        if len(self.position_names) == 1:
            self.code.add_code("loop_begin",
                """    int %s = prev_sparse_list(cell_idx);""" % self.position_names)
        elif len(self.position_names) == 2:
            self.code.add_code("loop_begin",
                """    int %(pos_a)s = prev_sparse_list(cell_idx) %% %(size_a)s;
                       int %(pos_b)s = prev_sparse_list(cell_idx) / %(size_b)s;""" %
                           dict(pos_a = self.position_names[0],
                                pos_b = self.position_names[1],
                                size_a = self.code.acc.size_names[0],
                                size_b = self.code.acc.size_names[1]))

        # FIXME use procer position names here
        # FIXME wrap positions around the borders
        if len(self.position_names) == 1:
            self.code.add_code("loop_end",
                    """if(was_active) {
                               %s
                       }""" % ("\n".join([
                           """
                           {int idx = i + %(offs_x)s;
                           if(!sparse_mask(idx)) {
                               sparse_list(sparse_cell_write_idx) = idx;
                               sparse_mask(idx) = true;
                               sparse_cell_write_idx++;
                           }}""" % dict(offs_x=offs[0])
                               for offs in self.code.neigh.offsets])))
        elif len(self.position_names) == 2:
            self.code.add_code("loop_end",
                    """if(was_active) {
                               %s
                       }""" % ("\n".join([
                           """
                           {int px = i + %(offs_x)s;
                           int py = j + %(offs_y)s;
                           int idx = px * %(size_x)s + py;
                           if(!sparse_mask(idx)) {
                               sparse_list(sparse_cell_write_idx) = idx;
                               sparse_mask(idx) = true;
                               sparse_cell_write_idx++;
                           }}""" % dict(offs_x=offs[0], offs_y=offs[1],
                                       size_x=self.code.acc.size_names[0])
                               for offs in self.code.neigh.offsets])))
        self.code.add_code("loop_end",
                """
                }
                sparse_mask = sparse_mask*0;""")

class OneDimSparseCellLoop(SparseCellLoop):
    def __init__(self):
        self.position_names = "loop_x"

class TwoDimSparseCellLoop(SparseCellLoop):
    def __init__(self):
        self.position_names = "loop_x", "loop_y"
<commit_msg>fix problem with position names longer than 1 char.<commit_after>from .bases import CellLoop
from .compatibility import one_dimension, two_dimensions, activity
from .utils import offset_pos

from itertools import product, izip
import numpy as np

class OneDimCellLoop(CellLoop):
    """The OneDimCellLoop iterates over all cells in order from 0 to sizeX."""

    requires_features = [one_dimension]

    def get_pos(self):
        return "loop_x",

    def visit(self):
        super(OneDimCellLoop, self).visit()
        self.code.add_code("loop_begin",
                """for(int loop_x=0; loop_x < sizeX; loop_x++) {""")
        self.code.add_code("loop_end",
                """}""")

    def get_iter(self):
        return iter(izip(xrange(0, self.code.acc.get_size_of())))

    def build_name(self, parts):
        parts.insert(0, "1d")

class TwoDimCellLoop(CellLoop):
    """The TwoDimCellLoop iterates over all cells from left to right, then from
    top to bottom."""

    requires_features = [two_dimensions]

    def get_pos(self):
        return "loop_x", "loop_y"

    def visit(self):
        super(TwoDimCellLoop, self).visit()
        size_names = self.code.acc.size_names
        self.code.add_code("loop_begin",
            """for(int loop_x=0; loop_x < %s; loop_x++) {
                for(int loop_y=0; loop_y < %s; loop_y++) {""" % (size_names))
        self.code.add_code("loop_end",
                """}
                }""")

    def get_iter(self):
        return iter(product(xrange(0, self.code.acc.get_size_of(0)),
                            xrange(0, self.code.acc.get_size_of(1))))

    def build_name(self, parts):
        parts.insert(0, "2d")

class SparseCellLoop(CellLoop):
    """The SparseCellLoopBase offers common code for loops that only calculate
    those fields, where the neighbours have changed in the last step.

    This is based on a list of positions called `sparse_list` as well as a mask
    of booleans called `sparse_mask`, that only internally gets used to make
    sure, that no fields are enlisted more than once.

    The `sparse_list` is duplicated into the property `prev_sparse_list`, from
    which reads are performed. Additionally, the number of active cells is
    taken from an "activity" stats before the loop to see how many valid
    entries exist in the array.

    For the pure-py version, a normal python set is used."""

    requires_features = [activity]

    def set_target(self, target):
        """Adds the activity mask and position list to the target attributes."""
        super(SparseCellLoop, self).set_target(target)

        if len(target.size) == 1:
            self.requires_features.append(one_dimension)
        elif len(target.size) == 2:
            self.requires_features.append(two_dimensions)
        else:
            raise NotImplementedError("SparseCellLoop has not been modified to"
                    " work with more than 2 dimensions.")

        size = self.calculate_size()
        target.sparse_mask = np.zeros(size, dtype=np.bool)
        target.sparse_list = np.zeros(size, dtype=np.int)
        target.sparse_set = set()

    def get_pos(self):
        return self.position_names

    def calculate_size(self):
        """Calculate how big the mask and list have to be.

        The current strategy is to just allocate one field for each field of the
        configuration."""
        return reduce(lambda a, b: a * b, self.target.size)

    def mark_cell_py(self, pos):
        positions = [offset_pos(pos, offs) for offs in
                     self.code.neigh.affected_cells()]
        positions = [pos
                         if self.code.border.is_position_valid(pos)
                         else self.code.border.correct_position(pos)
                     for pos in positions]
        self.target.sparse_set.update(positions)

    def new_config(self):
        size = self.calculate_size()
        print size
        self.target.sparse_mask = np.ones(size, dtype=np.bool)
        self.target.sparse_list = np.array(range(size), dtype=np.int)
        self.target.prev_sparse_list = self.target.sparse_list.copy()
        print self.target.prev_sparse_list
        self.target.sparse_set = set(product(*[range(siz) for siz in self.target.size]))

    def get_iter(self):
        # iterate over a copy of the set, so that it can be modified while running
        print len(self.target.sparse_set)
        the_list = list(self.target.sparse_set)
        self.target.sparse_set.clear()
        return iter(the_list)

    def visit(self):
        super(SparseCellLoop, self).visit()

        self.code.attrs.append("sparse_mask")
        self.code.attrs.append("sparse_list")
        self.code.attrs.append("prev_sparse_list")

        self.code.add_py_hook("loop_end",
            """if was_active: self.loop.mark_cell_py(pos)""")

        self.code.add_code("localvars",
            """int valid_positions = activity(0);
               if(valid_positions == -1) {
                   valid_positions = %(size_a)s * %(size_b)s;
               }
               int sparse_cell_write_idx = 0;""" %
                   dict(size_a=self.code.acc.size_names[0],
                        size_b=self.code.acc.size_names[1] if len(self.code.acc.size_names) == 2 else "1"))
        self.code.add_code("loop_begin",
            """for(int cell_idx=0; cell_idx < valid_positions; cell_idx++) {""")
        if len(self.position_names) == 1:
            self.code.add_code("loop_begin",
                """    int %s = prev_sparse_list(cell_idx);""" % self.position_names)
        elif len(self.position_names) == 2:
            self.code.add_code("loop_begin",
                """    int %(pos_a)s = prev_sparse_list(cell_idx) %% %(size_a)s;
                       int %(pos_b)s = prev_sparse_list(cell_idx) / %(size_b)s;""" %
                           dict(pos_a = self.position_names[0],
                                pos_b = self.position_names[1],
                                size_a = self.code.acc.size_names[0],
                                size_b = self.code.acc.size_names[1]))

        # FIXME use procer position names here
        # FIXME wrap positions around the borders
        if len(self.position_names) == 1:
            self.code.add_code("loop_end",
                    """if(was_active) {
                               %s
                       }""" % ("\n".join([
                           """
                           {int idx = i + %(offs_x)s;
                           if(!sparse_mask(idx)) {
                               sparse_list(sparse_cell_write_idx) = idx;
                               sparse_mask(idx) = true;
                               sparse_cell_write_idx++;
                           }}""" % dict(offs_x=offs[0])
                               for offs in self.code.neigh.offsets])))
        elif len(self.position_names) == 2:
            self.code.add_code("loop_end",
                    """if(was_active) {
                               %s
                       }""" % ("\n".join([
                           """
                           {int px = i + %(offs_x)s;
                           int py = j + %(offs_y)s;
                           int idx = px * %(size_x)s + py;
                           if(!sparse_mask(idx)) {
                               sparse_list(sparse_cell_write_idx) = idx;
                               sparse_mask(idx) = true;
                               sparse_cell_write_idx++;
                           }}""" % dict(offs_x=offs[0], offs_y=offs[1],
                                       size_x=self.code.acc.size_names[0])
                               for offs in self.code.neigh.offsets])))
        self.code.add_code("loop_end",
                """
                }
                sparse_mask = sparse_mask*0;""")

class OneDimSparseCellLoop(SparseCellLoop):
    def __init__(self):
        self.position_names = "loop_x",

class TwoDimSparseCellLoop(SparseCellLoop):
    def __init__(self):
        self.position_names = "loop_x", "loop_y"
<|endoftext|>