<commit_before># Copyright (C) 2013-2015 Nippon Telegraph and Telephone Corporation.
# Copyright (C) 2013-2015 YAMAMOTO Takashi <yamamoto at valinux co jp>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numbers
import time
import unittest
from nose.tools import raises
import six

from ryu.lib import hub
hub.patch()
from ryu.lib import rpc

class MyException(BaseException):
    pass

class Test_rpc(unittest.TestCase):
    """ Test case for ryu.lib.rpc
    """

    def _handle_request(self, m):
        e = self._server_endpoint
        msgid, method, params = m
        if method == b'resp':
            e.send_response(msgid, result=params[0])
        elif method == b'err':
            e.send_response(msgid, error=params[0])
        elif method == b'callback':
            n, cb, v = params
            assert n > 0
            self._requests.add(e.send_request(cb, [msgid, n, cb, v]))
        elif method == b'notify1':
            e.send_notification(params[1], params[2])
            e.send_response(msgid, result=params[0])
        elif method == b'shutdown':
            import socket
            # Though six.text_type is not needed in python2, it is
            # unconditionally applied for code simplicityp
            how = getattr(socket, six.text_type(params[0], 'utf-8'))
            self._server_sock.shutdown(how)
            e.send_response(msgid, result=method)
        else:
            raise Exception("unknown method %s" % method)

    def _handle_notification(self, m):
        e = self._server_endpoint
        method, params = m
        if method == b'notify2':
            e.send_notification(params[0], params[1])

    def _handle_response(self, m):
        e = self._server_endpoint
        msgid, error, result = m
        assert error is None
        self._requests.remove(msgid)
        omsgid, n, cb, v = result
        assert n >= 0
        if n == 0:
            e.send_response(omsgid, result=v)
        else:
            self._requests.add(e.send_request(cb, [omsgid, n, cb, v]))

    def setUp(self):
        import socket

        self._server_sock, self._client_sock = socket.socketpair()
        table = {
            rpc.MessageType.REQUEST: self._handle_request,
            rpc.MessageType.RESPONSE: self._handle_response,
            rpc.MessageType.NOTIFY: self._handle_notification
        }
        self._requests = set()
        self._server_sock.setblocking(0)
        self._server_endpoint = rpc.EndPoint(self._server_sock,
                                             disp_table=table)
        self._server_thread = hub.spawn(self._server_endpoint.serve)

    def tearDown(self):
        hub.kill(self._server_thread)
        hub.joinall([self._server_thread])

    def test_0_call_str(self):
        c = rpc.Client(self._client_sock)
        obj = b'hoge'
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_0_call_int(self):
        c = rpc.Client(self._client_sock)
        obj = 12345
        assert isinstance(obj, int)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_int2(self):
        c = rpc.Client(self._client_sock)
        obj = six.MAXSIZE
        assert isinstance(obj, int)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_int3(self):
        c = rpc.Client(self._client_sock)
        obj = - six.MAXSIZE - 1
        assert isinstance(obj, int)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_long(self):
        c = rpc.Client(self._client_sock)
        obj = 0xffffffffffffffff  # max value for msgpack
        _long = int if six.PY3 else long
        assert isinstance(obj, _long)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_long2(self):
        c = rpc.Client(self._client_sock)
        # NOTE: the python type of this value is int for 64-bit arch
        obj = -0x8000000000000000  # min value for msgpack
        assert isinstance(obj, numbers.Integral)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    @raises(TypeError)
    def test_0_call_bytearray(self):
        c = rpc.Client(self._client_sock)
        obj = bytearray(b'foo')
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_1_shutdown_wr(self):
        # test if the server shutdown on disconnect
        import socket
        self._client_sock.shutdown(socket.SHUT_WR)
        hub.joinall([self._server_thread])

    @raises(EOFError)
    def test_1_client_shutdown_wr(self):
        c = rpc.Client(self._client_sock)
        c.call(b'shutdown', [b'SHUT_WR'])

    def test_1_call_True(self):
        c = rpc.Client(self._client_sock)
        obj = True
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_None(self):
        c = rpc.Client(self._client_sock)
        obj = None
        assert c.call(b'resp', [obj]) is None

    def test_2_call_False(self):
        c = rpc.Client(self._client_sock)
        obj = False
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_dict(self):
        c = rpc.Client(self._client_sock)
        obj = {b'hoge': 1, b'fuga': 2}
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_empty_dict(self):
        c = rpc.Client(self._client_sock)
        obj = {}
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_array(self):
        c = rpc.Client(self._client_sock)
        obj = [1, 2, 3, 4]
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_empty_array(self):
        c = rpc.Client(self._client_sock)
        obj = []
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_tuple(self):
        c = rpc.Client(self._client_sock)
        # note: msgpack library implicitly convert a tuple into a list
        obj = (1, 2, 3)
        assert c.call(b'resp', [obj]) == list(obj)

    @raises(TypeError)
    def test_2_call_unicode(self):
        c = rpc.Client(self._client_sock)
        # note: on-wire msgpack has no notion of encoding.
        # the msgpack library implicitly converts unicode to
        # utf-8 encoded bytes by default.
        # we don't want to rely on the behaviour though because
        # it seems to be going to change.
        # https://gist.github.com/methane/5022403
        obj = u"hoge"
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_2_call_small_binary(self):
        import struct
        c = rpc.Client(self._client_sock)
        obj = struct.pack("100x")
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_3_call_complex(self):
        c = rpc.Client(self._client_sock)
        obj = [1, b'hoge', {b'foo': 1, 3: b'bar'}]
        assert c.call(b'resp', [obj]) == list(obj)

    def test_4_call_large_binary(self):
        import struct

        c = rpc.Client(self._client_sock)
        obj = struct.pack("10000000x")
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_0_notification1(self):
        l = []

        def callback(n):
            l.append(n)
        c = rpc.Client(self._client_sock, notification_callback=callback)
        obj = b'hogehoge'
        robj = b'fugafuga'
        assert c.call(b'notify1', [robj, b'notify_hoge', [obj]]) == robj
        c.receive_notification()
        assert len(l) == 1
        n = l.pop(0)
        assert n is not None
        method, params = n
        assert method == b'notify_hoge'
        assert params[0] == obj

    def test_0_notification2(self):
        l = []

        def callback(n):
            l.append(n)
        c = rpc.Client(self._client_sock, notification_callback=callback)
        obj = b'hogehogehoge'
        c.send_notification(b'notify2', [b'notify_hoge', [obj]])
        c.receive_notification()
        assert len(l) == 1
        n = l.pop(0)
        assert n is not None
        method, params = n
        assert method == b'notify_hoge'
        assert params[0] == obj

    def test_0_call_error(self):
        c = rpc.Client(self._client_sock)
        obj = b'hoge'
        try:
            c.call(b'err', [obj])
            raise Exception("unexpected")
        except rpc.RPCError as e:
            assert e.get_value() == obj

    def test_0_call_error_notification(self):
        l = []

        def callback(n):
            l.append(n)
        c = rpc.Client(self._client_sock, notification_callback=callback)
        c.send_notification(b'notify2', [b'notify_foo', []])
        hub.sleep(0.5)  # give the peer a chance to run
        obj = b'hoge'
        try:
            c.call(b'err', [obj])
            raise Exception("unexpected")
        except rpc.RPCError as e:
            assert e.get_value() == obj
        assert len(l) == 1
        n = l.pop(0)
        method, params = n
        assert method == b'notify_foo'
        assert params == []

    def test_4_async_call(self):
        """send a bunch of requests and then wait for responses
        """
        num_calls = 9999
        old_blocking = self._client_sock.setblocking(0)
        try:
            e = rpc.EndPoint(self._client_sock)
            s = set()
            for i in range(1, num_calls + 1):
                s.add(e.send_request(b'resp', [i]))
            sum = 0
            while s:
                e.block()
                e.process()
                done = set()
                for x in s:
                    r = e.get_response(x)
                    if r is None:
                        continue
                    res, error = r
                    assert error is None
                    sum += res
                    done.add(x)
                assert done.issubset(s)
                s -= done
            assert sum == (1 + num_calls) * num_calls / 2
        finally:
            self._client_sock.setblocking(old_blocking)

    def test_4_async_call2(self):
        """both sides act as rpc client and server
        """
        assert not self._requests
        num_calls = 100
        old_blocking = self._client_sock.setblocking(0)
        try:
            e = rpc.EndPoint(self._client_sock)
            s = set()
            for i in range(1, num_calls + 1):
                s.add(e.send_request(b'callback', [i, b'ourcallback', 0]))
            sum = 0
            while s:
                e.block()
                e.process()
                done = set()
                for x in s:
                    r = e.get_response(x)
                    if r is None:
                        continue
                    res, error = r
                    assert error is None
                    sum += res
                    done.add(x)
                assert done.issubset(s)
                s -= done
                r = e.get_request()
                if r is not None:
                    msgid, method, params = r
                    assert method == b'ourcallback'
                    omsgid, n, cb, v = params
                    assert omsgid in s
                    assert cb == b'ourcallback'
                    assert n > 0
                    e.send_response(msgid, result=[omsgid, n - 1, cb, v + 1])
            assert sum == (1 + num_calls) * num_calls / 2
        finally:
            self._client_sock.setblocking(old_blocking)
        assert not self._requests
<commit_msg>test: disable rpc unittest with large file for now<commit_after># Copyright (C) 2013-2015 Nippon Telegraph and Telephone Corporation.
# Copyright (C) 2013-2015 YAMAMOTO Takashi <yamamoto at valinux co jp>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numbers
import time
import sys
if sys.version_info < (2, 7):
    import unittest2 as unittest
else:
    import unittest
from nose.tools import raises
import six

from ryu.lib import hub
hub.patch()
from ryu.lib import rpc

class MyException(BaseException):
    pass

class Test_rpc(unittest.TestCase):
    """ Test case for ryu.lib.rpc
    """

    def _handle_request(self, m):
        e = self._server_endpoint
        msgid, method, params = m
        if method == b'resp':
            e.send_response(msgid, result=params[0])
        elif method == b'err':
            e.send_response(msgid, error=params[0])
        elif method == b'callback':
            n, cb, v = params
            assert n > 0
            self._requests.add(e.send_request(cb, [msgid, n, cb, v]))
        elif method == b'notify1':
            e.send_notification(params[1], params[2])
            e.send_response(msgid, result=params[0])
        elif method == b'shutdown':
            import socket
            # Though six.text_type is not needed in python2, it is
            # unconditionally applied for code simplicityp
            how = getattr(socket, six.text_type(params[0], 'utf-8'))
            self._server_sock.shutdown(how)
            e.send_response(msgid, result=method)
        else:
            raise Exception("unknown method %s" % method)

    def _handle_notification(self, m):
        e = self._server_endpoint
        method, params = m
        if method == b'notify2':
            e.send_notification(params[0], params[1])

    def _handle_response(self, m):
        e = self._server_endpoint
        msgid, error, result = m
        assert error is None
        self._requests.remove(msgid)
        omsgid, n, cb, v = result
        assert n >= 0
        if n == 0:
            e.send_response(omsgid, result=v)
        else:
            self._requests.add(e.send_request(cb, [omsgid, n, cb, v]))

    def setUp(self):
        import socket

        self._server_sock, self._client_sock = socket.socketpair()
        table = {
            rpc.MessageType.REQUEST: self._handle_request,
            rpc.MessageType.RESPONSE: self._handle_response,
            rpc.MessageType.NOTIFY: self._handle_notification
        }
        self._requests = set()
        self._server_sock.setblocking(0)
        self._server_endpoint = rpc.EndPoint(self._server_sock,
                                             disp_table=table)
        self._server_thread = hub.spawn(self._server_endpoint.serve)

    def tearDown(self):
        hub.kill(self._server_thread)
        hub.joinall([self._server_thread])

    def test_0_call_str(self):
        c = rpc.Client(self._client_sock)
        obj = b'hoge'
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_0_call_int(self):
        c = rpc.Client(self._client_sock)
        obj = 12345
        assert isinstance(obj, int)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_int2(self):
        c = rpc.Client(self._client_sock)
        obj = six.MAXSIZE
        assert isinstance(obj, int)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_int3(self):
        c = rpc.Client(self._client_sock)
        obj = - six.MAXSIZE - 1
        assert isinstance(obj, int)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_long(self):
        c = rpc.Client(self._client_sock)
        obj = 0xffffffffffffffff  # max value for msgpack
        _long = int if six.PY3 else long
        assert isinstance(obj, _long)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    def test_0_call_long2(self):
        c = rpc.Client(self._client_sock)
        # NOTE: the python type of this value is int for 64-bit arch
        obj = -0x8000000000000000  # min value for msgpack
        assert isinstance(obj, numbers.Integral)
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, type(obj))

    @raises(TypeError)
    def test_0_call_bytearray(self):
        c = rpc.Client(self._client_sock)
        obj = bytearray(b'foo')
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_1_shutdown_wr(self):
        # test if the server shutdown on disconnect
        import socket
        self._client_sock.shutdown(socket.SHUT_WR)
        hub.joinall([self._server_thread])

    @raises(EOFError)
    def test_1_client_shutdown_wr(self):
        c = rpc.Client(self._client_sock)
        c.call(b'shutdown', [b'SHUT_WR'])

    def test_1_call_True(self):
        c = rpc.Client(self._client_sock)
        obj = True
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_None(self):
        c = rpc.Client(self._client_sock)
        obj = None
        assert c.call(b'resp', [obj]) is None

    def test_2_call_False(self):
        c = rpc.Client(self._client_sock)
        obj = False
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_dict(self):
        c = rpc.Client(self._client_sock)
        obj = {b'hoge': 1, b'fuga': 2}
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_empty_dict(self):
        c = rpc.Client(self._client_sock)
        obj = {}
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_array(self):
        c = rpc.Client(self._client_sock)
        obj = [1, 2, 3, 4]
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_empty_array(self):
        c = rpc.Client(self._client_sock)
        obj = []
        assert c.call(b'resp', [obj]) == obj

    def test_2_call_tuple(self):
        c = rpc.Client(self._client_sock)
        # note: msgpack library implicitly convert a tuple into a list
        obj = (1, 2, 3)
        assert c.call(b'resp', [obj]) == list(obj)

    @raises(TypeError)
    def test_2_call_unicode(self):
        c = rpc.Client(self._client_sock)
        # note: on-wire msgpack has no notion of encoding.
        # the msgpack library implicitly converts unicode to
        # utf-8 encoded bytes by default.
        # we don't want to rely on the behaviour though because
        # it seems to be going to change.
        # https://gist.github.com/methane/5022403
        obj = u"hoge"
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_2_call_small_binary(self):
        import struct
        c = rpc.Client(self._client_sock)
        obj = struct.pack("100x")
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_3_call_complex(self):
        c = rpc.Client(self._client_sock)
        obj = [1, b'hoge', {b'foo': 1, 3: b'bar'}]
        assert c.call(b'resp', [obj]) == list(obj)

    @unittest.skip("doesn't work with eventlet 0.18 and later")
    def test_4_call_large_binary(self):
        import struct

        c = rpc.Client(self._client_sock)
        obj = struct.pack("10000000x")
        result = c.call(b'resp', [obj])
        assert result == obj
        assert isinstance(result, bytes)

    def test_0_notification1(self):
        l = []

        def callback(n):
            l.append(n)
        c = rpc.Client(self._client_sock, notification_callback=callback)
        obj = b'hogehoge'
        robj = b'fugafuga'
        assert c.call(b'notify1', [robj, b'notify_hoge', [obj]]) == robj
        c.receive_notification()
        assert len(l) == 1
        n = l.pop(0)
        assert n is not None
        method, params = n
        assert method == b'notify_hoge'
        assert params[0] == obj

    def test_0_notification2(self):
        l = []

        def callback(n):
            l.append(n)
        c = rpc.Client(self._client_sock, notification_callback=callback)
        obj = b'hogehogehoge'
        c.send_notification(b'notify2', [b'notify_hoge', [obj]])
        c.receive_notification()
        assert len(l) == 1
        n = l.pop(0)
        assert n is not None
        method, params = n
        assert method == b'notify_hoge'
        assert params[0] == obj

    def test_0_call_error(self):
        c = rpc.Client(self._client_sock)
        obj = b'hoge'
        try:
            c.call(b'err', [obj])
            raise Exception("unexpected")
        except rpc.RPCError as e:
            assert e.get_value() == obj

    def test_0_call_error_notification(self):
        l = []

        def callback(n):
            l.append(n)
        c = rpc.Client(self._client_sock, notification_callback=callback)
        c.send_notification(b'notify2', [b'notify_foo', []])
        hub.sleep(0.5)  # give the peer a chance to run
        obj = b'hoge'
        try:
            c.call(b'err', [obj])
            raise Exception("unexpected")
        except rpc.RPCError as e:
            assert e.get_value() == obj
        assert len(l) == 1
        n = l.pop(0)
        method, params = n
        assert method == b'notify_foo'
        assert params == []

    def test_4_async_call(self):
        """send a bunch of requests and then wait for responses
        """
        num_calls = 9999
        old_blocking = self._client_sock.setblocking(0)
        try:
            e = rpc.EndPoint(self._client_sock)
            s = set()
            for i in range(1, num_calls + 1):
                s.add(e.send_request(b'resp', [i]))
            sum = 0
            while s:
                e.block()
                e.process()
                done = set()
                for x in s:
                    r = e.get_response(x)
                    if r is None:
                        continue
                    res, error = r
                    assert error is None
                    sum += res
                    done.add(x)
                assert done.issubset(s)
                s -= done
            assert sum == (1 + num_calls) * num_calls / 2
        finally:
            self._client_sock.setblocking(old_blocking)

    def test_4_async_call2(self):
        """both sides act as rpc client and server
        """
        assert not self._requests
        num_calls = 100
        old_blocking = self._client_sock.setblocking(0)
        try:
            e = rpc.EndPoint(self._client_sock)
            s = set()
            for i in range(1, num_calls + 1):
                s.add(e.send_request(b'callback', [i, b'ourcallback', 0]))
            sum = 0
            while s:
                e.block()
                e.process()
                done = set()
                for x in s:
                    r = e.get_response(x)
                    if r is None:
                        continue
                    res, error = r
                    assert error is None
                    sum += res
                    done.add(x)
                assert done.issubset(s)
                s -= done
                r = e.get_request()
                if r is not None:
                    msgid, method, params = r
                    assert method == b'ourcallback'
                    omsgid, n, cb, v = params
                    assert omsgid in s
                    assert cb == b'ourcallback'
                    assert n > 0
                    e.send_response(msgid, result=[omsgid, n - 1, cb, v + 1])
            assert sum == (1 + num_calls) * num_calls / 2
        finally:
            self._client_sock.setblocking(old_blocking)
        assert not self._requests
<|endoftext|>


<commit_before>#!/usr/bin/env python3
import os
import sys
import subprocess
import logging
from collections import defaultdict
import multiprocessing
import argparse

from fasta_utils import read_fasta_file_handle, format_seq
from fastq_utils import read_fastq_file_handle

from assembler_factory import AssemblerFactory

logger = logging.getLogger(__name__)

def extract_reads_by_component(fastq, read_metanode_component_filepath):
    if not os.path.isfile(fastq):
        logger.fatal('The input reads file does not exists:%s' % fastq)
        sys.exit("An error occured. Can't extract component's reads")

    if not os.path.isfile(read_metanode_component_filepath):
        logger.fatal('The file storing the correspondance between reads and components does not exists:%s' % read_metanode_component_filepath)
        sys.exit("An error occured. Can't extract component's reads")

    # Reading read --> component file
    logger.debug('Reading read-->component from {}'.format(read_metanode_component_filepath))
    read_component_dict = dict()
    with open(read_metanode_component_filepath, 'r') as read_metanode_component_fh:
        read_component_dict = {t[0]:t[2] for t in (l.split() for l in read_metanode_component_fh) if t[2] != 'NULL'}

    # Storing reads for each component
    logger.debug('Storing reads by component from {}'.format(fastq))
    component_reads_dict = defaultdict(list)
    with open(fastq, 'r') as fastq_fh:
        for header, seq, qual in read_fastq_file_handle(fastq_fh):
            try:
                component_reads_dict[read_component_dict[header]].append(tuple((header, seq, qual)))
            except KeyError:
                pass
    return component_reads_dict

def extract_lca_by_component(components_lca_filepath):
    if not os.path.isfile(components_lca_filepath):
        logger.fatal('The file storing the correspondance between lca info and components does not exists:%s' % components_lca_filepath)
        sys.exit("An error occured. Can't extract component's lca")

    # Reading components LCA and storing them in a dict
    logger.debug('Reading components LCA assignment from {0}'.format(components_lca_filepath))
    component_lca_dict = dict()
    with open(components_lca_filepath, 'r') as component_lca_fh:
        component_lca_dict = {t[0]:t[1] for t in (l.split() for l in component_lca_fh) if len(t) == 2}
    return component_lca_dict

def save_components(component_reads_dict, directory):
    """
    Take a dict (key=component_id, value=[header,seq,qual])
    and save each component to a file into the given directory:
    directory/
        component_%s.fq % component_id

    Return a dict (key=component_id, value=fastq_path)
    """

    try:
        os.mkdir(directory)
    except FileExistsError as fee:
        pass

    components_fq = {}
    for component_id, reads_list in component_reads_dict.items():
        fq_name = "component%s_reads.fq" % component_id
        fq_path = os.path.join(directory, fq_name)
        #logger.debug("Save component %s into %s" % (component_id, fq_path))
        components_fq[component_id] = fq_path

        with open(fq_path, 'w') as component_fh:
            for (header, seq, qual) in reads_list:
                component_fh.write('@{}\n{}\n+\n{}\n'.format(header, seq, qual))
    return components_fq

def assemble_component(assembler_name,
                       in_fastq, workdir,
                       read_correction, cpu):
    #logger.debug('Assembling: %s' % in_fastq)
    assembler_factory = AssemblerFactory()
    assembler = assembler_factory.get(assembler_name)
    assembler.build_command_line(in_fastq, workdir, read_correction, cpu)

    fasta_file = assembler.run()
    return fasta_file

def concat_components_fasta_with_lca(assembled_components_fasta, contigs_fasta, component_lca_dict):
    if os.path.isfile(contigs_fasta):
        #logger.debug("Remove old contig fasta file:%s" % contigs_fasta)
        os.unlink(contigs_fasta)

    component_lca = 'NULL'
    contigs_fh = open(contigs_fasta, 'w')
    contig_count = 0
    for component_id, component_fasta in assembled_components_fasta.items():
        if component_id in component_lca_dict:
            component_lca = component_lca_dict[component_id]
        with open(component_fasta, 'r') as sga_contigs_fh:
            for header, seq in read_fasta_file_handle(sga_contigs_fh):
                if len(seq):
                    contig_count += 1
                    contigs_fh.write('>{0} component={1} '.format(contig_count, component_id))
                    contigs_fh.write('lca={0}\n{1}\n'.format(component_lca, format_seq(seq)))

    contigs_fh.close()

def _get_workdir(fq):
    """
    Convenient function to build a workdir from the name of the fastqfile
    """
    sga_wkdir_basename, _ = os.path.splitext(fq)
    return '%s_assembly_wkdir' % sga_wkdir_basename

def assemble_all_components(assembler_name,
                            fastq, read_metanode_component_filepath, components_lca_filepath,
                            out_contigs_fasta, workdir,
                            cpu, read_correction):

    logger.info("Save components to fastq files")
    components_dict = extract_reads_by_component(fastq, read_metanode_component_filepath)
    components_reads_fq = save_components(components_dict, workdir).items()
    assembled_components_fasta = {}

    logger.info("Assemble components")

    # Foreach component, build the parameters used by assemble_component and save
    # them into a list to be able to apply a map function on it
    params = []
    component_id_list = []
    for component_id, fq in components_reads_fq:
        params.append((assembler_name, fq, _get_workdir(fq), read_correction, 1))
        component_id_list.append(component_id)

    with multiprocessing.Pool(processes=cpu) as pool:
        try:
            fasta_list = pool.starmap(assemble_component, params)
        except subprocess.CalledProcessError as cpe:
            logger.fatal('Command %s returned non-zero exit status %s' % (cpe.cmd, cpe.returncode))
            sys.exit('Components assembly step failed')

    # Make the correspondance between the component_id and the fasta file
    assembled_components_fasta = dict(zip(component_id_list, fasta_list))

    lca_dict = extract_lca_by_component(components_lca_filepath)
    logger.info("Pool components contigs into: %s" % out_contigs_fasta)
    concat_components_fasta_with_lca(assembled_components_fasta,
                                     out_contigs_fasta, lca_dict)

if __name__ == '__main__':

    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(description='Assemble components')
    parser.add_argument('-a', '--assembler',
                        choices=[a.name() for a in AssemblerFactory.ASSEMBLER_ENGINES],
                        help="Select the assembler to be used. Default is %(default)s",
                        default="SGA")
    parser.add_argument('-i', '--input_fastq',
                        type=argparse.FileType('r'),
                        help='input fastq file',
                        required=True)
    parser.add_argument('-m', '--reads_metanode',
                        type=argparse.FileType('r'),
                        help='This file makes the correspondance between reads and the components',
                        required=True)
    parser.add_argument('-l', '--components_lca',
                        type=argparse.FileType('r'),
                        help='This file make the correspondance between lca and the components',
                        required=True)
    parser.add_argument('-w', '--workdir',
                        action = 'store',
                        type = str,
                        help = 'Working  directory')
    parser.add_argument('-o', '--output_fasta',
                        help='output fasta file',
                        required=True)
    parser.add_argument('--cpu',
                        action = 'store',
                        metavar = 'CPU',
                        type = int,
                        default = 1,
                        help = 'Max number of CPU to use. '
                        'Default is %(default)s cpu')
    parser.add_argument('--read_correction',
                        action = 'store',
                        type = str,
                        choices = ['no', 'yes', 'auto'],
                        default = 'no',
                        help = 'Set the assembler read correction step. '
                        'Default is %(default)s')

    args = parser.parse_args()
    assemble_all_components(args.assembler,
                            args.input_fastq.name, args.reads_metanode.name, args.components_lca.name,
                            args.output_fasta, args.workdir,
                            args.cpu, args.read_correction)
<commit_msg>Remove last occurences to SGA<commit_after>#!/usr/bin/env python3
import os
import sys
import subprocess
import logging
from collections import defaultdict
import multiprocessing
import argparse

from fasta_utils import read_fasta_file_handle, format_seq
from fastq_utils import read_fastq_file_handle

from assembler_factory import AssemblerFactory

logger = logging.getLogger(__name__)

def extract_reads_by_component(fastq, read_metanode_component_filepath):
    if not os.path.isfile(fastq):
        logger.fatal('The input reads file does not exists:%s' % fastq)
        sys.exit("An error occured. Can't extract component's reads")

    if not os.path.isfile(read_metanode_component_filepath):
        logger.fatal('The file storing the correspondance between reads and components does not exists:%s' % read_metanode_component_filepath)
        sys.exit("An error occured. Can't extract component's reads")

    # Reading read --> component file
    logger.debug('Reading read-->component from {}'.format(read_metanode_component_filepath))
    read_component_dict = dict()
    with open(read_metanode_component_filepath, 'r') as read_metanode_component_fh:
        read_component_dict = {t[0]:t[2] for t in (l.split() for l in read_metanode_component_fh) if t[2] != 'NULL'}

    # Storing reads for each component
    logger.debug('Storing reads by component from {}'.format(fastq))
    component_reads_dict = defaultdict(list)
    with open(fastq, 'r') as fastq_fh:
        for header, seq, qual in read_fastq_file_handle(fastq_fh):
            try:
                component_reads_dict[read_component_dict[header]].append(tuple((header, seq, qual)))
            except KeyError:
                pass
    return component_reads_dict

def extract_lca_by_component(components_lca_filepath):
    if not os.path.isfile(components_lca_filepath):
        logger.fatal('The file storing the correspondance between lca info and components does not exists:%s' % components_lca_filepath)
        sys.exit("An error occured. Can't extract component's lca")

    # Reading components LCA and storing them in a dict
    logger.debug('Reading components LCA assignment from {0}'.format(components_lca_filepath))
    component_lca_dict = dict()
    with open(components_lca_filepath, 'r') as component_lca_fh:
        component_lca_dict = {t[0]:t[1] for t in (l.split() for l in component_lca_fh) if len(t) == 2}
    return component_lca_dict

def save_components(component_reads_dict, directory):
    """
    Take a dict (key=component_id, value=[header,seq,qual])
    and save each component to a file into the given directory:
    directory/
        component_%s.fq % component_id

    Return a dict (key=component_id, value=fastq_path)
    """

    try:
        os.mkdir(directory)
    except FileExistsError as fee:
        pass

    components_fq = {}
    for component_id, reads_list in component_reads_dict.items():
        fq_name = "component%s_reads.fq" % component_id
        fq_path = os.path.join(directory, fq_name)
        #logger.debug("Save component %s into %s" % (component_id, fq_path))
        components_fq[component_id] = fq_path

        with open(fq_path, 'w') as component_fh:
            for (header, seq, qual) in reads_list:
                component_fh.write('@{}\n{}\n+\n{}\n'.format(header, seq, qual))
    return components_fq

def assemble_component(assembler_name,
                       in_fastq, workdir,
                       read_correction, cpu):
    #logger.debug('Assembling: %s' % in_fastq)
    assembler_factory = AssemblerFactory()
    assembler = assembler_factory.get(assembler_name)
    assembler.build_command_line(in_fastq, workdir, read_correction, cpu)

    fasta_file = assembler.run()
    return fasta_file

def concat_components_fasta_with_lca(assembled_components_fasta, contigs_fasta, component_lca_dict):
    if os.path.isfile(contigs_fasta):
        #logger.debug("Remove old contig fasta file:%s" % contigs_fasta)
        os.unlink(contigs_fasta)

    component_lca = 'NULL'
    contigs_fh = open(contigs_fasta, 'w')
    contig_count = 0
    for component_id, component_fasta in assembled_components_fasta.items():
        if component_id in component_lca_dict:
            component_lca = component_lca_dict[component_id]
        with open(component_fasta, 'r') as assembler_contigs_fh:
            for header, seq in read_fasta_file_handle(assembler_contigs_fh):
                if len(seq):
                    contig_count += 1
                    contigs_fh.write('>{0} component={1} '.format(contig_count, component_id))
                    contigs_fh.write('lca={0}\n{1}\n'.format(component_lca, format_seq(seq)))

    contigs_fh.close()

def _get_workdir(fq):
    """
    Convenient function to build a workdir from the name of the fastqfile
    """
    assembler_wkdir_basename, _ = os.path.splitext(fq)
    return '%s_assembly_wkdir' % assembler_wkdir_basename

def assemble_all_components(assembler_name,
                            fastq, read_metanode_component_filepath, components_lca_filepath,
                            out_contigs_fasta, workdir,
                            cpu, read_correction):

    logger.info("Save components to fastq files")
    components_dict = extract_reads_by_component(fastq, read_metanode_component_filepath)
    components_reads_fq = save_components(components_dict, workdir).items()
    assembled_components_fasta = {}

    logger.info("Assemble components")

    # Foreach component, build the parameters used by assemble_component and save
    # them into a list to be able to apply a map function on it
    params = []
    component_id_list = []
    for component_id, fq in components_reads_fq:
        params.append((assembler_name, fq, _get_workdir(fq), read_correction, 1))
        component_id_list.append(component_id)

    with multiprocessing.Pool(processes=cpu) as pool:
        try:
            fasta_list = pool.starmap(assemble_component, params)
        except subprocess.CalledProcessError as cpe:
            logger.fatal('Command %s returned non-zero exit status %s' % (cpe.cmd, cpe.returncode))
            sys.exit('Components assembly step failed')

    # Make the correspondance between the component_id and the fasta file
    assembled_components_fasta = dict(zip(component_id_list, fasta_list))

    lca_dict = extract_lca_by_component(components_lca_filepath)
    logger.info("Pool components contigs into: %s" % out_contigs_fasta)
    concat_components_fasta_with_lca(assembled_components_fasta,
                                     out_contigs_fasta, lca_dict)

if __name__ == '__main__':

    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(description='Assemble components')
    parser.add_argument('-a', '--assembler',
                        choices=[a.name() for a in AssemblerFactory.ASSEMBLER_ENGINES],
                        help="Select the assembler to be used. Default is %(default)s",
                        default="SGA")
    parser.add_argument('-i', '--input_fastq',
                        type=argparse.FileType('r'),
                        help='input fastq file',
                        required=True)
    parser.add_argument('-m', '--reads_metanode',
                        type=argparse.FileType('r'),
                        help='This file makes the correspondance between reads and the components',
                        required=True)
    parser.add_argument('-l', '--components_lca',
                        type=argparse.FileType('r'),
                        help='This file make the correspondance between lca and the components',
                        required=True)
    parser.add_argument('-w', '--workdir',
                        action = 'store',
                        type = str,
                        help = 'Working  directory')
    parser.add_argument('-o', '--output_fasta',
                        help='output fasta file',
                        required=True)
    parser.add_argument('--cpu',
                        action = 'store',
                        metavar = 'CPU',
                        type = int,
                        default = 1,
                        help = 'Max number of CPU to use. '
                        'Default is %(default)s cpu')
    parser.add_argument('--read_correction',
                        action = 'store',
                        type = str,
                        choices = ['no', 'yes', 'auto'],
                        default = 'no',
                        help = 'Set the assembler read correction step. '
                        'Default is %(default)s')

    args = parser.parse_args()
    assemble_all_components(args.assembler,
                            args.input_fastq.name, args.reads_metanode.name, args.components_lca.name,
                            args.output_fasta, args.workdir,
                            args.cpu, args.read_correction)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4

# The MIT License (MIT)
# Copyright (C) 2015 - George Y. Kussumoto <contato@georgeyk.com.br>
#
# 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.
#

# Not an extensive implementation, more information at:
# http://docs.ckan.org/en/ckan-2.2.2/api.html#get-able-api-functions

from __future__ import unicode_literals
import logging

try:
    from urllib.parse import urlencode
except ImportError:
    from urllib import urlencode

import requests

from .exceptions import DadosGovBRException
from .factory import model_dict_factory, model_list_factory

logger = logging.getLogger(__name__)

class DadosGovBR(object):
    base_api_url = 'http://dados.gov.br/api'
    api_version = 3
    endpoint = '{}/{}/action'.format(base_api_url, api_version)
    timeout = 4

    def _request(self, resource, **params):
        url = '{endpoint}/{resource}?{params}'.format(
            endpoint=self.endpoint, resource=resource, params=urlencode(params))
        logger.info('Request URL: {}'.format(url))

        response = requests.get(url, timeout=self.timeout)
        logger.info('Response: status {}'.format(response.status_code))

        data = response.json()
        if data['success']:
            return data['result']
        else:
            logger.debug('Response error: headers={}, content={}, '.format(response.headers, response.content))
            raise DadosGovBRException(data['error'])

    #
    # Public API
    #

    def search_resources(self, query, search_key='name', **kwargs):
        query = '{}:{}'.format(search_key, query)
        result = self._request('resource_search', query=query, **kwargs)
        for resource in result['results']:
            yield model_dict_factory('resource', resource)

    def search_dataset(self, query, **kwargs):
        result = self._request('package_search', q=query, **kwargs)
        for dataset in result['results']:
            yield model_dict_factory('dataset', dataset)

    def get_datasets(self):
        return model_list_factory('dataset item', self._request('package_list'))

    def get_dataset(self, dataset_id):
        dataset = self._request('package_show', id=dataset_id)
        return model_dict_factory('dataset', dataset)

    def get_tags(self):
        tags = self._request('tag_list')
        return model_list_factory('tag item', tags)

    def get_tag(self, tag_id):
        tag = self._request('tag_show', id=tag_id)
        return model_dict_factory('tag', tag)
<commit_msg>Refactor: simplify code<commit_after># -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4

# The MIT License (MIT)
# Copyright (C) 2015 - George Y. Kussumoto <contato@georgeyk.com.br>
#
# 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.
#

# Not an extensive implementation, more information at:
# http://docs.ckan.org/en/ckan-2.2.2/api.html#get-able-api-functions

from __future__ import unicode_literals
import logging

try:
    from urllib.parse import urlencode
except ImportError:
    from urllib import urlencode

import requests

from .exceptions import DadosGovBRException
from .factory import model_dict_factory, model_list_factory

logger = logging.getLogger(__name__)

class DadosGovBR(object):
    base_api_url = 'http://dados.gov.br/api'
    api_version = 3
    endpoint = '{}/{}/action'.format(base_api_url, api_version)
    timeout = 4

    def _request(self, resource, **params):
        url = '{endpoint}/{resource}?{params}'.format(
            endpoint=self.endpoint, resource=resource, params=urlencode(params))
        logger.info('Request URL: {}'.format(url))

        response = requests.get(url, timeout=self.timeout)
        logger.info('Response: status {}'.format(response.status_code))

        data = response.json()
        if data['success']:
            return data['result']
        else:
            logger.debug('Response error: headers={}, content={}, '.format(response.headers, response.content))
            raise DadosGovBRException(data['error'])

    #
    # Public API
    #

    def search_resources(self, query, search_key='name', **kwargs):
        query = '{}:{}'.format(search_key, query)
        result = self._request('resource_search', query=query, **kwargs)
        return model_list_factory('resource', result['results'])

    def search_datasets(self, query, **kwargs):
        result = self._request('package_search', q=query, **kwargs)
        return model_list_factory('dataset', result['results'])

    def get_datasets(self):
        return model_list_factory('dataset item', self._request('package_list'))

    def get_dataset(self, dataset_id):
        dataset = self._request('package_show', id=dataset_id)
        return model_dict_factory('dataset', dataset)

    def get_tags(self):
        tags = self._request('tag_list')
        return model_list_factory('tag item', tags)

    def get_tag(self, tag_id):
        tag = self._request('tag_show', id=tag_id)
        return model_dict_factory('tag', tag)
<|endoftext|>


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

requires the XenServer Python module:
http://community.citrix.com/cdn/xs/sdks/
"""

import sys
import socket
import XenAPI

class XenServerAPI:
    def __init__(self, host, user, passwd):
        self.session = XenAPI.Session('https://%s:443' % host)
        try:
            self.session.xenapi.login_with_password(user, passwd)
        except XenAPI.Failure, e:
            # If slave host from pool chosed, contact master
            exec 'err=%s' % str(e)
            self.session = XenAPI.Session('https://%s:443' % err[1])
            self.session.xenapi.login_with_password(user, passwd)
        except socket.error, e:
            # Try unqdn if connection refused
            self.session = XenAPI.Session('https://%s:443' % host.split('.')[0])
            self.session.xenapi.login_with_password(user, passwd)
        except Exception, e:
            print str(e)
            raise

    def choose_storage(self):
        srs = self.session.xenapi.SR.get_all_records()
        storage = {}
        for s in srs:
            key = srs[s]['name_label']
            if 'Local' in srs[s]['name_label']:
                block = self.session.xenapi.SR.get_PBDs(s)
                if block:
                    key += ':%s' % self.session.xenapi.host.get_name_label(
                        self.session.xenapi.PBD.get_host(block[0]))
                else:
                    continue
            if 'Removable' not in srs[s]['name_label'] and \
                'torage' in key and key not in storage:
                storage[key] = {
                    'size': long(srs[s]['physical_size']),
                    'used': long(srs[s]['physical_utilisation']),
                    'free': long(srs[s]['physical_size']) \
                          - long(srs[s]['physical_utilisation']),
                    }
        if len(storage.keys()) == 1:
            return storage.keys()[0].split(':')
        if len(storage.keys()) > 1:
            print 'There are multiple storage resources available:'
            skeys = sorted(storage.keys())
            svals = map(storage.get, skeys)
            storage = zip(skeys, svals)
            for s in storage:
                print '%d) %-35s (%5d GB free)' % (storage.index(s),
                    s[0], s[1]['free']/1024/1024/1024)
            ans = raw_input('Which one do you want to use: ')
            if int(ans) < 0 or int(ans) >= len(storage):
                print 'Storage selection aborted.'
                return False
            else:
                print 'User selected "%s"' % storage[int(ans)][0]
                if len(storage[int(ans)][0].split(':')) > 1:
                    return tuple(storage[int(ans)][0].split(':'))
                else:
                    return tuple([storage[int(ans)][0], None])
        else:
            return False

    def check_memory(self, need=None, xenhost=None):
        maxfree = 0
        hosts = self.session.xenapi.host.get_all()
        for h in hosts:
            free = long(self.session.xenapi.host.compute_free_memory(h))/1024/1024
            #print '%-20s: %5d MB Memory Free' % (self.session.xenapi.host.get_hostname(h), free)
            if (free > maxfree):
                maxfree = free
            if xenhost and xenhost == self.session.xenapi.host.get_hostname(h):
                return free
        if need:
            return maxfree>=need
        else:
            return maxfree

    def check_vcpus(self, need=None):
        hcpus = 0
        vcpus = 0
        vms = self.session.xenapi.VM.get_all_records()
        for v in vms:
            if vms[v]['power_state'] == 'Running':
                if 'Control domain' in vms[v]['name_label']:
                    hcpus += int(vms[v]['VCPUs_max'])
                else:
                    vcpus += int(vms[v]['VCPUs_max'])
        #print 'Pool/Host is currently using %d out of %d CPUs' % (int(vcpus), int(hcpus))
        if need:
            return int(hcpus)-int(vcpus)>=need
        else:
            return int(hcpus)-int(vcpus)

    def check_storage(self, disk=None, need=None):
        if not disk:
            return False
        maxfree = 0
        srs = self.session.xenapi.SR.get_all()
        for s in srs:
            name = self.session.xenapi.SR.get_name_label(s)
            if name == disk:
                free = (long(self.session.xenapi.SR.get_physical_size(s)) \
                    - long(self.session.xenapi.SR.get_physical_utilisation(s))  \
                    ) / 1024 / 1024 / 1024
                #print '%-20s: %5d GB Storage Free' % (name, free)
                if (free > maxfree):
                    maxfree = free
        if need:
            return maxfree>=need
        else:
            return maxfree

    def triple_check(self, info):
        info['storage'], xenhost = self.choose_storage()
        if xenhost:
            info['power_switch'] = xenhost
        freecores = self.check_vcpus()
        if freecores < info['cores']:
            print '%s does not have enough VCPUs for %d cores (%d available)!' \
                % (info['power_switch'], info['cores'], freecores)
            return False
        freeram = self.check_memory(xenhost)
        if freeram < (info['ram']*1024):
            print '%s does not have enough memory for %d MB (%d available)!' \
                % (info['power_switch'], info['ram']*1024, freeram)
            return False
        freedisk = self.check_storage(info['storage'])
        if freedisk < info['disk']:
            print '%s does not have enough storage for %d GB (%d available)!' \
                % (info['power_switch'], info['disk'], freedisk)
            return False
        return info['power_switch']
<commit_msg>typo, forgot named parameters for check_memory<commit_after># Copyright 2011 Gilt Groupe, INC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
mothership.xen
for interacting with the Citrix XenAPI

requires the XenServer Python module:
http://community.citrix.com/cdn/xs/sdks/
"""

import sys
import socket
import XenAPI

class XenServerAPI:
    def __init__(self, host, user, passwd):
        self.session = XenAPI.Session('https://%s:443' % host)
        try:
            self.session.xenapi.login_with_password(user, passwd)
        except XenAPI.Failure, e:
            # If slave host from pool chosed, contact master
            exec 'err=%s' % str(e)
            self.session = XenAPI.Session('https://%s:443' % err[1])
            self.session.xenapi.login_with_password(user, passwd)
        except socket.error, e:
            # Try unqdn if connection refused
            self.session = XenAPI.Session('https://%s:443' % host.split('.')[0])
            self.session.xenapi.login_with_password(user, passwd)
        except Exception, e:
            print str(e)
            raise

    def choose_storage(self):
        srs = self.session.xenapi.SR.get_all_records()
        storage = {}
        for s in srs:
            key = srs[s]['name_label']
            if 'Local' in srs[s]['name_label']:
                block = self.session.xenapi.SR.get_PBDs(s)
                if block:
                    key += ':%s' % self.session.xenapi.host.get_name_label(
                        self.session.xenapi.PBD.get_host(block[0]))
                else:
                    continue
            if 'Removable' not in srs[s]['name_label'] and \
                'torage' in key and key not in storage:
                storage[key] = {
                    'size': long(srs[s]['physical_size']),
                    'used': long(srs[s]['physical_utilisation']),
                    'free': long(srs[s]['physical_size']) \
                          - long(srs[s]['physical_utilisation']),
                    }
        if len(storage.keys()) == 1:
            return storage.keys()[0].split(':')
        if len(storage.keys()) > 1:
            print 'There are multiple storage resources available:'
            skeys = sorted(storage.keys())
            svals = map(storage.get, skeys)
            storage = zip(skeys, svals)
            for s in storage:
                print '%d) %-35s (%5d GB free)' % (storage.index(s),
                    s[0], s[1]['free']/1024/1024/1024)
            ans = raw_input('Which one do you want to use: ')
            if int(ans) < 0 or int(ans) >= len(storage):
                print 'Storage selection aborted.'
                return False
            else:
                print 'User selected "%s"' % storage[int(ans)][0]
                if len(storage[int(ans)][0].split(':')) > 1:
                    return tuple(storage[int(ans)][0].split(':'))
                else:
                    return tuple([storage[int(ans)][0], None])
        else:
            return False

    def check_memory(self, need=None, xenhost=None):
        maxfree = 0
        hosts = self.session.xenapi.host.get_all()
        for h in hosts:
            free = long(self.session.xenapi.host.compute_free_memory(h))/1024/1024
            print '%-20s: %5d MB Memory Free' % (self.session.xenapi.host.get_hostname(h), free)
            if (free > maxfree):
                maxfree = free
            if xenhost and xenhost == self.session.xenapi.host.get_hostname(h):
                return free
        if need:
            return maxfree>=need
        else:
            return maxfree

    def check_vcpus(self, need=None):
        hcpus = 0
        vcpus = 0
        vms = self.session.xenapi.VM.get_all_records()
        for v in vms:
            if vms[v]['power_state'] == 'Running':
                if 'Control domain' in vms[v]['name_label']:
                    hcpus += int(vms[v]['VCPUs_max'])
                else:
                    vcpus += int(vms[v]['VCPUs_max'])
        #print 'Pool/Host is currently using %d out of %d CPUs' % (int(vcpus), int(hcpus))
        if need:
            return int(hcpus)-int(vcpus)>=need
        else:
            return int(hcpus)-int(vcpus)

    def check_storage(self, disk=None, need=None):
        if not disk:
            return False
        maxfree = 0
        srs = self.session.xenapi.SR.get_all()
        for s in srs:
            name = self.session.xenapi.SR.get_name_label(s)
            if name == disk:
                free = (long(self.session.xenapi.SR.get_physical_size(s)) \
                    - long(self.session.xenapi.SR.get_physical_utilisation(s))  \
                    ) / 1024 / 1024 / 1024
                #print '%-20s: %5d GB Storage Free' % (name, free)
                if (free > maxfree):
                    maxfree = free
        if need:
            return maxfree>=need
        else:
            return maxfree

    def triple_check(self, info):
        info['storage'], xenhost = self.choose_storage()
        if xenhost:
            info['power_switch'] = xenhost
        freecores = self.check_vcpus()
        if freecores < info['cores']:
            print '%s does not have enough VCPUs for %d cores (%d available)!' \
                % (info['power_switch'], info['cores'], freecores)
            return False
        freeram = self.check_memory(xenhost=xenhost)
        if freeram < (info['ram']*1024):
            print '%s does not have enough memory for %d MB (%d available)!' \
                % (info['power_switch'], info['ram']*1024, freeram)
            return False
        freedisk = self.check_storage(info['storage'])
        if freedisk < info['disk']:
            print '%s does not have enough storage for %d GB (%d available)!' \
                % (info['power_switch'], info['disk'], freedisk)
            return False
        return info['power_switch']
<|endoftext|>


<commit_before>import ConfigParser
import boto.exception
import boto.s3.connection
import bunch
import itertools
import os
import random
import string

from .utils import region_sync_meta

s3 = bunch.Bunch()
config = bunch.Bunch()
targets = bunch.Bunch()

# this will be assigned by setup()
prefix = None

calling_formats = dict(
    ordinary=boto.s3.connection.OrdinaryCallingFormat(),
    subdomain=boto.s3.connection.SubdomainCallingFormat(),
    vhost=boto.s3.connection.VHostCallingFormat(),
    )

def get_prefix():
    assert prefix is not None
    return prefix

def is_slow_backend():
    return slow_backend

def choose_bucket_prefix(template, max_len=30):
    """
    Choose a prefix for our test buckets, so they're easy to identify.

    Use template and feed it more and more random filler, until it's
    as long as possible but still below max_len.
    """
    rand = ''.join(
        random.choice(string.ascii_lowercase + string.digits)
        for c in range(255)
        )

    while rand:
        s = template.format(random=rand)
        if len(s) <= max_len:
            return s
        rand = rand[:-1]

    raise RuntimeError(
        'Bucket prefix template is impossible to fulfill: {template!r}'.format(
            template=template,
            ),
        )

def nuke_prefixed_buckets_on_conn(prefix, name, conn):
    print 'Cleaning buckets from connection {name} prefix {prefix!r}.'.format(
        name=name,
        prefix=prefix,
        )
    for bucket in conn.get_all_buckets():
        if bucket.name.startswith(prefix):
            print 'Cleaning bucket {bucket}'.format(bucket=bucket)
            try:
                bucket.set_canned_acl('private')
                for key in bucket.list():
                    print 'Cleaning bucket {bucket} key {key}'.format(
                        bucket=bucket,
                        key=key,
                        )
                    key.set_canned_acl('private')
                    key.delete()
                bucket.delete()
            except boto.exception.S3ResponseError as e:
                if e.error_code != 'AccessDenied':
                    print 'GOT UNWANTED ERROR', e.error_code
                    raise
                # seems like we're not the owner of the bucket; ignore
                pass

def nuke_prefixed_buckets(prefix):
    # If no regions are specified, use the simple method
    if targets.main.master == None:
        for name, conn in s3.items():
            print 'Deleting buckets on {name}'.format(name=name)
            nuke_prefixed_buckets_on_conn(prefix, name, conn)
    else:
		    # First, delete all buckets on the master connection
		    for name, conn in s3.items():
		        if conn == targets.main.master.connection:
		            print 'Deleting buckets on {name} (master)'.format(name=name)
		            nuke_prefixed_buckets_on_conn(prefix, name, conn)

		    # Then sync to propagate deletes to secondaries
		    region_sync_meta(targets.main, targets.main.master.connection)
		    print 'region-sync in nuke_prefixed_buckets'

		    # Now delete remaining buckets on any other connection
		    for name, conn in s3.items():
		        if conn != targets.main.master.connection:
		            print 'Deleting buckets on {name} (non-master)'.format(name=name)
		            nuke_prefixed_buckets_on_conn(prefix, name, conn)

    print 'Done with cleanup of test buckets.'

class TargetConfig:
    def __init__(self, cfg, section):
        self.port = None
        self.api_name = ''
        self.is_master = False
        self.is_secure = False
        self.sync_agent_addr = None
        self.sync_agent_port = 0
        self.sync_meta_wait = 0
        try:
            self.api_name = cfg.get(section, 'api_name')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass
        try:
            self.port = cfg.getint(section, 'port')
        except ConfigParser.NoOptionError:
            pass
        try:
            self.host=cfg.get(section, 'host')
        except ConfigParser.NoOptionError:
            raise RuntimeError(
                'host not specified for section {s}'.format(s=section)
                )
        try:
            self.is_master=cfg.getboolean(section, 'is_master')
        except ConfigParser.NoOptionError:
            pass

        try:
            self.is_secure=cfg.getboolean(section, 'is_secure')
        except ConfigParser.NoOptionError:
            pass

        try:
            raw_calling_format = cfg.get(section, 'calling_format')
        except ConfigParser.NoOptionError:
            raw_calling_format = 'ordinary'

        try:
            self.sync_agent_addr = cfg.get(section, 'sync_agent_addr')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass

        try:
            self.sync_agent_port = cfg.getint(section, 'sync_agent_port')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass

        try:
            self.sync_meta_wait = cfg.getint(section, 'sync_meta_wait')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass

        try:
            self.calling_format = calling_formats[raw_calling_format]
        except KeyError:
            raise RuntimeError(
                'calling_format unknown: %r' % raw_calling_format
                )

class TargetConnection:
    def __init__(self, conf, conn):
        self.conf = conf
        self.connection = conn

class RegionsInfo:
    def __init__(self):
        self.m = bunch.Bunch()
        self.master = None
        self.secondaries = []

    def add(self, name, region_config):
        self.m[name] = region_config
        if (region_config.is_master):
            if not self.master is None:
                raise RuntimeError(
                    'multiple regions defined as master'
                    )
            self.master = region_config
        else:
            self.secondaries.append(region_config)
    def get(self, name):
        return self.m[name]
    def get(self):
        return self.m
    def iteritems(self):
        return self.m.iteritems()

regions = RegionsInfo()

class RegionsConn:
    def __init__(self):
        self.m = bunch.Bunch()
        self.default = None
        self.master = None
        self.secondaries = []

    def iteritems(self):
        return self.m.iteritems()

    def set_default(self, conn):
        self.default = conn

    def add(self, name, conn):
        self.m[name] = conn
        if not self.default:
            self.default = conn
        if (conn.conf.is_master):
            self.master = conn
        else:
            self.secondaries.append(conn)

# nosetests --processes=N with N>1 is safe
_multiprocess_can_split_ = True

def setup():

    cfg = ConfigParser.RawConfigParser()
    try:
        path = os.environ['S3TEST_CONF']
    except KeyError:
        raise RuntimeError(
            'To run tests, point environment '
            + 'variable S3TEST_CONF to a config file.',
            )
    with file(path) as f:
        cfg.readfp(f)

    global prefix
    global targets
    global slow_backend

    try:
        template = cfg.get('fixtures', 'bucket prefix')
    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
        template = 'test-{random}-'
    prefix = choose_bucket_prefix(template=template)

    try:
        slow_backend = cfg.getboolean('fixtures', 'slow backend')
    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
        slow_backend = False

    # pull the default_region out, if it exists
    try:
        default_region = cfg.get('fixtures', 'default_region')
    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
        default_region = None

    s3.clear()
    config.clear()

    for section in cfg.sections():
        try:
            (type_, name) = section.split(None, 1)
        except ValueError:
            continue
        if type_ != 'region':
            continue
        regions.add(name, TargetConfig(cfg, section))

    for section in cfg.sections():
        try:
            (type_, name) = section.split(None, 1)
        except ValueError:
            continue
        if type_ != 's3':
            continue

        if len(regions.get()) == 0:
            regions.add("default", TargetConfig(cfg, section))

        config[name] = bunch.Bunch()
        for var in [
            'user_id',
            'display_name',
            'email',
            ]:
            try:
                config[name][var] = cfg.get(section, var)
            except ConfigParser.NoOptionError:
                pass

        targets[name] = RegionsConn()

        for (k, conf) in regions.iteritems():
            conn = boto.s3.connection.S3Connection(
                aws_access_key_id=cfg.get(section, 'access_key'),
                aws_secret_access_key=cfg.get(section, 'secret_key'),
                is_secure=conf.is_secure,
                port=conf.port,
                host=conf.host,
                # TODO test vhost calling format
                calling_format=conf.calling_format,
                )

            temp_targetConn = TargetConnection(conf, conn)
            targets[name].add(k, temp_targetConn)

            # Explicitly test for and set the default region, if specified.
            # If it was not specified, use the 'is_master' flag to set it.
            if default_region:
                if default_region == name:
                    targets[name].set_default(temp_targetConn)
            elif conf.is_master:
                targets[name].set_default(temp_targetConn)

        s3[name] = targets[name].default.connection

    # WARNING! we actively delete all buckets we see with the prefix
    # we've chosen! Choose your prefix with care, and don't reuse
    # credentials!

    # We also assume nobody else is going to use buckets with that
    # prefix. This is racy but given enough randomness, should not
    # really fail.
    nuke_prefixed_buckets(prefix=prefix)

def teardown():
    # remove our buckets here also, to avoid littering
    nuke_prefixed_buckets(prefix=prefix)

bucket_counter = itertools.count(1)

def get_new_bucket_name():
    """
    Get a bucket name that probably does not exist.

    We make every attempt to use a unique random prefix, so if a
    bucket by this name happens to exist, it's ok if tests give
    false negatives.
    """
    name = '{prefix}{num}'.format(
        prefix=prefix,
        num=next(bucket_counter),
        )
    return name

def get_new_bucket(target=None, name=None, headers=None):
    """
    Get a bucket that exists and is empty.

    Always recreates a bucket from scratch. This is useful to also
    reset ACLs and such.
    """
    if target is None:
        target = targets.main.default
    connection = target.connection
    if name is None:
        name = get_new_bucket_name()
    # the only way for this to fail with a pre-existing bucket is if
    # someone raced us between setup nuke_prefixed_buckets and here;
    # ignore that as astronomically unlikely
    bucket = connection.create_bucket(name, location=target.conf.api_name, headers=headers)
    return bucket
<commit_msg>teardown: clean object versions<commit_after>import ConfigParser
import boto.exception
import boto.s3.connection
import bunch
import itertools
import os
import random
import string

from .utils import region_sync_meta

s3 = bunch.Bunch()
config = bunch.Bunch()
targets = bunch.Bunch()

# this will be assigned by setup()
prefix = None

calling_formats = dict(
    ordinary=boto.s3.connection.OrdinaryCallingFormat(),
    subdomain=boto.s3.connection.SubdomainCallingFormat(),
    vhost=boto.s3.connection.VHostCallingFormat(),
    )

def get_prefix():
    assert prefix is not None
    return prefix

def is_slow_backend():
    return slow_backend

def choose_bucket_prefix(template, max_len=30):
    """
    Choose a prefix for our test buckets, so they're easy to identify.

    Use template and feed it more and more random filler, until it's
    as long as possible but still below max_len.
    """
    rand = ''.join(
        random.choice(string.ascii_lowercase + string.digits)
        for c in range(255)
        )

    while rand:
        s = template.format(random=rand)
        if len(s) <= max_len:
            return s
        rand = rand[:-1]

    raise RuntimeError(
        'Bucket prefix template is impossible to fulfill: {template!r}'.format(
            template=template,
            ),
        )

def nuke_prefixed_buckets_on_conn(prefix, name, conn):
    print 'Cleaning buckets from connection {name} prefix {prefix!r}.'.format(
        name=name,
        prefix=prefix,
        )

    for bucket in conn.get_all_buckets():
        print 'prefix=',prefix
        if bucket.name.startswith(prefix):
            print 'Cleaning bucket {bucket}'.format(bucket=bucket)
            try:
                # bucket.set_canned_acl('private')
                for key in bucket.list_versions():
                    print 'Cleaning bucket {bucket} key {key}'.format(
                        bucket=bucket,
                        key=key,
                        )
                    # key.set_canned_acl('private')
                    key.delete()
                bucket.delete()
            except boto.exception.S3ResponseError as e:
                if e.error_code != 'AccessDenied':
                    print 'GOT UNWANTED ERROR', e.error_code
                    raise
                # seems like we're not the owner of the bucket; ignore
                pass

def nuke_prefixed_buckets(prefix):
    # If no regions are specified, use the simple method
    if targets.main.master == None:
        for name, conn in s3.items():
            print 'Deleting buckets on {name}'.format(name=name)
            nuke_prefixed_buckets_on_conn(prefix, name, conn)
    else:
		    # First, delete all buckets on the master connection
		    for name, conn in s3.items():
		        if conn == targets.main.master.connection:
		            print 'Deleting buckets on {name} (master)'.format(name=name)
		            nuke_prefixed_buckets_on_conn(prefix, name, conn)

		    # Then sync to propagate deletes to secondaries
		    region_sync_meta(targets.main, targets.main.master.connection)
		    print 'region-sync in nuke_prefixed_buckets'

		    # Now delete remaining buckets on any other connection
		    for name, conn in s3.items():
		        if conn != targets.main.master.connection:
		            print 'Deleting buckets on {name} (non-master)'.format(name=name)
		            nuke_prefixed_buckets_on_conn(prefix, name, conn)

    print 'Done with cleanup of test buckets.'

class TargetConfig:
    def __init__(self, cfg, section):
        self.port = None
        self.api_name = ''
        self.is_master = False
        self.is_secure = False
        self.sync_agent_addr = None
        self.sync_agent_port = 0
        self.sync_meta_wait = 0
        try:
            self.api_name = cfg.get(section, 'api_name')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass
        try:
            self.port = cfg.getint(section, 'port')
        except ConfigParser.NoOptionError:
            pass
        try:
            self.host=cfg.get(section, 'host')
        except ConfigParser.NoOptionError:
            raise RuntimeError(
                'host not specified for section {s}'.format(s=section)
                )
        try:
            self.is_master=cfg.getboolean(section, 'is_master')
        except ConfigParser.NoOptionError:
            pass

        try:
            self.is_secure=cfg.getboolean(section, 'is_secure')
        except ConfigParser.NoOptionError:
            pass

        try:
            raw_calling_format = cfg.get(section, 'calling_format')
        except ConfigParser.NoOptionError:
            raw_calling_format = 'ordinary'

        try:
            self.sync_agent_addr = cfg.get(section, 'sync_agent_addr')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass

        try:
            self.sync_agent_port = cfg.getint(section, 'sync_agent_port')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass

        try:
            self.sync_meta_wait = cfg.getint(section, 'sync_meta_wait')
        except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
            pass

        try:
            self.calling_format = calling_formats[raw_calling_format]
        except KeyError:
            raise RuntimeError(
                'calling_format unknown: %r' % raw_calling_format
                )

class TargetConnection:
    def __init__(self, conf, conn):
        self.conf = conf
        self.connection = conn

class RegionsInfo:
    def __init__(self):
        self.m = bunch.Bunch()
        self.master = None
        self.secondaries = []

    def add(self, name, region_config):
        self.m[name] = region_config
        if (region_config.is_master):
            if not self.master is None:
                raise RuntimeError(
                    'multiple regions defined as master'
                    )
            self.master = region_config
        else:
            self.secondaries.append(region_config)
    def get(self, name):
        return self.m[name]
    def get(self):
        return self.m
    def iteritems(self):
        return self.m.iteritems()

regions = RegionsInfo()

class RegionsConn:
    def __init__(self):
        self.m = bunch.Bunch()
        self.default = None
        self.master = None
        self.secondaries = []

    def iteritems(self):
        return self.m.iteritems()

    def set_default(self, conn):
        self.default = conn

    def add(self, name, conn):
        self.m[name] = conn
        if not self.default:
            self.default = conn
        if (conn.conf.is_master):
            self.master = conn
        else:
            self.secondaries.append(conn)

# nosetests --processes=N with N>1 is safe
_multiprocess_can_split_ = True

def setup():

    cfg = ConfigParser.RawConfigParser()
    try:
        path = os.environ['S3TEST_CONF']
    except KeyError:
        raise RuntimeError(
            'To run tests, point environment '
            + 'variable S3TEST_CONF to a config file.',
            )
    with file(path) as f:
        cfg.readfp(f)

    global prefix
    global targets
    global slow_backend

    try:
        template = cfg.get('fixtures', 'bucket prefix')
    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
        template = 'test-{random}-'
    prefix = choose_bucket_prefix(template=template)

    try:
        slow_backend = cfg.getboolean('fixtures', 'slow backend')
    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
        slow_backend = False

    # pull the default_region out, if it exists
    try:
        default_region = cfg.get('fixtures', 'default_region')
    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError):
        default_region = None

    s3.clear()
    config.clear()

    for section in cfg.sections():
        try:
            (type_, name) = section.split(None, 1)
        except ValueError:
            continue
        if type_ != 'region':
            continue
        regions.add(name, TargetConfig(cfg, section))

    for section in cfg.sections():
        try:
            (type_, name) = section.split(None, 1)
        except ValueError:
            continue
        if type_ != 's3':
            continue

        if len(regions.get()) == 0:
            regions.add("default", TargetConfig(cfg, section))

        config[name] = bunch.Bunch()
        for var in [
            'user_id',
            'display_name',
            'email',
            ]:
            try:
                config[name][var] = cfg.get(section, var)
            except ConfigParser.NoOptionError:
                pass

        targets[name] = RegionsConn()

        for (k, conf) in regions.iteritems():
            conn = boto.s3.connection.S3Connection(
                aws_access_key_id=cfg.get(section, 'access_key'),
                aws_secret_access_key=cfg.get(section, 'secret_key'),
                is_secure=conf.is_secure,
                port=conf.port,
                host=conf.host,
                # TODO test vhost calling format
                calling_format=conf.calling_format,
                )

            temp_targetConn = TargetConnection(conf, conn)
            targets[name].add(k, temp_targetConn)

            # Explicitly test for and set the default region, if specified.
            # If it was not specified, use the 'is_master' flag to set it.
            if default_region:
                if default_region == name:
                    targets[name].set_default(temp_targetConn)
            elif conf.is_master:
                targets[name].set_default(temp_targetConn)

        s3[name] = targets[name].default.connection

    # WARNING! we actively delete all buckets we see with the prefix
    # we've chosen! Choose your prefix with care, and don't reuse
    # credentials!

    # We also assume nobody else is going to use buckets with that
    # prefix. This is racy but given enough randomness, should not
    # really fail.
    nuke_prefixed_buckets(prefix=prefix)

def teardown():
    # remove our buckets here also, to avoid littering
    nuke_prefixed_buckets(prefix=prefix)

bucket_counter = itertools.count(1)

def get_new_bucket_name():
    """
    Get a bucket name that probably does not exist.

    We make every attempt to use a unique random prefix, so if a
    bucket by this name happens to exist, it's ok if tests give
    false negatives.
    """
    name = '{prefix}{num}'.format(
        prefix=prefix,
        num=next(bucket_counter),
        )
    return name

def get_new_bucket(target=None, name=None, headers=None):
    """
    Get a bucket that exists and is empty.

    Always recreates a bucket from scratch. This is useful to also
    reset ACLs and such.
    """
    if target is None:
        target = targets.main.default
    connection = target.connection
    if name is None:
        name = get_new_bucket_name()
    # the only way for this to fail with a pre-existing bucket is if
    # someone raced us between setup nuke_prefixed_buckets and here;
    # ignore that as astronomically unlikely
    bucket = connection.create_bucket(name, location=target.conf.api_name, headers=headers)
    return bucket
<|endoftext|>


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

###
# Copyright (c) 2002, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

from fix import *

import re
import string
import fnmatch

import world

def isUserHostmask(s):
    p1 = s.find('!')
    p2 = s.find('@')
    if p1 < p2-1 and p1 >= 1 and p2 >= 3 and len(s) > p2+1:
        return True
    else:
        return False

def isServerHostmask(s):
    return (not isUserHostmask(s) and s.find('!') == -1 and s.find('@') == -1)

def nickFromHostmask(hostmask):
    assert isUserHostmask(hostmask)
    return nick(hostmask.split('!', 1)[0])

def userFromHostmask(hostmask):
    assert isUserHostmask(hostmask)
    return hostmask.split('!', 1)[1].split('@', 1)[0]

def hostFromHostmask(hostmask):
    assert isUserHostmask(hostmask)
    return hostmask.split('@', 1)[1]

def splitHostmask(hostmask):
    assert isUserHostmask(hostmask)
    nck, rest = hostmask.split('!', 1)
    user, host = rest.split('@', 1)
    return (nick(nck), user, host)

def joinHostmask(nick, ident, host):
    assert nick and ident and host
    return '%s!%s@%s' % (nick, ident, host)

_lowertrans = string.maketrans(string.ascii_uppercase + r'\[]~',
                               string.ascii_lowercase + r'|{}^')
def toLower(nick):
    return nick.translate(_lowertrans)

def nickEqual(nick1, nick2):
    return toLower(nick1) == toLower(nick2)

nickchars = string.ascii_lowercase + string.ascii_uppercase + r'-[]\\`^{}'
_nickre = re.compile(r'^[%s]+$' % re.escape(nickchars))
def isNick(s):
    if re.match(_nickre, s):
        return True
    else:
        return False

def isChannel(s):
    return (s and s[0] in '#&+!' and len(s) <= 50 and \
            '\x07' not in s and ',' not in s and ' ' not in s)

def hostmaskPatternEqual(pattern, hostmask):
    return fnmatch.fnmatch(toLower(hostmask), toLower(pattern))

_ipchars = string.digits + '.'
def isIP(s):
    """Not quite perfect, but close enough until I can find the regexp I want.

    >>> isIP('255.255.255.255')
    1

    >>> isIP('abc.abc.abc.abc')
    0
    """
    return (s.translate(string.ascii, _ipchars) == "")

def banmask(hostmask):
    """Returns a properly generic banning hostmask for a hostmask.

    >>> banmask('nick!user@host.domain.tld')
    '*!*@*.domain.tld'

    >>> banmask('nick!user@10.0.0.1')
    '*!*@10.0.0.*'
    """
    host = hostFromHostmask(hostmask)
    if isIP(host):
        return ('*!*@%s.*' % host[:host.rfind('.')])
    else:
        return ('*!*@*%s' % host[host.find('.'):])

_argModes = 'ovhblkqe'
def separateModes(args):
    """Separates modelines into single mode change tuples.

    Examples:

    >>> separateModes(['+ooo', 'jemfinch', 'StoneTable', 'philmes'])
    [('+o', 'jemfinch'), ('+o', 'StoneTable'), ('+o', 'philmes')]

    >>> separateModes(['+o-o', 'jemfinch', 'PeterB'])
    [('+o', 'jemfinch'), ('-o', 'PeterB')]

    >>> separateModes(['+s-o', 'test'])
    [('+s', None), ('-o', 'test')]

    >>> separateModes(['+sntl', '100'])
    [('+s', None), ('+n', None), ('+t', None), ('+l', '100')]
    """
    modes = args[0]
    args = list(args[1:])
    ret = []
    index = 0
    length = len(modes)
    while index < length:
        if modes[index] in '+-':
            last = modes[index]
            index += 1
        else:
            if modes[index] in _argModes:
                ret.append((last + modes[index], args.pop(0)))
            else:
                ret.append((last + modes[index], None))
            index += 1
    return ret

def bold(s):
    return "\x02%s\x02" % s

def isValidArgument(s):
    return '\r' not in s and '\n' not in s and '\x00' not in s

notFunky = string.printable+'\x02'
def safeArgument(s):
    if isValidArgument(s) and s.translate(string.ascii, notFunky) == '':
        return s
    else:
        return repr(s)

def replyTo(msg):
    if isChannel(msg.args[0]):
        return msg.args[0]
    else:
        return msg.nick

class nick(str):
    """This class does case-insensitive comparisons of nicks."""
    def __init__(self, s):
        self.lowered = toLower(s)

    def __eq__(self, s):
        try:
            return toLower(s) == self.lowered
        except:
            return False

    def __hash__(self):
        return hash(self.lowered)

if __name__ == '__main__':
    import sys, doctest
    doctest.testmod(sys.modules['__main__'])
# vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78:
<commit_msg>Added joinModes function; docstring to bold; assert to separateModes.<commit_after>#!/usr/bin/env python

###
# Copyright (c) 2002, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

from fix import *

import re
import string
import fnmatch

import debug
import world

def isUserHostmask(s):
    p1 = s.find('!')
    p2 = s.find('@')
    if p1 < p2-1 and p1 >= 1 and p2 >= 3 and len(s) > p2+1:
        return True
    else:
        return False

def isServerHostmask(s):
    return (not isUserHostmask(s) and s.find('!') == -1 and s.find('@') == -1)

def nickFromHostmask(hostmask):
    assert isUserHostmask(hostmask)
    return nick(hostmask.split('!', 1)[0])

def userFromHostmask(hostmask):
    assert isUserHostmask(hostmask)
    return hostmask.split('!', 1)[1].split('@', 1)[0]

def hostFromHostmask(hostmask):
    assert isUserHostmask(hostmask)
    return hostmask.split('@', 1)[1]

def splitHostmask(hostmask):
    assert isUserHostmask(hostmask)
    nck, rest = hostmask.split('!', 1)
    user, host = rest.split('@', 1)
    return (nick(nck), user, host)

def joinHostmask(nick, ident, host):
    assert nick and ident and host
    return '%s!%s@%s' % (nick, ident, host)

_lowertrans = string.maketrans(string.ascii_uppercase + r'\[]~',
                               string.ascii_lowercase + r'|{}^')
def toLower(nick):
    return nick.translate(_lowertrans)

def nickEqual(nick1, nick2):
    return toLower(nick1) == toLower(nick2)

nickchars = string.ascii_lowercase + string.ascii_uppercase + r'-[]\\`^{}'
_nickre = re.compile(r'^[%s]+$' % re.escape(nickchars))
def isNick(s):
    if re.match(_nickre, s):
        return True
    else:
        return False

def isChannel(s):
    return (s and s[0] in '#&+!' and len(s) <= 50 and \
            '\x07' not in s and ',' not in s and ' ' not in s)

def hostmaskPatternEqual(pattern, hostmask):
    return fnmatch.fnmatch(toLower(hostmask), toLower(pattern))

_ipchars = string.digits + '.'
def isIP(s):
    """Not quite perfect, but close enough until I can find the regexp I want.

    >>> isIP('255.255.255.255')
    1

    >>> isIP('abc.abc.abc.abc')
    0
    """
    return (s.translate(string.ascii, _ipchars) == "")

def banmask(hostmask):
    """Returns a properly generic banning hostmask for a hostmask.

    >>> banmask('nick!user@host.domain.tld')
    '*!*@*.domain.tld'

    >>> banmask('nick!user@10.0.0.1')
    '*!*@10.0.0.*'
    """
    host = hostFromHostmask(hostmask)
    if isIP(host):
        return ('*!*@%s.*' % host[:host.rfind('.')])
    else:
        return ('*!*@*%s' % host[host.find('.'):])

_argModes = 'ovhblkqe'
def separateModes(args):
    """Separates modelines into single mode change tuples.

    Examples:

    >>> separateModes(['+ooo', 'jemfinch', 'StoneTable', 'philmes'])
    [('+o', 'jemfinch'), ('+o', 'StoneTable'), ('+o', 'philmes')]

    >>> separateModes(['+o-o', 'jemfinch', 'PeterB'])
    [('+o', 'jemfinch'), ('-o', 'PeterB')]

    >>> separateModes(['+s-o', 'test'])
    [('+s', None), ('-o', 'test')]

    >>> separateModes(['+sntl', '100'])
    [('+s', None), ('+n', None), ('+t', None), ('+l', '100')]
    """
    modes = args[0]
    assert modes[0] in '+-', 'Invalid args: %r' % args
    args = list(args[1:])
    ret = []
    index = 0
    length = len(modes)
    while index < length:
        if modes[index] in '+-':
            last = modes[index]
            index += 1
        else:
            if modes[index] in _argModes:
                ret.append((last + modes[index], args.pop(0)))
            else:
                ret.append((last + modes[index], None))
            index += 1
    return ret

def joinModes(modes):
    """Joins modes of the same form as returned by separateModes."""
    args = []
    modeChars = []
    currentMode = '\x00'
    for (mode, arg) in modes:
        if arg is not None:
            args.append(arg)
        if not mode.startswith(currentMode):
            currentMode = mode[0]
            modeChars.append(mode[0])
        modeChars.append(mode[1])
    args.insert(0, ''.join(modeChars))
    return args

def bold(s):
    """Returns the string s, bolded."""
    return "\x02%s\x02" % s

def isValidArgument(s):
    return '\r' not in s and '\n' not in s and '\x00' not in s

notFunky = string.printable+'\x02'
def safeArgument(s):
    if isValidArgument(s) and s.translate(string.ascii, notFunky) == '':
        return s
    else:
        return repr(s)

def replyTo(msg):
    if isChannel(msg.args[0]):
        return msg.args[0]
    else:
        return msg.nick

class nick(str):
    """This class does case-insensitive comparisons of nicks."""
    def __init__(self, s):
        self.lowered = toLower(s)

    def __eq__(self, s):
        try:
            return toLower(s) == self.lowered
        except:
            return False

    def __hash__(self):
        return hash(self.lowered)

if __name__ == '__main__':
    import sys, doctest
    doctest.testmod(sys.modules['__main__'])
# vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78:
<|endoftext|>