<commit_before>'''
Created on 2017/8/14

:author: hubo
'''

import asyncio
import functools
import queue

def wrap_callback(callback, loop):
    @functools.wraps(callback)
    def _callback(*args, **kwargs):
        loop.call_soon_threadsafe(functools.partial(callback, *args, **kwargs))
    return _callback

def wrap_in_executor(func, loop, executor=None):
    @functools.wraps(func)
    async def _func(*args, **kwargs):
        return await loop.run_in_executor(executor, functools.partial(func, *args, **kwargs))
    return _func

def wrap_future(grpc_fut, loop):
    fut_ = loop.create_future()
    def _set_state(grpc_fut, fut_):
        assert grpc_fut.done()
        if fut_.cancelled():
            return
        assert not fut_.done()
        if grpc_fut.cancelled():
            fut_.cancel()
        else:
            exception = grpc_fut.exception()
            if exception is not None:
                fut_.set_exception(exception)
            else:
                result = grpc_fut.result()
                fut_.set_result(result)

    def _call_check_cancel(fut_):
        if fut_.cancelled():
            grpc_fut.cancel()

    def _call_set_state(grpc_fut):
        loop.call_soon_threadsafe(_set_state, grpc_fut, fut_)

    fut_.add_done_callback(_call_check_cancel)
    grpc_fut.add_done_callback(_call_set_state)
    return fut_

def copy_members(source, dest, member_list, wrapper=None):
    for m in member_list:
        f = getattr(source, m, None)
        if f is None:
            continue
        if wrapper is not None:
            f = wrapper(f)
        setattr(dest, m, f)

def wrap_future_call(grpc_fut, loop, executor=None):
    fut_ = wrap_future(grpc_fut, loop)
    # Copy extra members
    copy_members(grpc_fut, fut_,
                 ['is_active',
                  'time_remaining'])
    @functools.wraps(grpc_fut.add_callback)
    def _add_callback(callback):
        grpc_fut.add_callback(wrap_callback(callback, loop))
    fut_.add_callback = _add_callback
    copy_members(grpc_fut, fut_,
                 ['initial_metadata',
                  'trailing_metadata',
                  'code',
                  'details'],
                 functools.partial(wrap_in_executor, loop=loop, executor=executor))
    return fut_

class WrappedIterator(object):
    """
    Wrap an grpc_iterator to an async iterator
    """
    def __init__(self, grpc_iterator, loop, executor=None, stream_executor=None):
        self._iterator = grpc_iterator
        self._loop = loop
        self._executor = executor
        if stream_executor is None:
            self._shared_executor = True
            stream_executor = executor
        else:
            self._shared_executor = False
            self._stream_executor = stream_executor
        self._next_future = None
        copy_members(grpc_iterator, self,
                     ['is_active',
                      'time_remaining',
                      'cancel'])
        @functools.wraps(grpc_iterator.add_callback)
        def _add_callback(callback):
            grpc_iterator.add_callback(wrap_callback(callback, loop))
        self.add_callback = _add_callback
        copy_members(grpc_iterator, self,
                     ['initial_metadata',
                      'trailing_metadata',
                      'code',
                      'details'],
                     functools.partial(wrap_in_executor, loop=loop, executor=executor))

    async def __aiter__(self):
        return self

    def _next(self):
        try:
            return next(self._iterator)
        except StopIteration:
            raise StopAsyncIteration

    async def __anext__(self):
        if self._next_future is None:
            self._next_future = self._loop.run_in_executor(self._stream_executor, self._next)
        try:
            return await asyncio.shield(self._next_future)
        finally:
            if self._next_future and self._next_future.done():
                self._next_future = None

    def __del__(self):
        if self._iterator is not None:
            self._iterator.cancel()
            self._iterator.close()
            self._iterator = None
        if self._next_future is not None:
            self._next_future.cancel()
            self._next_future = None
        if not self._shared_executor:
            self._stream_executor.shutdown()

    async def aclose(self):
        self.__del__()

class IteratorScope(object):
    def __init__(self, _iter):
        self._iter = _iter

    async def __aenter__(self):
        return self._iter

    async def __aexit__(self, exc_val, exc_typ, exc_tb):
        await self._iter.aclose()

class WrappedAsyncIterator(object):
    """
    Wrap an async iterator to an iterator for grpc input
    """
    def __init__(self, async_iter, loop):
        self._async_iter = async_iter
        self._loop = loop
        self._q = queue.Queue()

    def __iter__(self):
        return self

    async def _next(self):
        try:
            self._q.put((await self._async_iter.__anext__(), False))
        except StopAsyncIteration:
            self._q.put(None, True)
        except Exception as exc:
            self._q.put(exc, True)

    def __next__(self):
        if self._async_iter is None:
            raise StopIteration
        try:
            r, is_exc = self._q.get_nowait()
        except queue.Empty:
            self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, self._next(), loop=self._loop))
            r, is_exc = self._q.get()
        if is_exc:
            if r is None:
                self._async_iter = None
                raise StopIteration
            else:
                raise r
        else:
            return r

    def close(self):
        self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, self._async_iter.aclose(), loop=self._loop))
        self._async_iter = None

<commit_msg>do not call close()<commit_after>'''
Created on 2017/8/14

:author: hubo
'''

import asyncio
import functools
import queue

def wrap_callback(callback, loop):
    @functools.wraps(callback)
    def _callback(*args, **kwargs):
        loop.call_soon_threadsafe(functools.partial(callback, *args, **kwargs))
    return _callback

def wrap_in_executor(func, loop, executor=None):
    @functools.wraps(func)
    async def _func(*args, **kwargs):
        return await loop.run_in_executor(executor, functools.partial(func, *args, **kwargs))
    return _func

def wrap_future(grpc_fut, loop):
    fut_ = loop.create_future()
    def _set_state(grpc_fut, fut_):
        assert grpc_fut.done()
        if fut_.cancelled():
            return
        assert not fut_.done()
        if grpc_fut.cancelled():
            fut_.cancel()
        else:
            exception = grpc_fut.exception()
            if exception is not None:
                fut_.set_exception(exception)
            else:
                result = grpc_fut.result()
                fut_.set_result(result)

    def _call_check_cancel(fut_):
        if fut_.cancelled():
            grpc_fut.cancel()

    def _call_set_state(grpc_fut):
        loop.call_soon_threadsafe(_set_state, grpc_fut, fut_)

    fut_.add_done_callback(_call_check_cancel)
    grpc_fut.add_done_callback(_call_set_state)
    return fut_

def copy_members(source, dest, member_list, wrapper=None):
    for m in member_list:
        f = getattr(source, m, None)
        if f is None:
            continue
        if wrapper is not None:
            f = wrapper(f)
        setattr(dest, m, f)

def wrap_future_call(grpc_fut, loop, executor=None):
    fut_ = wrap_future(grpc_fut, loop)
    # Copy extra members
    copy_members(grpc_fut, fut_,
                 ['is_active',
                  'time_remaining'])
    @functools.wraps(grpc_fut.add_callback)
    def _add_callback(callback):
        grpc_fut.add_callback(wrap_callback(callback, loop))
    fut_.add_callback = _add_callback
    copy_members(grpc_fut, fut_,
                 ['initial_metadata',
                  'trailing_metadata',
                  'code',
                  'details'],
                 functools.partial(wrap_in_executor, loop=loop, executor=executor))
    return fut_

class WrappedIterator(object):
    """
    Wrap an grpc_iterator to an async iterator
    """
    def __init__(self, grpc_iterator, loop, executor=None, stream_executor=None):
        self._iterator = grpc_iterator
        self._loop = loop
        self._executor = executor
        if stream_executor is None:
            self._shared_executor = True
            stream_executor = executor
        else:
            self._shared_executor = False
            self._stream_executor = stream_executor
        self._next_future = None
        copy_members(grpc_iterator, self,
                     ['is_active',
                      'time_remaining',
                      'cancel'])
        @functools.wraps(grpc_iterator.add_callback)
        def _add_callback(callback):
            grpc_iterator.add_callback(wrap_callback(callback, loop))
        self.add_callback = _add_callback
        copy_members(grpc_iterator, self,
                     ['initial_metadata',
                      'trailing_metadata',
                      'code',
                      'details'],
                     functools.partial(wrap_in_executor, loop=loop, executor=executor))

    async def __aiter__(self):
        return self

    def _next(self):
        try:
            return next(self._iterator)
        except StopIteration:
            raise StopAsyncIteration

    async def __anext__(self):
        if self._next_future is None:
            self._next_future = self._loop.run_in_executor(self._stream_executor, self._next)
        try:
            return await asyncio.shield(self._next_future)
        finally:
            if self._next_future and self._next_future.done():
                self._next_future = None

    def __del__(self):
        if self._iterator is not None:
            self._iterator.cancel()
            self._iterator = None
        if self._next_future is not None:
            self._next_future.cancel()
            self._next_future = None
        if not self._shared_executor:
            self._stream_executor.shutdown()

    async def aclose(self):
        self.__del__()

class IteratorScope(object):
    def __init__(self, _iter):
        self._iter = _iter

    async def __aenter__(self):
        return self._iter

    async def __aexit__(self, exc_val, exc_typ, exc_tb):
        await self._iter.aclose()

class WrappedAsyncIterator(object):
    """
    Wrap an async iterator to an iterator for grpc input
    """
    def __init__(self, async_iter, loop):
        self._async_iter = async_iter
        self._loop = loop
        self._q = queue.Queue()

    def __iter__(self):
        return self

    async def _next(self):
        try:
            self._q.put((await self._async_iter.__anext__(), False))
        except StopAsyncIteration:
            self._q.put(None, True)
        except Exception as exc:
            self._q.put(exc, True)

    def __next__(self):
        if self._async_iter is None:
            raise StopIteration
        try:
            r, is_exc = self._q.get_nowait()
        except queue.Empty:
            self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, self._next(), loop=self._loop))
            r, is_exc = self._q.get()
        if is_exc:
            if r is None:
                self._async_iter = None
                raise StopIteration
            else:
                raise r
        else:
            return r

    def close(self):
        self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, self._async_iter.aclose(), loop=self._loop))
        self._async_iter = None

<|endoftext|>


<commit_before>'''
Created on 2017/8/14

:author: hubo
'''

import asyncio
import functools
import queue
from asyncio import CancelledError

def wrap_callback(callback, loop):
    @functools.wraps(callback)
    def _callback(*args, **kwargs):
        if not loop.is_closed():
            loop.call_soon_threadsafe(functools.partial(callback, *args, **kwargs))
    return _callback

def wrap_active_test(func, test, loop, executor=None):
    @functools.wraps(func)
    async def _func(*args, **kwargs):
        if test():
            return await loop.run_in_executor(executor, functools.partial(func, *args, **kwargs))
        else:
            return func(*args, **kwargs)
    return _func

def wrap_future(grpc_fut, loop):
    fut_ = loop.create_future()
    def _set_state(grpc_fut, fut_):
        assert grpc_fut.done()
        if fut_.cancelled():
            return
        assert not fut_.done()
        if grpc_fut.cancelled():
            fut_.cancel()
        else:
            exception = grpc_fut.exception()
            if exception is not None:
                fut_.set_exception(exception)
            else:
                result = grpc_fut.result()
                fut_.set_result(result)

    def _call_check_cancel(fut_):
        if fut_.cancelled():
            grpc_fut.cancel()

    def _call_set_state(grpc_fut):
        if not loop.is_closed():
            loop.call_soon_threadsafe(_set_state, grpc_fut, fut_)

    fut_.add_done_callback(_call_check_cancel)
    grpc_fut.add_done_callback(_call_set_state)
    return fut_

def copy_members(source, dest, member_list, wrapper=None):
    for m in member_list:
        f = getattr(source, m, None)
        if f is None:
            continue
        if wrapper is not None:
            f = wrapper(f)
        setattr(dest, m, f)

def wrap_future_call(grpc_fut, loop, executor=None):
    fut_ = wrap_future(grpc_fut, loop)
    # Copy extra members
    copy_members(grpc_fut, fut_,
                 ['is_active',
                  'time_remaining'])
    @functools.wraps(grpc_fut.add_callback)
    def _add_callback(callback):
        grpc_fut.add_callback(wrap_callback(callback, loop))
    fut_.add_callback = _add_callback
    copy_members(grpc_fut, fut_,
                 ['initial_metadata',
                  'trailing_metadata',
                  'code',
                  'details'],
                 functools.partial(wrap_active_test, test=grpc_fut.is_active, loop=loop, executor=executor))
    return fut_

class WrappedIterator(object):
    """
    Wrap an grpc_iterator to an async iterator
    """
    def __init__(self, grpc_iterator, loop, executor=None, stream_executor=None):
        self._iterator = grpc_iterator
        self._loop = loop
        self._executor = executor
        if stream_executor is None:
            self._shared_executor = True
            self._stream_executor = executor
        else:
            self._shared_executor = False
            self._stream_executor = stream_executor
        self._next_future = None
        copy_members(grpc_iterator, self,
                     ['is_active',
                      'time_remaining',
                      'cancel'])
        @functools.wraps(grpc_iterator.add_callback)
        def _add_callback(callback):
            grpc_iterator.add_callback(wrap_callback(callback, loop))
        self.add_callback = _add_callback
        copy_members(grpc_iterator, self,
                     ['initial_metadata',
                      'trailing_metadata',
                      'code',
                      'details'],
                     functools.partial(wrap_active_test, test=grpc_iterator.is_active, loop=loop, executor=executor))

    def __aiter__(self):
        return self

    def _next(self):
        if self._iterator is None:
            raise StopAsyncIteration
        try:
            return next(self._iterator)
        except StopIteration:
            raise StopAsyncIteration
        except Exception:
            raise

    async def __anext__(self):
        if self._next_future is None:
            if self._iterator is None:
                raise StopAsyncIteration
            self._next_future = self._loop.run_in_executor(self._stream_executor, self._next)
        try:
            return await asyncio.shield(self._next_future)
        finally:
            if self._next_future and self._next_future.done():
                self._next_future = None

    def __del__(self):
        if self._iterator is not None:
            self.cancel()
            self._iterator = None
        if self._next_future is not None:
            if not self._loop.is_closed():
                self._loop.call_soon_threadsafe(lambda f=self._next_future: f.cancel())
            self._next_future = None
        if not self._shared_executor and self._stream_executor is not None:
            self._stream_executor.shutdown()
            self._stream_executor = None

    async def aclose(self):
        self.__del__()

class IteratorScope(object):
    def __init__(self, _iter):
        self._iter = _iter

    async def __aenter__(self):
        return self._iter

    async def __aexit__(self, exc_val, exc_typ, exc_tb):
        await self._iter.aclose()

class WrappedAsyncIterator(object):
    """
    Wrap an async iterator to an iterator for grpc input
    """
    def __init__(self, async_iter, loop):
        self._async_iter = async_iter
        self._loop = loop
        self._q = queue.Queue()
        self._stop_future = loop.create_future()
        self._next_future = None
        self._closed = False

    def __iter__(self):
        return self

    async def _next(self):
        if self._async_iter is None:
            # An edge condition
            self._q.put((None, True))
            return
        if self._next_future is None:
            self._next_future = asyncio.ensure_future(self._async_iter.__anext__())
        try:
            done, _ = await asyncio.wait([self._stop_future, self._next_future],
                                   return_when=asyncio.FIRST_COMPLETED)
            if self._stop_future in done:
                self._q.put((await self._stop_future, True))
                self._next_future.cancel()
                try:
                    await self._next_future
                except CancelledError:
                    pass
                finally:
                    self._next_future = None
            else:
                nf = self._next_future
                self._next_future = None
                self._q.put((await nf, False))
        except StopAsyncIteration:
            self._q.put((None, True))
        except Exception as exc:
            self._q.put((exc, True))

    def __next__(self):
        if self._async_iter is None:
            raise StopIteration
        try:
            r, is_exc = self._q.get_nowait()
        except queue.Empty:
            if not self._loop.is_closed():
                self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, self._next()))
            r, is_exc = self._q.get()
        if is_exc:
            if r is None:
                self._async_iter = None
                raise StopIteration
            else:
                raise r
        else:
            return r

    def close(self):
        if self._async_iter is not None:
            async def async_close():
                if not self._stop_future.done():
                    self._stop_future.set_result(None)
                await self._async_iter.aclose()
            try:
                if not self._loop.is_closed():
                    self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, async_close()))
            finally:
                # Ensure __next__ ends
                self._q.put((None, True))
                self._async_iter = None

    def cancel(self, exception=True):
        if exception:
            exc = CancelledError()
        else:
            exc = None
        def _set_result():
            if not self._stop_future.done():
                self._stop_future.set_result(exc)
        # Ensure __next__ ends. Sometimes the loop is already closing, so the exit result may not be written
        # to the queue
        self._q.put((exc, True))
        if not self._loop.is_closed():
            self._loop.call_soon_threadsafe(_set_result)
<commit_msg>Fix async for "Future exception was never retrieved" when task cancelled #16<commit_after>'''
Created on 2017/8/14

:author: hubo
'''

import asyncio
import functools
import queue
from asyncio import CancelledError
import logging
import grpc

logger = logging.getLogger(__name__)

def wrap_callback(callback, loop):
    @functools.wraps(callback)
    def _callback(*args, **kwargs):
        if not loop.is_closed():
            loop.call_soon_threadsafe(functools.partial(callback, *args, **kwargs))
    return _callback

def wrap_active_test(func, test, loop, executor=None):
    @functools.wraps(func)
    async def _func(*args, **kwargs):
        if test():
            return await loop.run_in_executor(executor, functools.partial(func, *args, **kwargs))
        else:
            return func(*args, **kwargs)
    return _func

def wrap_future(grpc_fut, loop):
    fut_ = loop.create_future()
    def _set_state(grpc_fut, fut_):
        assert grpc_fut.done()
        if fut_.cancelled():
            return
        assert not fut_.done()
        if grpc_fut.cancelled():
            fut_.cancel()
        else:
            exception = grpc_fut.exception()
            if exception is not None:
                fut_.set_exception(exception)
            else:
                result = grpc_fut.result()
                fut_.set_result(result)

    def _call_check_cancel(fut_):
        if fut_.cancelled():
            grpc_fut.cancel()

    def _call_set_state(grpc_fut):
        if not loop.is_closed():
            loop.call_soon_threadsafe(_set_state, grpc_fut, fut_)

    fut_.add_done_callback(_call_check_cancel)
    grpc_fut.add_done_callback(_call_set_state)
    return fut_

def copy_members(source, dest, member_list, wrapper=None):
    for m in member_list:
        f = getattr(source, m, None)
        if f is None:
            continue
        if wrapper is not None:
            f = wrapper(f)
        setattr(dest, m, f)

def wrap_future_call(grpc_fut, loop, executor=None):
    fut_ = wrap_future(grpc_fut, loop)
    # Copy extra members
    copy_members(grpc_fut, fut_,
                 ['is_active',
                  'time_remaining'])
    @functools.wraps(grpc_fut.add_callback)
    def _add_callback(callback):
        grpc_fut.add_callback(wrap_callback(callback, loop))
    fut_.add_callback = _add_callback
    copy_members(grpc_fut, fut_,
                 ['initial_metadata',
                  'trailing_metadata',
                  'code',
                  'details'],
                 functools.partial(wrap_active_test, test=grpc_fut.is_active, loop=loop, executor=executor))
    return fut_

class WrappedIterator(object):
    """
    Wrap an grpc_iterator to an async iterator
    """
    def __init__(self, grpc_iterator, loop, executor=None, stream_executor=None):
        self._iterator = grpc_iterator
        self._loop = loop
        self._executor = executor
        if stream_executor is None:
            self._shared_executor = True
            self._stream_executor = executor
        else:
            self._shared_executor = False
            self._stream_executor = stream_executor
        self._next_future = None
        copy_members(grpc_iterator, self,
                     ['is_active',
                      'time_remaining',
                      'cancel'])
        @functools.wraps(grpc_iterator.add_callback)
        def _add_callback(callback):
            grpc_iterator.add_callback(wrap_callback(callback, loop))
        self.add_callback = _add_callback
        copy_members(grpc_iterator, self,
                     ['initial_metadata',
                      'trailing_metadata',
                      'code',
                      'details'],
                     functools.partial(wrap_active_test, test=grpc_iterator.is_active, loop=loop, executor=executor))

    def __aiter__(self):
        return self

    def _next(self):
        if self._iterator is None:
            raise StopAsyncIteration
        try:
            return next(self._iterator)
        except StopIteration:
            raise StopAsyncIteration
        except Exception:
            raise

    async def __anext__(self):
        if self._next_future is None:
            if self._iterator is None:
                raise StopAsyncIteration
            self._next_future = self._loop.run_in_executor(self._stream_executor, self._next)

            def cb(fut):
                try:
                    fut.result()
                except grpc.RpcError as ex:
                    if ex.code() != grpc.StatusCode.CANCELLED:
                        logger.exception("__anext__ grpc exception")

            self._next_future.add_done_callback(cb)
        try:
            return await asyncio.shield(self._next_future)
        finally:
            if self._next_future and self._next_future.done():
                self._next_future = None

    def __del__(self):
        if self._iterator is not None:
            self.cancel()
            self._iterator = None
        if self._next_future is not None:
            if not self._loop.is_closed():
                self._loop.call_soon_threadsafe(lambda f=self._next_future: f.cancel())
            self._next_future = None
        if not self._shared_executor and self._stream_executor is not None:
            self._stream_executor.shutdown()
            self._stream_executor = None

    async def aclose(self):
        self.__del__()

class IteratorScope(object):
    def __init__(self, _iter):
        self._iter = _iter

    async def __aenter__(self):
        return self._iter

    async def __aexit__(self, exc_val, exc_typ, exc_tb):
        await self._iter.aclose()

class WrappedAsyncIterator(object):
    """
    Wrap an async iterator to an iterator for grpc input
    """
    def __init__(self, async_iter, loop):
        self._async_iter = async_iter
        self._loop = loop
        self._q = queue.Queue()
        self._stop_future = loop.create_future()
        self._next_future = None
        self._closed = False

    def __iter__(self):
        return self

    async def _next(self):
        if self._async_iter is None:
            # An edge condition
            self._q.put((None, True))
            return
        if self._next_future is None:
            self._next_future = asyncio.ensure_future(self._async_iter.__anext__())
        try:
            done, _ = await asyncio.wait([self._stop_future, self._next_future],
                                   return_when=asyncio.FIRST_COMPLETED)
            if self._stop_future in done:
                self._q.put((await self._stop_future, True))
                self._next_future.cancel()
                try:
                    await self._next_future
                except CancelledError:
                    pass
                finally:
                    self._next_future = None
            else:
                nf = self._next_future
                self._next_future = None
                self._q.put((await nf, False))
        except StopAsyncIteration:
            self._q.put((None, True))
        except Exception as exc:
            self._q.put((exc, True))

    def __next__(self):
        if self._async_iter is None:
            raise StopIteration
        try:
            r, is_exc = self._q.get_nowait()
        except queue.Empty:
            if not self._loop.is_closed():
                self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, self._next()))
            r, is_exc = self._q.get()
        if is_exc:
            if r is None:
                self._async_iter = None
                raise StopIteration
            else:
                raise r
        else:
            return r

    def close(self):
        if self._async_iter is not None:
            async def async_close():
                if not self._stop_future.done():
                    self._stop_future.set_result(None)
                await self._async_iter.aclose()
            try:
                if not self._loop.is_closed():
                    self._loop.call_soon_threadsafe(functools.partial(asyncio.ensure_future, async_close()))
            finally:
                # Ensure __next__ ends
                self._q.put((None, True))
                self._async_iter = None

    def cancel(self, exception=True):
        if exception:
            exc = CancelledError()
        else:
            exc = None
        def _set_result():
            if not self._stop_future.done():
                self._stop_future.set_result(exc)
        # Ensure __next__ ends. Sometimes the loop is already closing, so the exit result may not be written
        # to the queue
        self._q.put((exc, True))
        if not self._loop.is_closed():
            self._loop.call_soon_threadsafe(_set_result)
<|endoftext|>


<commit_before>import asyncio
import asyncio.events
import errno
import re
import threading
import zmq

from asyncio.unix_events import SelectorEventLoop, SafeChildWatcher
from asyncio.transports import _FlowControlMixin
from collections import deque, Iterable, Set
from ipaddress import ip_address

from .selector import ZmqSelector
from .interface import ZmqTransport

__all__ = ['ZmqEventLoop', 'ZmqEventLoopPolicy']

class ZmqEventLoop(SelectorEventLoop):
    """ZeroMQ event loop.

    Follows asyncio.AbstractEventLoop specification, in addition implements
    create_zmq_connection method for working with ZeroMQ sockets.
    """

    def __init__(self, *, io_threads=1):
        super().__init__(selector=ZmqSelector())
        self._zmq_context = zmq.Context(io_threads)

    def close(self):
        super().close()
        self._zmq_context.destroy()

    @asyncio.coroutine
    def create_zmq_connection(self, protocol_factory, zmq_type, *,
                              bind=None, connect=None, zmq_sock=None):
        """A coroutine which creates a ZeroMQ connection endpoint.

        The return value is a pair of (transport, protocol),
        where transport support ZmqTransport interface.

        protocol_factory should instantiate object with ZmqProtocol interface.

        zmq_type is type of ZeroMQ socket (zmq.REQ, zmq.REP, zmq.PUB, zmq.SUB,
        zmq.PAIR, zmq.DEALER, zmq.ROUTER, zmq.PULL, zmq.PUSH, etc.)

        bind is string or iterable of strings that specifies enpoints.
        Every endpoint creates ending for acceptin connections
        and binds it to the transport.
        Other side should use connect parameter to connect to this transport.
        See http://api.zeromq.org/master:zmq-bind for details.

        connect is string or iterable of strings that specifies enpoints.
        Every endpoint connects transport to specified transport.
        Other side should use bind parameter to wait for incoming connections.
        See http://api.zeromq.org/master:zmq-connect for details.

        endpoint is a string consisting of two parts as follows:
        transport://address.

        The transport part specifies the underlying transport protocol to use.
        The meaning of the address part is specific to the underlying
        transport protocol selected.

        The following transports are defined:

        inproc - local in-process (inter-thread) communication transport,
        see http://api.zeromq.org/master:zmq-inproc.

        ipc - local inter-process communication transport,
        see http://api.zeromq.org/master:zmq-ipc

        tcp - unicast transport using TCP,
        see http://api.zeromq.org/master:zmq_tcp

        pgm, epgm - reliable multicast transport using PGM,
        see http://api.zeromq.org/master:zmq_pgm

        zmq_sock is a zmq.Socket instance to use preexisting object
        with created transport.

        The only one of bind, connect or zmq_sock should be specified.
        """
        if 1 != sum(
                1 if i is not None else 0 for i in [zmq_sock, bind, connect]):
            raise ValueError(
                "the only bind, connect or zmq_sock should be specified "
                "at the same time", bind, connect, zmq_sock)

        try:
            if zmq_sock is None:
                zmq_sock = self._zmq_context.socket(zmq_type)
            elif zmq_sock.getsockopt(zmq.TYPE) != zmq_type:
                raise ValueError('Invalid zmq_sock type')
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc

        protocol = protocol_factory()
        waiter = asyncio.Future(loop=self)
        transport = _ZmqTransportImpl(self, zmq_sock, protocol, waiter)
        yield from waiter

        try:
            if bind is not None:
                if isinstance(bind, str):
                    bind = [bind]
                else:
                    if not isinstance(bind, Iterable):
                        raise ValueError('bind should be str or iterable')
                for endpoint in bind:
                    transport.bind(endpoint)
            elif connect is not None:
                if isinstance(connect, str):
                    connect = [connect]
                else:
                    if not isinstance(connect, Iterable):
                        raise ValueError('connect should be '
                                         'str or iterable')
                for endpoint in connect:
                    transport.connect(endpoint)

            return transport, protocol
        except OSError:
            # don't care if zmq_sock.close can raise exception
            # that should never happen
            zmq_sock.close()
            raise

class _EndpointsSet(Set):

    __slots__ = ('_collection',)

    def __init__(self, collection):
        self._collection = collection

    def __len__(self):
        return len(self._collection)

    def __contains__(self, endpoint):
        return endpoint in self._collection

    def __iter__(self):
        return iter(self._collection)

    def __repr__(self):
        return '{' + ', '.join(sorted(self._collection)) + '}'

    __str__ = __repr__

class _ZmqTransportImpl(ZmqTransport, _FlowControlMixin):

    _TCP_RE = re.compile('^tcp://(.+):(\d+)|\*$')

    def __init__(self, loop, zmq_sock, protocol, waiter=None):
        super().__init__(None)
        self._extra['zmq_socket'] = zmq_sock
        self._loop = loop
        self._zmq_sock = zmq_sock
        self._protocol = protocol
        self._closing = False
        self._buffer = deque()
        self._buffer_size = 0
        self._listeners = set()
        self._connections = set()

        self._loop.add_reader(self._zmq_sock, self._read_ready)
        self._loop.call_soon(self._protocol.connection_made, self)
        if waiter is not None:
            self._loop.call_soon(waiter.set_result, None)

    def _read_ready(self):
        try:
            try:
                data = self._zmq_sock.recv_multipart(zmq.NOBLOCK)
            except zmq.ZMQError as exc:
                if exc.errno in (errno.EAGAIN, errno.EINTR):
                    return
                else:
                    raise OSError(exc.errno, exc.strerror) from exc
        except Exception as exc:
            self._fatal_error(exc, 'Fatal read error on zmq socket transport')
        else:
            self._protocol.msg_received(tuple(data))

    def write(self, data):
        if not data:
            return
        for part in data:
            if not isinstance(part, (bytes, bytearray, memoryview)):
                raise TypeError('data argument must be iterable of '
                                'byte-ish (%r)' % data)
        data_len = sum(len(part) for part in data)
        if not data_len:
            return

        if not self._buffer:
            try:
                try:
                    self._zmq_sock.send_multipart(data, zmq.DONTWAIT)
                    return
                except zmq.ZMQError as exc:
                    if exc.errno in (errno.EAGAIN, errno.EINTR):
                        self._loop.add_writer(self._zmq_sock,
                                              self._write_ready)
                    else:
                        raise OSError(exc.errno, exc.strerror) from exc
            except Exception as exc:
                self._fatal_error(exc,
                                  'Fatal write error on zmq socket transport')
                return

        self._buffer.append(data)
        self._buffer_size += data_len
        self._maybe_pause_protocol()

    def _write_ready(self):
        assert self._buffer, 'Data should not be empty'

        try:
            try:
                self._zmq_sock.send_multipart(self._buffer[0], zmq.DONTWAIT)
            except zmq.ZMQError as exc:
                if exc.errno in (errno.EAGAIN, errno.EINTR):
                    pass
                else:
                    raise OSError(exc.errno, exc.strerror) from exc
        except Exception as exc:
            self._loop.remove_writer(self._zmq_sock)
            self._buffer.clear()
            self._buffer_size = 0
            self._fatal_error(exc,
                              'Fatal write error on zmq socket transport')
        else:
            sent_data = self._buffer.popleft()
            self._buffer_size -= sum(len(part) for part in sent_data)

            self._maybe_resume_protocol()

            if not self._buffer:
                self._loop.remove_writer(self._zmq_sock)
                if self._closing:
                    self._call_connection_lost(None)

    def can_write_eof(self):
        return False

    def abort(self):
        self._force_close(None)

    def close(self):
        if self._closing:
            return
        self._closing = True
        self._loop.remove_reader(self._zmq_sock)
        if not self._buffer:
            self._loop.call_soon(self._call_connection_lost, None)

    def _fatal_error(self, exc, message='Fatal error on transport'):
        # Should be called from exception handler only.
        self._loop.call_exception_handler({
            'message': message,
            'exception': exc,
            'transport': self,
            'protocol': self._protocol,
            })
        self._force_close(exc)

    def _force_close(self, exc):
        if self._buffer:
            self._buffer.clear()
            self._buffer_size = 0
            self._loop.remove_writer(self._zmq_sock)
        if not self._closing:
            self._closing = True
            self._loop.remove_reader(self._zmq_sock)
        self._loop.call_soon(self._call_connection_lost, exc)

    def _call_connection_lost(self, exc):
        try:
            self._protocol.connection_lost(exc)
        finally:
            self._zmq_sock.close()
            self._zmq_sock = None
            self._protocol = None
            self._loop = None

    def getsockopt(self, option):
        try:
            return self._zmq_sock.getsockopt(option)
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc

    def setsockopt(self, option, value):
        try:
            self._zmq_sock.setsockopt(option, value)
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc

    def get_write_buffer_size(self):
        return self._buffer_size

    def bind(self, endpoint):
        try:
            self._zmq_sock.bind(endpoint)
            breal_endpoint = self._zmq_sock.getsockopt(zmq.LAST_ENDPOINT)
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc
        else:
            real_endpoint = breal_endpoint.decode('utf-8').rstrip('\x00')
            self._listeners.add(real_endpoint)
            return real_endpoint

    def unbind(self, endpoint):
        try:
            self._zmq_sock.unbind(endpoint)
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc
        else:
            self._listeners.discard(endpoint)

    def listeners(self):
        return _EndpointsSet(self._listeners)

    def connect(self, endpoint):
        match = self._TCP_RE.match(endpoint)
        if match:
            ip_address(match.group(1))  # check for correct IPv4 or IPv6
        try:
            self._zmq_sock.connect(endpoint)
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc
        else:
            self._connections.add(endpoint)
            return endpoint

    def disconnect(self, endpoint):
        try:
            self._zmq_sock.disconnect(endpoint)
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc
        else:
            self._connections.discard(endpoint)

    def connections(self):
        return _EndpointsSet(self._connections)

class ZmqEventLoopPolicy(asyncio.AbstractEventLoopPolicy):
    """ZeroMQ policy implementation for accessing the event loop.

    In this policy, each thread has its own event loop.  However, we
    only automatically create an event loop by default for the main
    thread; other threads by default have no event loop.
    """

    class _Local(threading.local):
        _loop = None
        _set_called = False

    def __init__(self, *, io_threads=1):
        self._local = self._Local()
        self._watcher = None
        self._io_threads = io_threads

    def get_event_loop(self):
        """Get the event loop.

        If current thread is the main thread and there are no
        registered event loop for current thread then the call creates
        new event loop and registers it.

        Return an instance of ZmqEventLoop.
        Raise RuntimeError if there is no registered event loop
        for current thread.
        """
        if (self._local._loop is None and
                not self._local._set_called and
                isinstance(threading.current_thread(), threading._MainThread)):
            self.set_event_loop(self.new_event_loop())
        if self._local._loop is None:
            raise RuntimeError('There is no current event loop in thread %r.' %
                               threading.current_thread().name)
        return self._local._loop

    def new_event_loop(self):
        """Create a new event loop.

        You must call set_event_loop() to make this the current event
        loop.
        """
        return ZmqEventLoop(io_threads=self._io_threads)

    def _init_watcher(self):
        with asyncio.events._lock:
            if self._watcher is None:  # pragma: no branch
                self._watcher = SafeChildWatcher()
                if isinstance(threading.current_thread(),
                              threading._MainThread):
                    self._watcher.attach_loop(self._local._loop)

    def set_event_loop(self, loop):
        """Set the event loop.

        As a side effect, if a child watcher was set before, then calling
        .set_event_loop() from the main thread will call .attach_loop(loop) on
        the child watcher.
        """

        self._local._set_called = True
        if loop is not None and not isinstance(loop,
                                               asyncio.AbstractEventLoop):
            raise TypeError("loop should be None "
                            "or AbstractEventLoop instance")
        self._local._loop = loop

        if (self._watcher is not None and
                isinstance(threading.current_thread(), threading._MainThread)):
            self._watcher.attach_loop(loop)

    def get_child_watcher(self):
        """Get the child watcher.

        If not yet set, a SafeChildWatcher object is automatically created.
        """
        if self._watcher is None:
            self._init_watcher()

        return self._watcher

    def set_child_watcher(self, watcher):
        """Set the child watcher."""

        if (watcher is not None and
                not isinstance(watcher, asyncio.AbstractChildWatcher)):
            raise TypeError("watcher should be None or AbstractChildWatcher "
                            "instance")

        if self._watcher is not None:
            self._watcher.close()

        self._watcher = watcher
<commit_msg>Process EINTR in zmq socket method calls from transport<commit_after>import asyncio
import asyncio.events
import errno
import re
import threading
import zmq

from asyncio.unix_events import SelectorEventLoop, SafeChildWatcher
from asyncio.transports import _FlowControlMixin
from collections import deque, Iterable, Set
from ipaddress import ip_address

from .selector import ZmqSelector
from .interface import ZmqTransport

__all__ = ['ZmqEventLoop', 'ZmqEventLoopPolicy']

class ZmqEventLoop(SelectorEventLoop):
    """ZeroMQ event loop.

    Follows asyncio.AbstractEventLoop specification, in addition implements
    create_zmq_connection method for working with ZeroMQ sockets.
    """

    def __init__(self, *, io_threads=1):
        super().__init__(selector=ZmqSelector())
        self._zmq_context = zmq.Context(io_threads)

    def close(self):
        super().close()
        self._zmq_context.destroy()

    @asyncio.coroutine
    def create_zmq_connection(self, protocol_factory, zmq_type, *,
                              bind=None, connect=None, zmq_sock=None):
        """A coroutine which creates a ZeroMQ connection endpoint.

        The return value is a pair of (transport, protocol),
        where transport support ZmqTransport interface.

        protocol_factory should instantiate object with ZmqProtocol interface.

        zmq_type is type of ZeroMQ socket (zmq.REQ, zmq.REP, zmq.PUB, zmq.SUB,
        zmq.PAIR, zmq.DEALER, zmq.ROUTER, zmq.PULL, zmq.PUSH, etc.)

        bind is string or iterable of strings that specifies enpoints.
        Every endpoint creates ending for acceptin connections
        and binds it to the transport.
        Other side should use connect parameter to connect to this transport.
        See http://api.zeromq.org/master:zmq-bind for details.

        connect is string or iterable of strings that specifies enpoints.
        Every endpoint connects transport to specified transport.
        Other side should use bind parameter to wait for incoming connections.
        See http://api.zeromq.org/master:zmq-connect for details.

        endpoint is a string consisting of two parts as follows:
        transport://address.

        The transport part specifies the underlying transport protocol to use.
        The meaning of the address part is specific to the underlying
        transport protocol selected.

        The following transports are defined:

        inproc - local in-process (inter-thread) communication transport,
        see http://api.zeromq.org/master:zmq-inproc.

        ipc - local inter-process communication transport,
        see http://api.zeromq.org/master:zmq-ipc

        tcp - unicast transport using TCP,
        see http://api.zeromq.org/master:zmq_tcp

        pgm, epgm - reliable multicast transport using PGM,
        see http://api.zeromq.org/master:zmq_pgm

        zmq_sock is a zmq.Socket instance to use preexisting object
        with created transport.

        The only one of bind, connect or zmq_sock should be specified.
        """
        if 1 != sum(
                1 if i is not None else 0 for i in [zmq_sock, bind, connect]):
            raise ValueError(
                "the only bind, connect or zmq_sock should be specified "
                "at the same time", bind, connect, zmq_sock)

        try:
            if zmq_sock is None:
                zmq_sock = self._zmq_context.socket(zmq_type)
            elif zmq_sock.getsockopt(zmq.TYPE) != zmq_type:
                raise ValueError('Invalid zmq_sock type')
        except zmq.ZMQError as exc:
            raise OSError(exc.errno, exc.strerror) from exc

        protocol = protocol_factory()
        waiter = asyncio.Future(loop=self)
        transport = _ZmqTransportImpl(self, zmq_sock, protocol, waiter)
        yield from waiter

        try:
            if bind is not None:
                if isinstance(bind, str):
                    bind = [bind]
                else:
                    if not isinstance(bind, Iterable):
                        raise ValueError('bind should be str or iterable')
                for endpoint in bind:
                    transport.bind(endpoint)
            elif connect is not None:
                if isinstance(connect, str):
                    connect = [connect]
                else:
                    if not isinstance(connect, Iterable):
                        raise ValueError('connect should be '
                                         'str or iterable')
                for endpoint in connect:
                    transport.connect(endpoint)

            return transport, protocol
        except OSError:
            # don't care if zmq_sock.close can raise exception
            # that should never happen
            zmq_sock.close()
            raise

class _EndpointsSet(Set):

    __slots__ = ('_collection',)

    def __init__(self, collection):
        self._collection = collection

    def __len__(self):
        return len(self._collection)

    def __contains__(self, endpoint):
        return endpoint in self._collection

    def __iter__(self):
        return iter(self._collection)

    def __repr__(self):
        return '{' + ', '.join(sorted(self._collection)) + '}'

    __str__ = __repr__

class _ZmqTransportImpl(ZmqTransport, _FlowControlMixin):

    _TCP_RE = re.compile('^tcp://(.+):(\d+)|\*$')

    def __init__(self, loop, zmq_sock, protocol, waiter=None):
        super().__init__(None)
        self._extra['zmq_socket'] = zmq_sock
        self._loop = loop
        self._zmq_sock = zmq_sock
        self._protocol = protocol
        self._closing = False
        self._buffer = deque()
        self._buffer_size = 0
        self._listeners = set()
        self._connections = set()

        self._loop.add_reader(self._zmq_sock, self._read_ready)
        self._loop.call_soon(self._protocol.connection_made, self)
        if waiter is not None:
            self._loop.call_soon(waiter.set_result, None)

    def _read_ready(self):
        try:
            try:
                data = self._zmq_sock.recv_multipart(zmq.NOBLOCK)
            except zmq.ZMQError as exc:
                if exc.errno in (errno.EAGAIN, errno.EINTR):
                    return
                else:
                    raise OSError(exc.errno, exc.strerror) from exc
        except Exception as exc:
            self._fatal_error(exc, 'Fatal read error on zmq socket transport')
        else:
            self._protocol.msg_received(tuple(data))

    def write(self, data):
        if not data:
            return
        for part in data:
            if not isinstance(part, (bytes, bytearray, memoryview)):
                raise TypeError('data argument must be iterable of '
                                'byte-ish (%r)' % data)
        data_len = sum(len(part) for part in data)
        if not data_len:
            return

        if not self._buffer:
            try:
                try:
                    self._zmq_sock.send_multipart(data, zmq.DONTWAIT)
                    return
                except zmq.ZMQError as exc:
                    if exc.errno in (errno.EAGAIN, errno.EINTR):
                        self._loop.add_writer(self._zmq_sock,
                                              self._write_ready)
                    else:
                        raise OSError(exc.errno, exc.strerror) from exc
            except Exception as exc:
                self._fatal_error(exc,
                                  'Fatal write error on zmq socket transport')
                return

        self._buffer.append(data)
        self._buffer_size += data_len
        self._maybe_pause_protocol()

    def _write_ready(self):
        assert self._buffer, 'Data should not be empty'

        try:
            try:
                self._zmq_sock.send_multipart(self._buffer[0], zmq.DONTWAIT)
            except zmq.ZMQError as exc:
                if exc.errno in (errno.EAGAIN, errno.EINTR):
                    pass
                else:
                    raise OSError(exc.errno, exc.strerror) from exc
        except Exception as exc:
            self._loop.remove_writer(self._zmq_sock)
            self._buffer.clear()
            self._buffer_size = 0
            self._fatal_error(exc,
                              'Fatal write error on zmq socket transport')
        else:
            sent_data = self._buffer.popleft()
            self._buffer_size -= sum(len(part) for part in sent_data)

            self._maybe_resume_protocol()

            if not self._buffer:
                self._loop.remove_writer(self._zmq_sock)
                if self._closing:
                    self._call_connection_lost(None)

    def can_write_eof(self):
        return False

    def abort(self):
        self._force_close(None)

    def close(self):
        if self._closing:
            return
        self._closing = True
        self._loop.remove_reader(self._zmq_sock)
        if not self._buffer:
            self._loop.call_soon(self._call_connection_lost, None)

    def _fatal_error(self, exc, message='Fatal error on transport'):
        # Should be called from exception handler only.
        self._loop.call_exception_handler({
            'message': message,
            'exception': exc,
            'transport': self,
            'protocol': self._protocol,
            })
        self._force_close(exc)

    def _force_close(self, exc):
        if self._buffer:
            self._buffer.clear()
            self._buffer_size = 0
            self._loop.remove_writer(self._zmq_sock)
        if not self._closing:
            self._closing = True
            self._loop.remove_reader(self._zmq_sock)
        self._loop.call_soon(self._call_connection_lost, exc)

    def _call_connection_lost(self, exc):
        try:
            self._protocol.connection_lost(exc)
        finally:
            self._zmq_sock.close()
            self._zmq_sock = None
            self._protocol = None
            self._loop = None

    def getsockopt(self, option):
        while True:
            try:
                return self._zmq_sock.getsockopt(option)
            except zmq.ZMQError as exc:
                if exc.errno == errno.EINTR:
                    continue
                raise OSError(exc.errno, exc.strerror) from exc

    def setsockopt(self, option, value):
        while True:
            try:
                self._zmq_sock.setsockopt(option, value)
                return
            except zmq.ZMQError as exc:
                if exc.errno == errno.EINTR:
                    continue
                raise OSError(exc.errno, exc.strerror) from exc

    def get_write_buffer_size(self):
        return self._buffer_size

    def bind(self, endpoint):
        while True:
            try:
                self._zmq_sock.bind(endpoint)
                breal_endpoint = self._zmq_sock.getsockopt(zmq.LAST_ENDPOINT)
            except zmq.ZMQError as exc:
                if exc.errno == errno.EINTR:
                    continue
                raise OSError(exc.errno, exc.strerror) from exc
            else:
                real_endpoint = breal_endpoint.decode('utf-8').rstrip('\x00')
                self._listeners.add(real_endpoint)
                return real_endpoint

    def unbind(self, endpoint):
        while True:
            try:
                self._zmq_sock.unbind(endpoint)
            except zmq.ZMQError as exc:
                if exc.errno == errno.EINTR:
                    continue
                raise OSError(exc.errno, exc.strerror) from exc
            else:
                self._listeners.discard(endpoint)
                return

    def listeners(self):
        return _EndpointsSet(self._listeners)

    def connect(self, endpoint):
        match = self._TCP_RE.match(endpoint)
        if match:
            ip_address(match.group(1))  # check for correct IPv4 or IPv6
        while True:
            try:
                self._zmq_sock.connect(endpoint)
            except zmq.ZMQError as exc:
                if exc.errno == errno.EINTR:
                    continue
                raise OSError(exc.errno, exc.strerror) from exc
            else:
                self._connections.add(endpoint)
                return endpoint

    def disconnect(self, endpoint):
        while True:
            try:
                self._zmq_sock.disconnect(endpoint)
            except zmq.ZMQError as exc:
                if exc.errno == errno.EINTR:
                    continue
                raise OSError(exc.errno, exc.strerror) from exc
            else:
                self._connections.discard(endpoint)
                return

    def connections(self):
        return _EndpointsSet(self._connections)

class ZmqEventLoopPolicy(asyncio.AbstractEventLoopPolicy):
    """ZeroMQ policy implementation for accessing the event loop.

    In this policy, each thread has its own event loop.  However, we
    only automatically create an event loop by default for the main
    thread; other threads by default have no event loop.
    """

    class _Local(threading.local):
        _loop = None
        _set_called = False

    def __init__(self, *, io_threads=1):
        self._local = self._Local()
        self._watcher = None
        self._io_threads = io_threads

    def get_event_loop(self):
        """Get the event loop.

        If current thread is the main thread and there are no
        registered event loop for current thread then the call creates
        new event loop and registers it.

        Return an instance of ZmqEventLoop.
        Raise RuntimeError if there is no registered event loop
        for current thread.
        """
        if (self._local._loop is None and
                not self._local._set_called and
                isinstance(threading.current_thread(), threading._MainThread)):
            self.set_event_loop(self.new_event_loop())
        if self._local._loop is None:
            raise RuntimeError('There is no current event loop in thread %r.' %
                               threading.current_thread().name)
        return self._local._loop

    def new_event_loop(self):
        """Create a new event loop.

        You must call set_event_loop() to make this the current event
        loop.
        """
        return ZmqEventLoop(io_threads=self._io_threads)

    def _init_watcher(self):
        with asyncio.events._lock:
            if self._watcher is None:  # pragma: no branch
                self._watcher = SafeChildWatcher()
                if isinstance(threading.current_thread(),
                              threading._MainThread):
                    self._watcher.attach_loop(self._local._loop)

    def set_event_loop(self, loop):
        """Set the event loop.

        As a side effect, if a child watcher was set before, then calling
        .set_event_loop() from the main thread will call .attach_loop(loop) on
        the child watcher.
        """

        self._local._set_called = True
        if loop is not None and not isinstance(loop,
                                               asyncio.AbstractEventLoop):
            raise TypeError("loop should be None "
                            "or AbstractEventLoop instance")
        self._local._loop = loop

        if (self._watcher is not None and
                isinstance(threading.current_thread(), threading._MainThread)):
            self._watcher.attach_loop(loop)

    def get_child_watcher(self):
        """Get the child watcher.

        If not yet set, a SafeChildWatcher object is automatically created.
        """
        if self._watcher is None:
            self._init_watcher()

        return self._watcher

    def set_child_watcher(self, watcher):
        """Set the child watcher."""

        if (watcher is not None and
                not isinstance(watcher, asyncio.AbstractChildWatcher)):
            raise TypeError("watcher should be None or AbstractChildWatcher "
                            "instance")

        if self._watcher is not None:
            self._watcher.close()

        self._watcher = watcher
<|endoftext|>


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

from aldryn_client import forms

class Form(forms.BaseForm):

    def to_settings(self, cleaned_data, settings_dict):
        import aldryn_geoip

        settings_dict.setdefault('GEOIP_PATH', aldryn_geoip.where())
        # django.contrib.gis.geoip.GeoIP.GEOIP_STANDARD
        settings_dict.setdefault('GEOIP_CACHE_METHOD', 0)

        country_mw = 'country_segment.middleware.ResolveCountryCodeMiddleware'
        if country_mw not in settings_dict['MIDDLEWARE_CLASSES']:
            for position, mw in enumerate(settings_dict['MIDDLEWARE_CLASSES']):
                #
                # Its not a strict requirement that the
                # ResolveCountryCodeMiddleware go after SessionMiddleware,
                # but, it seems like a pretty nice place.
                #
                if mw == 'django.contrib.sessions.middleware.SessionMiddleware':
                    settings_dict['MIDDLEWARE_CLASSES'].insert(position + 1, country_mw)
                    break
            else:
                #
                # B'okay, not sure how this CMS installation works, but...
                # let's just put it at the top.
                #
                settings_dict['MIDDLEWARE_CLASSES'].insert(0, country_mw)

        return settings_dict
<commit_msg>use X-Forwarded-For on aldryn<commit_after># -*- coding: utf-8 -*-

from aldryn_client import forms

class Form(forms.BaseForm):

    def to_settings(self, cleaned_data, settings_dict):
        import aldryn_geoip

        settings_dict.setdefault('GEOIP_PATH', aldryn_geoip.where())
        # django.contrib.gis.geoip.GeoIP.GEOIP_STANDARD
        settings_dict.setdefault('GEOIP_CACHE_METHOD', 0)

        # On Aldryn default to the X-Forwarded-For header
        settings_dict.setdefault(
            'GEOIP_REQUEST_IP_RESOLVER',
            'aldryn_geoip.utils.x_forwarded_ip'
        )

        country_mw = 'country_segment.middleware.ResolveCountryCodeMiddleware'
        if country_mw not in settings_dict['MIDDLEWARE_CLASSES']:
            for position, mw in enumerate(settings_dict['MIDDLEWARE_CLASSES']):
                #
                # Its not a strict requirement that the
                # ResolveCountryCodeMiddleware go after SessionMiddleware,
                # but, it seems like a pretty nice place.
                #
                if mw == 'django.contrib.sessions.middleware.SessionMiddleware':
                    settings_dict['MIDDLEWARE_CLASSES'].insert(position + 1, country_mw)
                    break
            else:
                #
                # B'okay, not sure how this CMS installation works, but...
                # let's just put it at the top.
                #
                settings_dict['MIDDLEWARE_CLASSES'].insert(0, country_mw)

        return settings_dict
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

def _previsit(v, previsited_d, clock):
    clock += 1
    previsited_d[v] = clock
    return previsited_d, clock

def _postvisit(v, postvisited_d, clock):
    clock += 1
    postvisited_d[v] = clock
    return postvisited_d, clock

def _dfs_explore(v, graph_adj_d, visited_d,
                 previsited_d, postvisited_d, clock, dag_bool):
    visited_d[v] = True
    previsited_d, clock = _previsit(v, previsited_d, clock)
    for v_neighbor in graph_adj_d[v]:
        print('{0} -> {1}'.format(v, v_neighbor))
        if (visited_d[v_neighbor]):
            if (previsited_d[v] > previsited_d[v_neighbor] and
                not v_neighbor in postvisited_d):
                dag_bool = False
        if not visited_d[v_neighbor]:
            previsited_d, postvisited_d, clock, dag_bool = (
                _dfs_explore(v_neighbor, graph_adj_d, visited_d,
                             previsited_d, postvisited_d, clock, dag_bool))
    postvisited_d, clock = _postvisit(v, postvisited_d, clock)
    return previsited_d, postvisited_d, clock, dag_bool

def check_dag(graph_adj_d):
    """Check Directed Acyclic Graph (DAG)."""
    visited_d = {v: False for v in graph_adj_d.keys()}
    previsited_d = {}
    postvisited_d = {}
    clock = 0
    dag_bool = True
    for v in graph_adj_d.keys():
        if not visited_d[v] and dag_bool:
            previsited_d, postvisited_d, clock, dag_bool = (
                _dfs_explore(v, graph_adj_d, visited_d,
                             previsited_d, postvisited_d, clock, dag_bool))
    return dag_bool

def main():
    # Graph adjacency dictionary for DAG.
    dag_adj_d = {
        'A': ['D'],
        'B': ['D'],
        'C': ['D'],
        'D': ['E', 'G'],
        'E': ['J'],
        'F': ['G'],
        'G': ['I'],
        'I': ['J'],
        'J': []
    }
    print('dag_adj_d: {}'.format(dag_adj_d))
    print('Is DAG: {}'.format(check_dag(dag_adj_d)))

    # Graph adjacency dictionary for non-DAG.
    nondag_adj_d = {
        'A': ['B'],
        'B': ['C', 'D'],
        'C': ['A', 'C'],
        'D': ['B', 'C']
    }
    print('nondag_adj_d: {}'.format(nondag_adj_d))
    print('Is DAG: {}'.format(check_dag(nondag_adj_d)))

if __name__ == '__main__':
    main()
<commit_msg>Revise DAG's adjacency dictionary<commit_after>from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

def _previsit(v, previsited_d, clock):
    clock += 1
    previsited_d[v] = clock
    return previsited_d, clock

def _postvisit(v, postvisited_d, clock):
    clock += 1
    postvisited_d[v] = clock
    return postvisited_d, clock

def _dfs_explore(v, graph_adj_d, visited_d,
                 previsited_d, postvisited_d, clock, dag_bool):
    visited_d[v] = True
    previsited_d, clock = _previsit(v, previsited_d, clock)
    for v_neighbor in graph_adj_d[v]:
        print('{0} -> {1}'.format(v, v_neighbor))
        if (visited_d[v_neighbor]):
            if (previsited_d[v] > previsited_d[v_neighbor] and
                not v_neighbor in postvisited_d):
                dag_bool = False
        if not visited_d[v_neighbor]:
            previsited_d, postvisited_d, clock, dag_bool = (
                _dfs_explore(v_neighbor, graph_adj_d, visited_d,
                             previsited_d, postvisited_d, clock, dag_bool))
    postvisited_d, clock = _postvisit(v, postvisited_d, clock)
    return previsited_d, postvisited_d, clock, dag_bool

def check_dag(graph_adj_d):
    """Check Directed Acyclic Graph (DAG)."""
    visited_d = {v: False for v in graph_adj_d.keys()}
    previsited_d = {}
    postvisited_d = {}
    clock = 0
    dag_bool = True
    for v in graph_adj_d.keys():
        if not visited_d[v] and dag_bool:
            previsited_d, postvisited_d, clock, dag_bool = (
                _dfs_explore(v, graph_adj_d, visited_d,
                             previsited_d, postvisited_d, clock, dag_bool))
    return dag_bool

def main():
    # DAG's adjacency dictionary.
    dag_adj_d = {
        'A': ['B', 'F'],
        'B': ['C', 'D', 'F'],
        'C': ['D'],
        'D': ['E', 'F'],
        'E': ['F'],
        'F': []
    }
    print('dag_adj_d: {}'.format(dag_adj_d))
    print('Is DAG: {}'.format(check_dag(dag_adj_d)))

    # non-DAG's adjacency dictionary.
    nondag_adj_d = {
        'A': ['B'],
        'B': ['C', 'D'],
        'C': ['A', 'C'],
        'D': ['B', 'C']
    }
    print('nondag_adj_d: {}'.format(nondag_adj_d))
    print('Is DAG: {}'.format(check_dag(nondag_adj_d)))

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


<commit_before>from __future__ import with_statement
import tempfile
import sys
import os
from sys import stdout
from contextlib import closing
from alloc import alloc

class mbox(alloc):

  one_line_help = "Download a task's emails to an mbox file."

  # Setup the options that this cli can accept
  ops = []
  ops.append((''  ,'help           ','Show this help.'))
  ops.append(('q' ,'quiet          ','Run with no output except errors.'))
  ops.append(('t:','task=ID|NAME   ','A task ID, or a fuzzy match for a task name.'))

  # Specify some header and footer text for the help text
  help_text = "Usage: %s [OPTIONS]\n"
  help_text+= one_line_help
  help_text+= '''\n\n%s

This program will automatically run $MAILER on the mbox file, if outputting to a TTY.

Examples:
alloc mbox -t 1234
alloc mbox -t 1234 > file.mbox'''

  def run(self):

    # Get the command line arguments into a dictionary
    o, remainder = self.get_args(self.ops, self.help_text)

    self.quiet = o['quiet']
    taskID = ''
    emailUIDs = ''

    # Got this far, then authenticate
    self.authenticate();

    # Get a taskID either passed via command line, or figured out from a task name
    tops = {}
    if self.is_num(o['task']):
      taskID = o['task']
    elif o['task']:
      tops = {}
      tops["taskName"] = o["task"]
      tops["taskView"] = "prioritised"
      taskID = self.search_for_task(tops)

    if taskID:

      # If we're redirecting stdout eg -t 123 >task123.html
      if not stdout.isatty():
        print self.make_request({"method":"search_emails","str":'SUBJECT "Task Comment: '+taskID+' "'})

      else:
        if not 'MAILER' in os.environ or not os.environ['MAILER']:
          self.die('The environment variable $MAILER has not been defined. Eg: export MAILER="mutt -f "')

        str = self.make_request({"method":"search_emails","str":'SUBJECT "Task Comment: '+taskID+' "'})

        fd, filepath = tempfile.mkstemp()
        with closing(os.fdopen(fd, 'wb')) as tf:
          tf.write(unicode(str).encode('utf-8'))

        command = os.environ['MAILER']+' "'+filepath+'"'
        self.msg('Running: '+command)
        os.system(command)
        self.msg('Removing: '+filepath)
        os.remove(filepath)

<commit_msg>18139 Added time sheet item comments to `alloc mbox`.<commit_after>from __future__ import with_statement
import tempfile
import sys
import os
from sys import stdout
from contextlib import closing
from alloc import alloc

class mbox(alloc):

  one_line_help = "Download a task's emails to an mbox file."

  # Setup the options that this cli can accept
  ops = []
  ops.append((''  ,'help           ','Show this help.'))
  ops.append(('q' ,'quiet          ','Run with no output except errors.'))
  ops.append(('t:','task=ID|NAME   ','A task ID, or a fuzzy match for a task name.'))

  # Specify some header and footer text for the help text
  help_text = "Usage: %s [OPTIONS]\n"
  help_text+= one_line_help
  help_text+= '''\n\n%s

This program will automatically run $MAILER on the mbox file, if outputting to a TTY.

Examples:
alloc mbox -t 1234
alloc mbox -t 1234 > file.mbox'''

  def run(self):

    # Get the command line arguments into a dictionary
    o, remainder = self.get_args(self.ops, self.help_text)

    self.quiet = o['quiet']
    taskID = ''
    emailUIDs = ''

    # Got this far, then authenticate
    self.authenticate();

    # Get a taskID either passed via command line, or figured out from a task name
    tops = {}
    if self.is_num(o['task']):
      taskID = o['task']
    elif o['task']:
      tops = {}
      tops["taskName"] = o["task"]
      tops["taskView"] = "prioritised"
      taskID = self.search_for_task(tops)

    if taskID:

      # If we're redirecting stdout eg -t 123 >task123.html
      if not stdout.isatty():
        str = self.make_request({"method":"search_emails","str":'SUBJECT "Task Comment: '+taskID+' "'})
        str += "\n\n" + self.make_request({"method":"get_timeSheetItem_comments","taskID":taskID})
        print str

      else:
        if not 'MAILER' in os.environ or not os.environ['MAILER']:
          self.die('The environment variable $MAILER has not been defined. Eg: export MAILER="mutt -f "')

        str = self.make_request({"method":"search_emails","str":'SUBJECT "Task Comment: '+taskID+' "'})
        str += "\n\n" + self.make_request({"method":"get_timeSheetItem_comments","taskID":taskID})

        fd, filepath = tempfile.mkstemp()
        with closing(os.fdopen(fd, 'wb')) as tf:
          tf.write(unicode(str).encode('utf-8'))

        command = os.environ['MAILER']+' "'+filepath+'"'
        self.msg('Running: '+command)
        os.system(command)
        self.msg('Removing: '+filepath)
        os.remove(filepath)

<|endoftext|>


<commit_before>"""

This file contains functions for reading in phylotree variants and haplogroups
and building a matrix to represent these markers.

Mon Apr  4 12:00:54 PDT 2016

"""

import sys
import collections

class Phylotree(object):
    """
    Class for representing haplogroups and the variants that define them in
    an explicit tree data structure. The general workflow of this is to load
    in the tree from a file based on the Phylotree mtDNA tree (phylotree.org),
    perform any filtering steps and produce a table of haplogroup IDs and
    the variants that define them.
    """
    class PhyloNode(object):
        """
        Class that represents a single node in the phylotree.
        """
        def __init__(self, hap='', parent=None, variants=None):
            """ Initialize new PhyloNode """
            self.hap_id = hap
            self.parent = parent
            self.children = list()
            self.anon_child = 0
            self.anon = (self.hap_id == '')
            if self.parent is not None:
                parent.children.append(self)
                if self.anon:
                    self.hap_id = parent.get_anon_name()
            if variants is None:
                self.variants = list()
            else:
                self.variants = variants

        def all_variants(self):
            """
            Produces a list of variants that define this haplogroup including
            variants inherited from the node's parent. Importantly, mutations
            that occur to the same position are masked by the most recent
            mutation (i.e. C152T will not be included if T152C occurs farther
            down in the tree).
            """
            summed_vars = dict()
            node = self
            while node is not None:
                for var in node.variants:
                    pos = pos_from_var(var)
                    if pos not in summed_vars:
                        summed_vars[pos] = var
                node = node.parent
            return [summed_vars[pos] for pos in sorted(summed_vars)]

        def dump(self, out=sys.stderr, indent=0):
            """
            Dumps the contents of this node and its children recursively.
            """
            prefix = ' ' * indent
            out.write('%sNode: %s\n' % (prefix, self.hap_id))
            out.write('%sVariants: %s\n' % (prefix, ','.join(self.variants)))
            out.write('%sChildren: %d\n' % (prefix, len(self.children)))
            for child in self.children:
                child.dump(out, indent + 2)

        def get_anon_name(self):
            """
            Returns a unique haplogroup name based on this one for a child with
            no specified name.
            """
            self.anon_child += 1
            return "%s[%d]" % (self.hap_id, self.anon_child)

    def __init__(self, phy_in=None, rm_unstable=False, rm_backmut=False):
        """ Initialize a blank Phylotree before reading from a file. """
        self.root  = None
        self.nodes = list()
        self.variants = collections.defaultdict(collections.Counter)
        self.ignore = None
        self.hap_var  = None
        if phy_in is not None:
            self.read_csv(phy_in)
            self.process_variants(rm_unstable, rm_backmut)
            self.process_haplotypes()

    def read_csv(self, phy_in):
        """ Reads tree structure and variants from input stream. """
        cur = -1
        node_stack = list([None])
        for line in phy_in:
            level, hap_id, raw_var = _read_phy_line(line)
            variants = [var for var in raw_var if is_snp(var)]
            while cur >= 0 and cur >= level:
                node_stack.pop()
                cur -= 1
            new_node = Phylotree.PhyloNode(hap_id, node_stack[-1], variants)
            self.nodes.append(new_node)
            node_stack.append(new_node)
            cur += 1
        else:
            # Set the root for the tree. First item on node_stack is None.
            self.root = node_stack[1]
        return

    def process_variants(self, rm_unstable, rm_backmut):
        """
        Builds a table of variant sites stored in Phylotree.variants after
        filtering based on phylotree annotation. This table keeps track of
        number of mutations and derived alleles as they occur.
        """
        var_pos = set()
        self.ignore = set()
        for node in self.nodes:
            for var in node.variants:
                pos = pos_from_var(var)
                if rm_unstable and is_unstable(var):
                    self.ignore.add(pos)
                elif rm_backmut and is_backmutation(var):
                    self.ignore.add(pos)
                else:
                    var_pos.add(pos)
                    der = der_allele(var)
                    self.variants[pos][der] += 1
        if rm_unstable or rm_backmut:
            var_pos -= self.ignore
            for pos in self.ignore:
                if pos in self.variants:
                    del self.variants[pos]
            for node in self.nodes:
                node.variants = [var for var in node.variants
                                 if pos_from_var(var) not in self.ignore]
        return

    def process_haplotypes(self):
        """
        Contructs a dictionary that maps haplogroup IDs to a list of
        variants associated with that haplogroup, using the Phylotree
        information and variants after any filtering. Merges haplogroups with
        identical variants.
        """
        haplotypes = collections.defaultdict(list)
        hap_tab = dict()
        for node in self.nodes:
            var_str = ','.join(node.all_variants())
            haplotypes[var_str].append(node)
        for var_str in haplotypes:
            variants = var_str.split(',')
            haplo_id = '/'.join([node.hap_id for node in haplotypes[var_str]])
            hap_tab[haplo_id] = variants
        self.hap_var = hap_tab
        return

    def get_variant_pos(self):
        """
        Returns a list of positions with variants after filtering.
        """
        return sorted(self.variants.keys())

    def add_custom_hap(self, hap_id, variants):
        """
        Adds a custom haplotype to the haplotype variant table. Does not add a
        node in the tree for this haplotype. Ignores positions that have been
        removed from consideration by process_variants().
        """
        if hap_id in self.hap_var:
            raise ValueError("Custom haplogroup name '%d' already in use."
                             % (hap_id))
        ok_vars = [var for var in variants
                   if pos_from_var(var) not in self.ignore]
        self.hap_var[hap_id] = ok_vars
        return

def pos_from_var(var):
    """ Returns the position of the variant """
    if var.startswith('('):
        var = var[1:-1]
    var = var.rstrip('!')
    return int(var[1:-1]) - 1 # 1-based to 0-based

def der_allele(var):
    """ Returns the derived allele of this variant. """
    var = var.rstrip(')!')
    return var[-1].upper()

def anc_allele(var):
    """ Returns the derived allele of this variant. """
    var = var.lstrip('(')
    return var[0].upper()

def is_snp(var):
    """ Returns true if var is a SNP or False if it is an indel """
    if '.' in var or 'd' in var:
        return False # Variant is an indel
    return True

def is_unstable(var):
    """ Returns true if var is annotated as unstable. """
    if var[0] == '(':
        return True
    return False

def is_backmutation(var):
    """ Returns true if var is annotated as a backmutation. """
    if '!' in var:
        return True
    return False

def rm_snp_annot(var):
    """
    Returns the SNP variant string, nicely formatted with annotation stripped.
    """
    if var.startsiwth('('):
        var = var[1:-1]
    var.rstrip('!')
    return var.upper()

def _read_phy_line(line):
    """
    Reads a single comma-separated line from phylotree and returns the
    indentation level, the haplogroup id, and variants. Importantly, some nodes
    do not have IDs, so we need to check that we do not accidentally grab the
    variant list instead of the blank id.
    """
    items = line.rstrip().split(',')
    level = 0
    while items[level] == '':
        level += 1
    hap_id = items[level]
    while ' ' in hap_id:
        level -= 1
        hap_id = items[level]
    variants = items[level + 1].split()
    return level, hap_id, variants

def main():
    """ Simple test of phylotree functions. """
    if len(sys.argv) > 1:
        phy_fn = sys.argv[1]
        with open(phy_fn, 'r') as phy_in:
            phy = Phylotree(phy_in)
            #phy.root.dump(sys.stdout)
            #phy.process_variants(rm_unstable=False, rm_backmut=False)
            for pos in sorted(phy.variants):
                print pos, phy.variants[pos], (sum(phy.variants[pos].values()))
#           for hap in phy.hap_var:
#               print hap, ','.join(phy.hap_var[hap])
    return 0

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Added stub for class function in Phylotree for ignoring troublesome sites.<commit_after>"""

This file contains functions for reading in phylotree variants and haplogroups
and building a matrix to represent these markers.

Mon Apr  4 12:00:54 PDT 2016

"""

import sys
import collections

class Phylotree(object):
    """
    Class for representing haplogroups and the variants that define them in
    an explicit tree data structure. The general workflow of this is to load
    in the tree from a file based on the Phylotree mtDNA tree (phylotree.org),
    perform any filtering steps and produce a table of haplogroup IDs and
    the variants that define them.
    """
    class PhyloNode(object):
        """
        Class that represents a single node in the phylotree.
        """
        def __init__(self, hap='', parent=None, variants=None):
            """ Initialize new PhyloNode """
            self.hap_id = hap
            self.parent = parent
            self.children = list()
            self.anon_child = 0
            self.anon = (self.hap_id == '')
            if self.parent is not None:
                parent.children.append(self)
                if self.anon:
                    self.hap_id = parent.get_anon_name()
            if variants is None:
                self.variants = list()
            else:
                self.variants = variants

        def all_variants(self):
            """
            Produces a list of variants that define this haplogroup including
            variants inherited from the node's parent. Importantly, mutations
            that occur to the same position are masked by the most recent
            mutation (i.e. C152T will not be included if T152C occurs farther
            down in the tree).
            """
            summed_vars = dict()
            node = self
            while node is not None:
                for var in node.variants:
                    pos = pos_from_var(var)
                    if pos not in summed_vars:
                        summed_vars[pos] = var
                node = node.parent
            return [summed_vars[pos] for pos in sorted(summed_vars)]

        def dump(self, out=sys.stderr, indent=0):
            """
            Dumps the contents of this node and its children recursively.
            """
            prefix = ' ' * indent
            out.write('%sNode: %s\n' % (prefix, self.hap_id))
            out.write('%sVariants: %s\n' % (prefix, ','.join(self.variants)))
            out.write('%sChildren: %d\n' % (prefix, len(self.children)))
            for child in self.children:
                child.dump(out, indent + 2)

        def get_anon_name(self):
            """
            Returns a unique haplogroup name based on this one for a child with
            no specified name.
            """
            self.anon_child += 1
            return "%s[%d]" % (self.hap_id, self.anon_child)

    def __init__(self, phy_in=None, rm_unstable=False, rm_backmut=False):
        """ Initialize a blank Phylotree before reading from a file. """
        self.root  = None
        self.nodes = list()
        self.variants = collections.defaultdict(collections.Counter)
        self.ignore = None
        self.hap_var  = None
        if phy_in is not None:
            self.read_csv(phy_in)
            self.process_variants(rm_unstable, rm_backmut)
            self.process_haplotypes()

    def read_csv(self, phy_in):
        """ Reads tree structure and variants from input stream. """
        cur = -1
        node_stack = list([None])
        for line in phy_in:
            level, hap_id, raw_var = _read_phy_line(line)
            variants = [var for var in raw_var if is_snp(var)]
            while cur >= 0 and cur >= level:
                node_stack.pop()
                cur -= 1
            new_node = Phylotree.PhyloNode(hap_id, node_stack[-1], variants)
            self.nodes.append(new_node)
            node_stack.append(new_node)
            cur += 1
        else:
            # Set the root for the tree. First item on node_stack is None.
            self.root = node_stack[1]
        return

    def process_variants(self, rm_unstable, rm_backmut):
        """
        Builds a table of variant sites stored in Phylotree.variants after
        filtering based on phylotree annotation. This table keeps track of
        number of mutations and derived alleles as they occur.
        """
        var_pos = set()
        self.ignore = set()
        for node in self.nodes:
            for var in node.variants:
                pos = pos_from_var(var)
                if rm_unstable and is_unstable(var):
                    self.ignore.add(pos)
                elif rm_backmut and is_backmutation(var):
                    self.ignore.add(pos)
                else:
                    var_pos.add(pos)
                    der = der_allele(var)
                    self.variants[pos][der] += 1
        if rm_unstable or rm_backmut:
            var_pos -= self.ignore
            for pos in self.ignore:
                if pos in self.variants:
                    del self.variants[pos]
            for node in self.nodes:
                node.variants = [var for var in node.variants
                                 if pos_from_var(var) not in self.ignore]
        return

    def process_haplotypes(self):
        """
        Contructs a dictionary that maps haplogroup IDs to a list of
        variants associated with that haplogroup, using the Phylotree
        information and variants after any filtering. Merges haplogroups with
        identical variants.
        """
        haplotypes = collections.defaultdict(list)
        hap_tab = dict()
        for node in self.nodes:
            var_str = ','.join(node.all_variants())
            haplotypes[var_str].append(node)
        for var_str in haplotypes:
            variants = var_str.split(',')
            haplo_id = '/'.join([node.hap_id for node in haplotypes[var_str]])
            hap_tab[haplo_id] = variants
        self.hap_var = hap_tab
        return

    def get_variant_pos(self):
        """
        Returns a list of positions with variants after filtering.
        """
        return sorted(self.variants.keys())

    def add_custom_hap(self, hap_id, variants):
        """
        Adds a custom haplotype to the haplotype variant table. Does not add a
        node in the tree for this haplotype. Ignores positions that have been
        removed from consideration by process_variants().
        """
        if hap_id in self.hap_var:
            raise ValueError("Custom haplogroup name '%d' already in use."
                             % (hap_id))
        ok_vars = [var for var in variants
                   if pos_from_var(var) not in self.ignore]
        self.hap_var[hap_id] = ok_vars
        return

    def ignore_sites(self, sites_str):
        """
        Adds the sites described by the argument into the ignore list and
        rebuilds the variant and haplotype tables.
        """
        # TODO: Add sites to ignore set.
        self.process_variants(self.rm_unstable, self.rm_backmut)
        self.process_haplotypes()
        return

def pos_from_var(var):
    """ Returns the position of the variant """
    if var.startswith('('):
        var = var[1:-1]
    var = var.rstrip('!')
    return int(var[1:-1]) - 1 # 1-based to 0-based

def der_allele(var):
    """ Returns the derived allele of this variant. """
    var = var.rstrip(')!')
    return var[-1].upper()

def anc_allele(var):
    """ Returns the derived allele of this variant. """
    var = var.lstrip('(')
    return var[0].upper()

def is_snp(var):
    """ Returns true if var is a SNP or False if it is an indel """
    if '.' in var or 'd' in var:
        return False # Variant is an indel
    return True

def is_unstable(var):
    """ Returns true if var is annotated as unstable. """
    if var[0] == '(':
        return True
    return False

def is_backmutation(var):
    """ Returns true if var is annotated as a backmutation. """
    if '!' in var:
        return True
    return False

def rm_snp_annot(var):
    """
    Returns the SNP variant string, nicely formatted with annotation stripped.
    """
    if var.startsiwth('('):
        var = var[1:-1]
    var.rstrip('!')
    return var.upper()

def _read_phy_line(line):
    """
    Reads a single comma-separated line from phylotree and returns the
    indentation level, the haplogroup id, and variants. Importantly, some nodes
    do not have IDs, so we need to check that we do not accidentally grab the
    variant list instead of the blank id.
    """
    items = line.rstrip().split(',')
    level = 0
    while items[level] == '':
        level += 1
    hap_id = items[level]
    while ' ' in hap_id:
        level -= 1
        hap_id = items[level]
    variants = items[level + 1].split()
    return level, hap_id, variants

def main():
    """ Simple test of phylotree functions. """
    if len(sys.argv) > 1:
        phy_fn = sys.argv[1]
        with open(phy_fn, 'r') as phy_in:
            phy = Phylotree(phy_in)
            #phy.root.dump(sys.stdout)
            #phy.process_variants(rm_unstable=False, rm_backmut=False)
            for pos in sorted(phy.variants):
                print pos, phy.variants[pos], (sum(phy.variants[pos].values()))
#           for hap in phy.hap_var:
#               print hap, ','.join(phy.hap_var[hap])
    return 0

if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>