<commit_before>"""\
@file coros.py
@author Donovan Preston

Copyright (c) 2007, Linden Research, Inc.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import collections
import sys
import time
import traceback

from eventlet import api
from eventlet import channel
from eventlet import pools
from eventlet import greenlib

try:
    set
except NameError:  # python 2.3 compatibility
    from sets import Set as set

class Cancelled(RuntimeError):
    pass

NOT_USED = object()

class event(object):
    """An abstraction where an arbitrary number of coroutines
    can wait for one event from another.

    Events differ from channels in two ways:
      1) calling send() does not unschedule the current coroutine
      2) send() can only be called once; use reset() to prepare the event for
         another send()
    They are ideal for communicating return values between coroutines.

    >>> from eventlet import coros, api
    >>> evt = coros.event()
    >>> def baz(b):
    ...     evt.send(b + 1)
    ...
    >>> _ = api.spawn(baz, 3)
    >>> evt.wait()
    4
    """
    _result = None
    def __init__(self):
        self.reset()

    def reset(self):
        """ Reset this event so it can be used to send again.
        Can only be called after send has been called.

        >>> from eventlet import coros
        >>> evt = coros.event()
        >>> evt.send(1)
        >>> evt.reset()
        >>> evt.send(2)
        >>> evt.wait()
        2

        Calling reset multiple times in a row is an error.

        >>> evt.reset()
        >>> evt.reset()
        Traceback (most recent call last):
        ...
        AssertionError: Trying to re-reset() a fresh event.

        """
        assert self._result is not NOT_USED, 'Trying to re-reset() a fresh event.'
        self.epoch = time.time()
        self._result = NOT_USED
        self._waiters = {}

    def wait(self):
        """Wait until another coroutine calls send.
        Returns the value the other coroutine passed to
        send.

        >>> from eventlet import coros, api
        >>> evt = coros.event()
        >>> def wait_on():
        ...    retval = evt.wait()
        ...    print "waited for", retval
        >>> _ = api.spawn(wait_on)
        >>> evt.send('result')
        >>> api.sleep(0)
        waited for result

        Returns immediately if the event has already
        occured.

        >>> evt.wait()
        'result'
        """
        if self._result is NOT_USED:
            self._waiters[api.getcurrent()] = True
            return api.get_hub().switch()
        if self._exc is not None:
            raise self._exc
        return self._result

    def cancel(self, waiter):
        """Raise an exception into a coroutine which called
        wait() an this event instead of returning a value
        from wait. Sends the eventlet.coros.Cancelled
        exception

        waiter: The greenlet (greenlet.getcurrent()) of the
            coroutine to cancel

        >>> from eventlet import coros, api
        >>> evt = coros.event()
        >>> def wait_on():
        ...    try:
        ...        print "received " + evt.wait()
        ...    except coros.Cancelled, c:
        ...        print "Cancelled"
        ...
        >>> waiter = api.spawn(wait_on)

        The cancel call works on coroutines that are in the wait() call.

        >>> api.sleep(0)  # enter the wait()
        >>> evt.cancel(waiter)
        >>> api.sleep(0)  # receive the exception
        Cancelled

        The cancel is invisible to coroutines that call wait() after cancel()
        is called.  This is different from send()'s behavior, where the result
        is passed to any waiter regardless of the ordering of the calls.

        >>> waiter = api.spawn(wait_on)
        >>> api.sleep(0)

        Cancels have no effect on the ability to send() to the event.

        >>> evt.send('stuff')
        >>> api.sleep(0)
        received stuff
        """
        if waiter in self._waiters:
            del self._waiters[waiter]
            api.get_hub().schedule_call(
                0, greenlib.switch, waiter, None, Cancelled())

    def send(self, result=None, exc=None):
        """Makes arrangements for the waiters to be woken with the
        result and then returns immediately to the parent.

        >>> from eventlet import coros, api
        >>> evt = coros.event()
        >>> def waiter():
        ...     print 'about to wait'
        ...     result = evt.wait()
        ...     print 'waited for', result
        >>> _ = api.spawn(waiter)
        >>> api.sleep(0)
        about to wait
        >>> evt.send('a')
        >>> api.sleep(0)
        waited for a

        It is an error to call send() multiple times on the same event.

        >>> evt.send('whoops')
        Traceback (most recent call last):
        ...
        AssertionError: Trying to re-send() an already-triggered event.

        Use reset() between send()s to reuse an event object.
        """
        assert self._result is NOT_USED, 'Trying to re-send() an already-triggered event.'
        self._result = result
        self._exc = exc
        hub = api.get_hub()
        for waiter in self._waiters:
            hub.schedule_call(0, greenlib.switch, waiter, self._result)

def execute(func, *args, **kw):
    """ Executes an operation asynchronously in a new coroutine, returning
    an event to retrieve the return value.

    This has the same api as the CoroutinePool.execute method; the only
    difference is that this one creates a new coroutine instead of drawing
    from a pool.

    >>> from eventlet import coros
    >>> evt = coros.execute(lambda a: ('foo', a), 1)
    >>> evt.wait()
    ('foo', 1)
    """
    evt = event()
    def _really_execute():
        evt.send(func(*args, **kw))
    api.spawn(_really_execute)
    return evt

class CoroutinePool(pools.Pool):
    """ Like a thread pool, but with coroutines.

    Coroutine pools are useful for splitting up tasks or globally controlling
    concurrency.  You don't retrieve the coroutines directly with get() --
    instead use the execute() and execute_async() methods to run code.

    >>> from eventlet import coros, api
    >>> p = coros.CoroutinePool(max_size=2)
    >>> def foo(a):
    ...   print "foo", a
    ...
    >>> evt = p.execute(foo, 1)
    >>> evt.wait()
    foo 1

    Once the pool is exhausted, calling an execute forces a yield.

    >>> p.execute_async(foo, 2)
    >>> p.execute_async(foo, 3)
    >>> p.free()
    0
    >>> p.execute_async(foo, 4)
    foo 2
    foo 3

    >>> api.sleep(0)
    foo 4
    """

    def __init__(self, min_size=0, max_size=4):
        self._greenlets = set()
        super(CoroutinePool, self).__init__(min_size, max_size)

    def _main_loop(self, sender):
        """ Private, infinite loop run by a pooled coroutine. """
        try:
            while True:
                recvd = sender.wait()
                sender = event()
                (evt, func, args, kw) = recvd
                self._safe_apply(evt, func, args, kw)
                api.get_hub().runloop.cancel_timers(api.getcurrent())
                self.put(sender)
        finally:
            # if we get here, something broke badly, and all we can really
            # do is try to keep the pool from leaking items
            self.put(self.create())

    def _safe_apply(self, evt, func, args, kw):
        """ Private method that runs the function, catches exceptions, and
        passes back the return value in the event."""
        try:
            result = func(*args, **kw)
            if evt is not None:
                evt.send(result)
        except api.GreenletExit, e:
            # we're printing this out to see if it ever happens
            # in practice
            print "GreenletExit raised in coroutine pool", e
            if evt is not None:
                evt.send(e)  # sent as a return value, not an exception
        except KeyboardInterrupt:
            raise  # allow program to exit
        except Exception, e:
            traceback.print_exc()
            if evt is not None:
                evt.send(exc=e)

    def _execute(self, evt, func, args, kw):
        """ Private implementation of the execute methods.
        """
        # if reentering an empty pool, don't try to wait on a coroutine freeing
        # itself -- instead, just execute in the current coroutine
        if self.free() == 0 and api.getcurrent() in self._greenlets:
            self._safe_apply(evt, func, args, kw)
        else:
            sender = self.get()
            sender.send((evt, func, args, kw))

    def create(self):
        """Private implementation of eventlet.pools.Pool
        interface. Creates an event and spawns the
        _main_loop coroutine, passing the event.
        The event is used to send a callable into the
        new coroutine, to be executed.
        """
        sender = event()
        self._greenlets.add(api.spawn(self._main_loop, sender))
        return sender

    def execute(self, func, *args, **kw):
        """Execute func in one of the coroutines maintained
        by the pool, when one is free.

        Immediately returns an eventlet.coros.event object which
        func's result will be sent to when it is available.

        >>> from eventlet import coros
        >>> p = coros.CoroutinePool()
        >>> evt = p.execute(lambda a: ('foo', a), 1)
        >>> evt.wait()
        ('foo', 1)
        """
        receiver = event()
        self._execute(receiver, func, args, kw)
        return receiver

    def execute_async(self, func, *args, **kw):
        """Execute func in one of the coroutines maintained
        by the pool, when one is free.

        No return value is provided.
        >>> from eventlet import coros, api
        >>> p = coros.CoroutinePool()
        >>> def foo(a):
        ...   print "foo", a
        ...
        >>> p.execute_async(foo, 1)
        >>> api.sleep(0)
        foo 1
        """
        self._execute(None, func, args, kw)

class pipe(object):
    """ Implementation of pipe using events.  Not tested!  Not used, either."""
    def __init__(self):
        self._event = event()
        self._buffer = ''

    def send(self, txt):
        self._buffer += txt
        evt, self._event = self._event, event()
        evt.send()

    def recv(self, num=16384):
        if not self._buffer:
            self._event.wait()
        if num >= len(self._buffer):
            buf, self._buffer = self._buffer, ''
        else:
            buf, self._buffer = self._buffer[:num], self._buffer[num:]
        return buf

class Actor(object):
    """ A free-running coroutine that accepts and processes messages.

    Kind of the equivalent of an Erlang process, really.  It processes
    a queue of messages in the order that they were sent.  You must
    subclass this and implement your own version of receive().

    The actor's reference count will never drop to zero while the
    coroutine exists; if you lose all references to the actor object
    it will never be freed.
    """
    def __init__(self, concurrency = 1):
        """ Constructs an Actor, kicking off a new coroutine to process the messages.

        The concurrency argument specifies how many messages the actor will try
        to process concurrently.  If it is 1, the actor will process messages
        serially.
        """
        self._mailbox = collections.deque()
        self._event = event()
        self._killer = api.spawn(self.run_forever)
        self._pool = CoroutinePool(min_size=0, max_size=concurrency)

    def run_forever(self):
        """ Loops forever, continually checking the mailbox. """
        while True:
            if not self._mailbox:
                self._event.wait()
                self._event = event()
            else:
                # leave the message in the mailbox until after it's
                # been processed so the event doesn't get triggered
                # while in the received method
                self._pool.execute_async(
                    self.received, self._mailbox[0])
                self._mailbox.popleft()

    def cast(self, message):
        """ Send a message to the actor.

        If the actor is busy, the message will be enqueued for later
        consumption.  There is no return value.

        >>> a = Actor()
        >>> a.received = lambda msg: msg
        >>> a.cast("hello")
        """
        self._mailbox.append(message)
        # if this is the only message, the coro could be waiting
        if len(self._mailbox) == 1:
            self._event.send()

    def received(self, message):
        """ Called to process each incoming message.

        The default implementation just raises an exception, so
        replace it with something useful!

        >>> class Greeter(Actor):
        ...     def received(self, (message, evt) ):
        ...         print "received", message
        ...         if evt: evt.send()
        ...
        >>> a = Greeter()

        This example uses events to synchronize between the actor and the main
        coroutine in a predictable manner, but this kinda defeats the point of
        the Actor, so don't do it in a real application.

        >>> evt = event()
        >>> a.cast( ("message 1", evt) )
        >>> evt.wait()  # force it to run at this exact moment
        received message 1
        >>> evt.reset()
        >>> a.cast( ("message 2", None) )
        >>> a.cast( ("message 3", evt) )
        >>> evt.wait()
        received message 2
        received message 3

        >>> api.kill(a._killer)   # test cleanup
        """
        raise NotImplementedError()

def _test():
    print "Running doctests.  There will be no further output if they succeed."
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    _test()
<commit_msg>[svn r111] add ready() method to event class, so you can tell in advance whether or not wait() is going to block.<commit_after>"""\
@file coros.py
@author Donovan Preston

Copyright (c) 2007, Linden Research, Inc.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import collections
import sys
import time
import traceback

from eventlet import api
from eventlet import channel
from eventlet import pools
from eventlet import greenlib

try:
    set
except NameError:  # python 2.3 compatibility
    from sets import Set as set

class Cancelled(RuntimeError):
    pass

NOT_USED = object()

class event(object):
    """An abstraction where an arbitrary number of coroutines
    can wait for one event from another.

    Events differ from channels in two ways:
      1) calling send() does not unschedule the current coroutine
      2) send() can only be called once; use reset() to prepare the event for
         another send()
    They are ideal for communicating return values between coroutines.

    >>> from eventlet import coros, api
    >>> evt = coros.event()
    >>> def baz(b):
    ...     evt.send(b + 1)
    ...
    >>> _ = api.spawn(baz, 3)
    >>> evt.wait()
    4
    """
    _result = None
    def __init__(self):
        self.reset()

    def reset(self):
        """ Reset this event so it can be used to send again.
        Can only be called after send has been called.

        >>> from eventlet import coros
        >>> evt = coros.event()
        >>> evt.send(1)
        >>> evt.reset()
        >>> evt.send(2)
        >>> evt.wait()
        2

        Calling reset multiple times in a row is an error.

        >>> evt.reset()
        >>> evt.reset()
        Traceback (most recent call last):
        ...
        AssertionError: Trying to re-reset() a fresh event.

        """
        assert self._result is not NOT_USED, 'Trying to re-reset() a fresh event.'
        self.epoch = time.time()
        self._result = NOT_USED
        self._waiters = {}

    def ready(self):
        """ Return true if the wait() call will return immediately.
        Used to avoid waiting for things that might take a while to time out.
        For example, you can put a bunch of events into a list, and then visit
        them all repeatedly, calling ready() until one returns True, and then
        you can wait() on that one."""
        return self._result is not NOT_USED

    def wait(self):
        """Wait until another coroutine calls send.
        Returns the value the other coroutine passed to
        send.

        >>> from eventlet import coros, api
        >>> evt = coros.event()
        >>> def wait_on():
        ...    retval = evt.wait()
        ...    print "waited for", retval
        >>> _ = api.spawn(wait_on)
        >>> evt.send('result')
        >>> api.sleep(0)
        waited for result

        Returns immediately if the event has already
        occured.

        >>> evt.wait()
        'result'
        """
        if self._result is NOT_USED:
            self._waiters[api.getcurrent()] = True
            return api.get_hub().switch()
        if self._exc is not None:
            raise self._exc
        return self._result

    def cancel(self, waiter):
        """Raise an exception into a coroutine which called
        wait() an this event instead of returning a value
        from wait. Sends the eventlet.coros.Cancelled
        exception

        waiter: The greenlet (greenlet.getcurrent()) of the
            coroutine to cancel

        >>> from eventlet import coros, api
        >>> evt = coros.event()
        >>> def wait_on():
        ...    try:
        ...        print "received " + evt.wait()
        ...    except coros.Cancelled, c:
        ...        print "Cancelled"
        ...
        >>> waiter = api.spawn(wait_on)

        The cancel call works on coroutines that are in the wait() call.

        >>> api.sleep(0)  # enter the wait()
        >>> evt.cancel(waiter)
        >>> api.sleep(0)  # receive the exception
        Cancelled

        The cancel is invisible to coroutines that call wait() after cancel()
        is called.  This is different from send()'s behavior, where the result
        is passed to any waiter regardless of the ordering of the calls.

        >>> waiter = api.spawn(wait_on)
        >>> api.sleep(0)

        Cancels have no effect on the ability to send() to the event.

        >>> evt.send('stuff')
        >>> api.sleep(0)
        received stuff
        """
        if waiter in self._waiters:
            del self._waiters[waiter]
            api.get_hub().schedule_call(
                0, greenlib.switch, waiter, None, Cancelled())

    def send(self, result=None, exc=None):
        """Makes arrangements for the waiters to be woken with the
        result and then returns immediately to the parent.

        >>> from eventlet import coros, api
        >>> evt = coros.event()
        >>> def waiter():
        ...     print 'about to wait'
        ...     result = evt.wait()
        ...     print 'waited for', result
        >>> _ = api.spawn(waiter)
        >>> api.sleep(0)
        about to wait
        >>> evt.send('a')
        >>> api.sleep(0)
        waited for a

        It is an error to call send() multiple times on the same event.

        >>> evt.send('whoops')
        Traceback (most recent call last):
        ...
        AssertionError: Trying to re-send() an already-triggered event.

        Use reset() between send()s to reuse an event object.
        """
        assert self._result is NOT_USED, 'Trying to re-send() an already-triggered event.'
        self._result = result
        self._exc = exc
        hub = api.get_hub()
        for waiter in self._waiters:
            hub.schedule_call(0, greenlib.switch, waiter, self._result)

def execute(func, *args, **kw):
    """ Executes an operation asynchronously in a new coroutine, returning
    an event to retrieve the return value.

    This has the same api as the CoroutinePool.execute method; the only
    difference is that this one creates a new coroutine instead of drawing
    from a pool.

    >>> from eventlet import coros
    >>> evt = coros.execute(lambda a: ('foo', a), 1)
    >>> evt.wait()
    ('foo', 1)
    """
    evt = event()
    def _really_execute():
        evt.send(func(*args, **kw))
    api.spawn(_really_execute)
    return evt

class CoroutinePool(pools.Pool):
    """ Like a thread pool, but with coroutines.

    Coroutine pools are useful for splitting up tasks or globally controlling
    concurrency.  You don't retrieve the coroutines directly with get() --
    instead use the execute() and execute_async() methods to run code.

    >>> from eventlet import coros, api
    >>> p = coros.CoroutinePool(max_size=2)
    >>> def foo(a):
    ...   print "foo", a
    ...
    >>> evt = p.execute(foo, 1)
    >>> evt.wait()
    foo 1

    Once the pool is exhausted, calling an execute forces a yield.

    >>> p.execute_async(foo, 2)
    >>> p.execute_async(foo, 3)
    >>> p.free()
    0
    >>> p.execute_async(foo, 4)
    foo 2
    foo 3

    >>> api.sleep(0)
    foo 4
    """

    def __init__(self, min_size=0, max_size=4):
        self._greenlets = set()
        super(CoroutinePool, self).__init__(min_size, max_size)

    def _main_loop(self, sender):
        """ Private, infinite loop run by a pooled coroutine. """
        try:
            while True:
                recvd = sender.wait()
                sender = event()
                (evt, func, args, kw) = recvd
                self._safe_apply(evt, func, args, kw)
                api.get_hub().runloop.cancel_timers(api.getcurrent())
                self.put(sender)
        finally:
            # if we get here, something broke badly, and all we can really
            # do is try to keep the pool from leaking items
            self.put(self.create())

    def _safe_apply(self, evt, func, args, kw):
        """ Private method that runs the function, catches exceptions, and
        passes back the return value in the event."""
        try:
            result = func(*args, **kw)
            if evt is not None:
                evt.send(result)
        except api.GreenletExit, e:
            # we're printing this out to see if it ever happens
            # in practice
            print "GreenletExit raised in coroutine pool", e
            if evt is not None:
                evt.send(e)  # sent as a return value, not an exception
        except KeyboardInterrupt:
            raise  # allow program to exit
        except Exception, e:
            traceback.print_exc()
            if evt is not None:
                evt.send(exc=e)

    def _execute(self, evt, func, args, kw):
        """ Private implementation of the execute methods.
        """
        # if reentering an empty pool, don't try to wait on a coroutine freeing
        # itself -- instead, just execute in the current coroutine
        if self.free() == 0 and api.getcurrent() in self._greenlets:
            self._safe_apply(evt, func, args, kw)
        else:
            sender = self.get()
            sender.send((evt, func, args, kw))

    def create(self):
        """Private implementation of eventlet.pools.Pool
        interface. Creates an event and spawns the
        _main_loop coroutine, passing the event.
        The event is used to send a callable into the
        new coroutine, to be executed.
        """
        sender = event()
        self._greenlets.add(api.spawn(self._main_loop, sender))
        return sender

    def execute(self, func, *args, **kw):
        """Execute func in one of the coroutines maintained
        by the pool, when one is free.

        Immediately returns an eventlet.coros.event object which
        func's result will be sent to when it is available.

        >>> from eventlet import coros
        >>> p = coros.CoroutinePool()
        >>> evt = p.execute(lambda a: ('foo', a), 1)
        >>> evt.wait()
        ('foo', 1)
        """
        receiver = event()
        self._execute(receiver, func, args, kw)
        return receiver

    def execute_async(self, func, *args, **kw):
        """Execute func in one of the coroutines maintained
        by the pool, when one is free.

        No return value is provided.
        >>> from eventlet import coros, api
        >>> p = coros.CoroutinePool()
        >>> def foo(a):
        ...   print "foo", a
        ...
        >>> p.execute_async(foo, 1)
        >>> api.sleep(0)
        foo 1
        """
        self._execute(None, func, args, kw)

class pipe(object):
    """ Implementation of pipe using events.  Not tested!  Not used, either."""
    def __init__(self):
        self._event = event()
        self._buffer = ''

    def send(self, txt):
        self._buffer += txt
        evt, self._event = self._event, event()
        evt.send()

    def recv(self, num=16384):
        if not self._buffer:
            self._event.wait()
        if num >= len(self._buffer):
            buf, self._buffer = self._buffer, ''
        else:
            buf, self._buffer = self._buffer[:num], self._buffer[num:]
        return buf

class Actor(object):
    """ A free-running coroutine that accepts and processes messages.

    Kind of the equivalent of an Erlang process, really.  It processes
    a queue of messages in the order that they were sent.  You must
    subclass this and implement your own version of receive().

    The actor's reference count will never drop to zero while the
    coroutine exists; if you lose all references to the actor object
    it will never be freed.
    """
    def __init__(self, concurrency = 1):
        """ Constructs an Actor, kicking off a new coroutine to process the messages.

        The concurrency argument specifies how many messages the actor will try
        to process concurrently.  If it is 1, the actor will process messages
        serially.
        """
        self._mailbox = collections.deque()
        self._event = event()
        self._killer = api.spawn(self.run_forever)
        self._pool = CoroutinePool(min_size=0, max_size=concurrency)

    def run_forever(self):
        """ Loops forever, continually checking the mailbox. """
        while True:
            if not self._mailbox:
                self._event.wait()
                self._event = event()
            else:
                # leave the message in the mailbox until after it's
                # been processed so the event doesn't get triggered
                # while in the received method
                self._pool.execute_async(
                    self.received, self._mailbox[0])
                self._mailbox.popleft()

    def cast(self, message):
        """ Send a message to the actor.

        If the actor is busy, the message will be enqueued for later
        consumption.  There is no return value.

        >>> a = Actor()
        >>> a.received = lambda msg: msg
        >>> a.cast("hello")
        """
        self._mailbox.append(message)
        # if this is the only message, the coro could be waiting
        if len(self._mailbox) == 1:
            self._event.send()

    def received(self, message):
        """ Called to process each incoming message.

        The default implementation just raises an exception, so
        replace it with something useful!

        >>> class Greeter(Actor):
        ...     def received(self, (message, evt) ):
        ...         print "received", message
        ...         if evt: evt.send()
        ...
        >>> a = Greeter()

        This example uses events to synchronize between the actor and the main
        coroutine in a predictable manner, but this kinda defeats the point of
        the Actor, so don't do it in a real application.

        >>> evt = event()
        >>> a.cast( ("message 1", evt) )
        >>> evt.wait()  # force it to run at this exact moment
        received message 1
        >>> evt.reset()
        >>> a.cast( ("message 2", None) )
        >>> a.cast( ("message 3", evt) )
        >>> evt.wait()
        received message 2
        received message 3

        >>> api.kill(a._killer)   # test cleanup
        """
        raise NotImplementedError()

def _test():
    print "Running doctests.  There will be no further output if they succeed."
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    _test()
<|endoftext|>


<commit_before>import nose.tools
from tests.common import WpsTestClient

def test_caps():
    wps = WpsTestClient()
    resp = wps.get(service='wps', request='getcapabilities')
    names = resp.xpath_text('/wps:Capabilities'
                            '/wps:ProcessOfferings'
                            '/wps:Process'
                            '/ows:Identifier')
    sorted_names = sorted(names.split())
    expected_names = ['WClass', 'analogs', 'ensembleRobustness', 'eobs_to_cordex', 'fetch', 'indices_single', 'sdm', 'segetalflora', 'subset_countries', 'visualisation']
    nose.tools.ok_(sorted_names == expected_names, "Expected: {0}, Got: {1}".format(expected_names, sorted_names))

<commit_msg>added extracpoints to wps_caps test<commit_after>import nose.tools
from tests.common import WpsTestClient

def test_caps():
    wps = WpsTestClient()
    resp = wps.get(service='wps', request='getcapabilities')
    names = resp.xpath_text('/wps:Capabilities'
                            '/wps:ProcessOfferings'
                            '/wps:Process'
                            '/ows:Identifier')
    sorted_names = sorted(names.split())
    expected_names = ['WClass', 'analogs', 'ensembleRobustness', 'eobs_to_cordex', 'extractpoints', 'fetch', 'indices_single', 'sdm', 'segetalflora', 'subset_countries', 'visualisation']
    nose.tools.ok_(sorted_names == expected_names, "Expected: {0}, Got: {1}".format(expected_names, sorted_names))

<|endoftext|>


<commit_before>import os
import subprocess
import threading
import xmlrpc.server as rpc
import time
import sys
import multiprocessing as mp

try:
    import pydivert
except ImportError:
    pass

class UnixDelayHandler:

    @staticmethod
    def get_interface_name():
        """ Get the first interface name that is not the localhost. """
        net = os.listdir("/sys/class/net")
        net.remove("lo")
        return net[0]

    def _set_delay(self, sec, duration=10.0):
        thr = threading.Timer(duration, self._remove_delay)
        thr.start()
        return subprocess.call(["tc", "qdisc", "add", "dev",
                                self.get_interface_name(), "root",
                                "handle", "1:0", "netem", "delay",
                                ("%dmsec" % (sec * 1000))])

    def set_delay(self, sec, duration=10.0):
        """ Set network delay, return with the call's result. """
        thr = threading.Timer(2.0, self._set_delay, (sec, duration))
        thr.start()
        return True

    def _remove_delay(self):
        return subprocess.call(["tc", "qdisc", "del", "dev",
                                self.get_interface_name(), "root"])

    def remove_delay(self):
        """ Remove network delay, return with the call's result. """
        thr = threading.Timer(1.0, self._remove_delay)
        thr.start()
        return True

class WinDelayHandler:
    proc = None

    def delay(self, sec, duration=10.0):
        netfil = "tcp.DstPort == 389 or tcp.SrcPort == 389"
        start = time.time()
        with pydivert.WinDivert(netfil) as divert:
            for packet in divert:
                time.sleep(sec)
                divert.send(packet)
                if time.time() - start >= duration:
                    break

    def set_delay(self, sec, duration=10.0):
        """ Set network delay, return with the call's result. """
        WinDelayHandler.stop = False
        self.proc = mp.Process(target=self.delay, args=(sec, duration))
        self.proc.start()
        return True

    def remove_delay(self):
        """ Remove network delay, return with the call's result. """
        if self.proc is not None and self.proc.is_alive():
            self.proc.terminate()
        return True

if __name__ == "__main__":
    if sys.platform == "win32":
        handler = WinDelayHandler()
    else:
        handler = UnixDelayHandler()
    server = rpc.SimpleXMLRPCServer(("0.0.0.0", 8000))
    server.register_function(handler.set_delay, "set_delay")
    server.register_function(handler.remove_delay, "remove_delay")
    server.serve_forever()<commit_msg>Add MacDelayHandler, rename UnixDelayHandler<commit_after>import os
import subprocess
import threading
import xmlrpc.server as rpc
import time
import sys
import multiprocessing as mp

try:
    import pydivert
except ImportError:
    pass

class LinuxDelayHandler:

    @staticmethod
    def get_interface_name():
        """ Get the first interface name that is not the localhost. """
        net = os.listdir("/sys/class/net")
        net.remove("lo")
        return net[0]

    def _set_delay(self, sec, duration=10.0):
        thr = threading.Timer(duration, self._remove_delay)
        thr.start()
        return subprocess.call(["tc", "qdisc", "add", "dev",
                                self.get_interface_name(), "root",
                                "handle", "1:0", "netem", "delay",
                                ("%dmsec" % (sec * 1000))])

    def set_delay(self, sec, duration=10.0):
        """ Set network delay, return with the call's result. """
        thr = threading.Timer(2.0, self._set_delay, (sec, duration))
        thr.start()
        return True

    def _remove_delay(self):
        return subprocess.call(["tc", "qdisc", "del", "dev",
                                self.get_interface_name(), "root"])

    def remove_delay(self):
        """ Remove network delay, return with the call's result. """
        thr = threading.Timer(1.0, self._remove_delay)
        thr.start()
        return True

class MacDelayHandler:
    def set_delay(self, sec, duration=10.0):
        with open('/etc/pf.conf') as fp:
            conf = fp.read()
            conf += '\ndummynet-anchor "mop"\nanchor "mop"\n'
            rule = "dummynet in quick proto tcp from any to any port {389, 636} pipe 1\n"

            try:
                subprocess.run(["pfctl", "-f", "-"], input=conf, encoding="utf-8", check=True)
                subprocess.run(["pfctl", "-a", "mop", "-f", "-"], input=rule,
                               encoding="utf-8", check=True)
                subprocess.check_call(["dnctl", "pipe", "1", "config", "delay",
                                       "%d" % int(sec * 1000)])

                thr = threading.Timer(duration, self.remove_delay)
                thr.start()

                return True
            except subprocess.CalledProcessError:
                return False

    def remove_delay(self):
        try:
            subprocess.check_call(["dnctl", "-q", "flush"])
            subprocess.check_call(["pfctl", "-f", "/etc/pf.conf"])
            return True
        except subprocess.CalledProcessError:
            return False

class WinDelayHandler:
    proc = None

    def delay(self, sec, duration=10.0):
        netfil = "tcp.DstPort == 389 or tcp.SrcPort == 389"
        start = time.time()
        with pydivert.WinDivert(netfil) as divert:
            for packet in divert:
                time.sleep(sec)
                divert.send(packet)
                if time.time() - start >= duration:
                    break

    def set_delay(self, sec, duration=10.0):
        """ Set network delay, return with the call's result. """
        WinDelayHandler.stop = False
        self.proc = mp.Process(target=self.delay, args=(sec, duration))
        self.proc.start()
        return True

    def remove_delay(self):
        """ Remove network delay, return with the call's result. """
        if self.proc is not None and self.proc.is_alive():
            self.proc.terminate()
        return True

if __name__ == "__main__":
    if sys.platform == "win32":
        handler = WinDelayHandler()
    elif sys.platform == "darwin":
        handler = MacDelayHandler()
    else:
        handler = LinuxDelayHandler()
    server = rpc.SimpleXMLRPCServer(("0.0.0.0", 8000))
    server.register_function(handler.set_delay, "set_delay")
    server.register_function(handler.remove_delay, "remove_delay")
    server.serve_forever()<|endoftext|>


<commit_before>import logging
from sys import version_info
from functools import wraps
from threading import Event

from enum import IntEnum

from ..watchmen import watch

def action_valid_decorator_factory(state, protocol):
    """
    This decorator factory  is used to generate decorators which  verify that
    requested actions on a target, such as step(), stop(), read_register(),
    write_register() and so on are actually executable.

    :param state: The required state of the Target
    :type state:  An entry of the Enum TargetStates
    :param protocol: The protocol required to execute the action.
    :type protocol: str
    """

    def decorator(func):
        @wraps(func)
        def check(self, *args, **kwargs):
            if getattr(self.protocols, protocol) is None:
                raise Exception(
                    "%s() requested but %s is undefined." %
                    (func.__name__, protocol))
            if not self.state.value & state.value:
                raise Exception("%s() requested but Target is %s" %
                                (func.__name__, TargetStates(self.state).name))
            return func(self, *args, **kwargs)

        return check

    return decorator

def synchronize_state(*states, **kwargs):
    """
    This decorator can be used to make sure that the target executed a desired
    set of state transitions in an particular order.
    This is useful, when the user explicitly requests the target to change
    it's state and need an update notification on the transition itself.
    Internally, this works by creating an event and using a watchmen to check
    whether it was triggered.
    :param *states: The desired states of the target
    :param transition_optional: Also allow to return if the target is already
                                in the desired states, even if the transition
                                didn't happen
    """
    def decorator(func):
        @wraps(func)
        def state_synchronizer(self, *args, **kwargs):
            state = states[-1]
            transition_optional = kwargs.get('transition_optional', False)

            blocking = kwargs.get('blocking', True)
            avatar = self.avatar
            if blocking is True:
                state_reached = Event()

                def state_synchronize_cb(avatar, message, *args, **kwargs):
                    if message.origin == self:
                        if message.state == state:
                            state_reached.set()
                        elif message.state == TargetStates.EXITED:
                            raise Exception("Target %s exited" % self.name)

                w = avatar.watchmen.add('UpdateState', when='after',
                                        callback=state_synchronize_cb)
            if len(states) == 1:
                ret = func(self, *args, **kwargs)
            else:
                ret = synchronize_state(*states[:-1])(func)(self, *args, **kwargs)
            if blocking is True:
                if not (transition_optional == True and self.state == state):
                    state_reached.wait()
                avatar.watchmen.remove_watchman('UpdateState', w)
            return ret

        return state_synchronizer

    return decorator

class TargetStates(IntEnum):
    """
    A simple Enum for the different states a target can be in.
    """
    CREATED = 0x1
    INITIALIZED = 0x2
    STOPPED = 0x4
    RUNNING = 0x8
    SYNCING = 0x10
    EXITED = 0x20
    NOT_RUNNING = INITIALIZED | STOPPED

class TargetRegs(object):
    def __init__(self, target, register_dict):
        self._target = target
        self.__dict__.update(register_dict)

    def __getattribute__(self, name):
        if name == '_get_names' or name == '__dict__':
            return super(TargetRegs, self).__getattribute__(name)
        elif name in self._get_names():
            return self._target.read_register(name)
        else:
            return super(TargetRegs, self).__getattribute__(name)

    def __setattr__(self, name, value):
        if name == '_target':
            return super(TargetRegs, self).__setattr__(name, value)
        elif name in self._get_names():
            return self._target.write_register(name, value)
        else:
            return super(TargetRegs, self).__setattr__(name, value)

    def _update(self, reg_dict):
        self.__dict__.update(reg_dict)

    def _get_nr_from_name(self, reg_name):
        return self.__dict__[reg_name]

    def _get_names(self):
        names = set(self.__dict__) ^ set(['_target'])
        return names

class TargetProtocolStore(object):
    """This class stores the various protocols associated to one target"""

    DEFAULT_PROTOCOLS = ['memory', 'registers', 'execution']

    def __init__(self, additional_protocols=None):
        self.protocols = set(TargetProtocolStore.DEFAULT_PROTOCOLS)
        self.protocols |= additional_protocols if additional_protocols else set()
        self.unique_protocols = {} #Stores protocol references and their count
        for p in self.protocols:
            setattr(self, p, None)

    def set_all(self, instance, only_defaults=False):
        """
        Sets an instantiated protocol object for either all protocols in this
        store, or only the default ones
        :param instance: the protocol instance
        """
        protocols = (TargetProtocolStore.DEFAULT_PROTOCOLS if only_defaults
                     else self.protocols
                    )
        for p in protocols:
            setattr(self, p, instance)

    def shutdown(self):
        """Shutsdown all the associated protocols"""
        for p in self.protocols:
            setattr(self, p, None)

    def __setattr__(self, name, value):
        if name == 'protocols' or name == 'unique_protocols':
            return super(TargetProtocolStore, self).__setattr__(name, value)

        # Check whether the protocol is already an attribute
        if hasattr(self, name) is False:
            self.protocols.add(name)
            saved_val = None
        else:
            saved_val = getattr(self, name)

        if value is not None and self.unique_protocols.get(value, None) is None:
            self.unique_protocols[value] = 0

        if value is None and saved_val is not None:
            self.unique_protocols[saved_val] -= 1
        elif value is not None and saved_val is None:
            self.unique_protocols[value] += 1
        elif value is not None and saved_val is not None:
            self.unique_protocols[value] += 1
            self.unique_protocols[saved_val] -= 1

        # if there is no reference left, let's shut the prot down
        if saved_val is not None and self.unique_protocols[saved_val] == 0:
            getattr(self, name).shutdown()

        return super(TargetProtocolStore, self).__setattr__(name, value)

class Target(object):
    """The Target object is one of Avatars core concept, as Avatar orchestrate
    different targets.
    While the generic target has no implementation, it provides an insight over
    all the functions a Target MUST implement
    """

    def __init__(self, avatar, name=None):  # type: ('Avatar', str) -> None
        """
        Creates a new instance of a Target.
        :param avatar: The avatar instance this target should be orchestrated by
        :param name: The name of this target, mainly for logging. This is optional and will be autogenerated otherwise.
        """
        super(Target, self).__init__()
        self.state = TargetStates.CREATED

        self.avatar = avatar
        self.name = name if name else self._get_unique_name()

        self.status = {}
        self._arch = avatar.arch
        self.protocols = TargetProtocolStore()

        self.state = TargetStates.CREATED
        self._no_state_update_pending = Event()

        self.log = logging.getLogger('%s.targets.%s' % (avatar.log.name, self.name))
        log_file = logging.FileHandler('%s/%s.log' % (avatar.output_directory, self.name))
        formatter = logging.Formatter('%(asctime)s | %(name)s.%(levelname)s | %(message)s')
        log_file.setFormatter(formatter)
        self.log.addHandler(log_file)

        self.regs = TargetRegs(self, self._arch.registers)

    def dictify(self):
        """
        Returns the memory range as *printable* dictionary for the config
        """

        ignore = ['state', 'status', 'regs', 'protocols', 'log', 'avatar']
        expected_types = (str, bool, int, list)
        if version_info < (3, 0): expected_types += (unicode, )

        t_dict = {'type': self.__class__.__name__,
                  'module': self.__module__}

        for k, v in self.__dict__.items():
            if k in ignore: continue
            if k.startswith('_'): continue
            if v is None: continue
            if not isinstance(v, expected_types):
                raise Exception(
                    "Unsupported type %s for dictifying %s for target %s" %
                    (type(v), k, self.name))
            t_dict[k] = v
        return t_dict

    @watch('TargetInit')
    def init(self):
        """
        Initializes the target to start the analyses
        """
        pass

    @watch('TargetShutdown')
    def shutdown(self):
        """
        Shutdowns the target
        """
        self.protocols.shutdown()

    @watch('TargetCont')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    @synchronize_state(TargetStates.RUNNING)
    def cont(self, blocking=True):
        """
        Continues the execution of the target
        :param blocking: if True, block until the target is RUNNING
        """
        return self.protocols.execution.cont()

    @watch('TargetStop')
    @action_valid_decorator_factory(TargetStates.RUNNING, 'execution')
    @synchronize_state(TargetStates.STOPPED, transition_optional=True)
    def stop(self, blocking=True):
        """
        Stops the execution of the target
        :param blocking: if True, block until the target is STOPPED
        """
        return self.protocols.execution.stop()

    @watch('TargetStep')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    @synchronize_state(TargetStates.RUNNING, TargetStates.STOPPED)
    def step(self, blocking=True):
        """
        Steps one instruction.
        :param blocking: if True, block until the target is STOPPED again
        """
        return self.protocols.execution.step()

    @watch('TargetSetFile')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    def set_file(self, elf):
        """
        Load an ELF file

        :param elf: ELF file to load
        :returns: True on success else False
        """
        if not hasattr(self.protocols.execution, 'set_file'):
            self.log.error('Protocol "' + type(self.protocols.execution).__name__ + '" does not support "set_file"')
            return False

        return self.protocols.execution.set_file(elf)

    @watch('TargetDownload')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    def download(self):
        """
        Download the loaded code to the Target

        :returns: True on success else False
        """
        if not hasattr(self.protocols.execution, 'download'):
            self.log.error('Protocol "' + type(self.protocols.execution).__name__ + '" does not support "download"')
            return False

        return self.protocols.execution.download()

    @watch('TargetGetSymbol')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'memory')
    def get_symbol(self, symbol):
        """
        Get the address of a symbol

        :param symbol:    The name of a symbol whose address is wanted
        :returns:         (True, Address) on success else False
        """
        return self.protocols.memory.get_symbol(symbol)

    @watch('TargetWriteMemory')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'memory')
    def write_memory(self, address, size, value, num_words=1, raw=False):
        """
        Writing to memory of the target

        :param address:   The address from where the memory-write should
                          start
        :param size:      The size of the memory write
        :param value:     The actual value written to memory
        :type val:        int if num_words == 1 and raw == False
                          list if num_words > 1 and raw == False
                          str or byte if raw == True
        :param num_words: The amount of words to read
        :param raw:       Specifies whether to write in raw or word mode
        :returns:         True on success else False
        """
        target_range = self.avatar.get_memory_range(address)
        if target_range.forwarded is True and target_range.forwarded_to != self:
            return target_range.forwarded_to.write_memory(address, size, value,
                                                          num_words, raw)

        return self.protocols.memory.write_memory(address, size, value,
                                                  num_words, raw)

    @watch('TargetReadMemory')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'memory')
    def read_memory(self, address, size, num_words=1, raw=False):
        """
        Reading from memory of the target

        :param address:     The address to read from
        :param size:        The size of a read word
        :param words:       The amount of words to read (default: 1)
        :param raw:         Whether the read memory is returned unprocessed
        :return:          The read memory
        """
        target_range = self.avatar.get_memory_range(address)
        if target_range.forwarded is True and target_range.forwarded_to != self:
            return target_range.forwarded_to.read_memory(address, size,
                                                         num_words, raw)

        return self.protocols.memory.read_memory(address, size, words, raw)

    @watch('TargetRegisterWrite')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'registers')
    def write_register(self, register, value):
        """
        Writing a register to the target

        :param register:     The name of the register
        :param value:        The actual value written to the register
        """
        return self.protocols.registers.write_register(register, value)

    @watch('TargetRegisterRead')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'registers')
    def read_register(self, register):
        """
        Reading a register from the target

        :param register:     The name of the register
        :return:             The actual value read from the register
        """
        return self.protocols.registers.read_register(register)

    @watch('TargetSetBreakpoint')
    @action_valid_decorator_factory(TargetStates.NOT_RUNNING, 'execution')
    def set_breakpoint(self, line, hardware=False, temporary=False, regex=False,
                       condition=None, ignore_count=0, thread=0, **kwargs):
        """Inserts a breakpoint

        :param bool hardware: Hardware breakpoint
        :param bool tempory:  Tempory breakpoint
        :param str regex:     If set, inserts breakpoints matching the regex
        :param str condition: If set, inserts a breakpoint with the condition
        :param int ignore_count: Amount of times the bp should be ignored
        :param int thread:    Threadno in which this breakpoints should be added
        """
        return self.protocols.execution.set_breakpoint(line, hardware=hardware,
                                                  temporary=temporary,
                                                  regex=regex,
                                                  condition=condition,
                                                  ignore_count=ignore_count,
                                                  thread=thread, **kwargs)

    @watch('TargetSetWatchPoint')
    @action_valid_decorator_factory(TargetStates.NOT_RUNNING, 'execution')
    def set_watchpoint(self, variable, write=True, read=False):
        """Inserts a watchpoint

        :param      variable: The name of a variable or an address to watch
        :param bool write:    Write watchpoint
        :param bool read:     Read watchpoint
        """
        return self.protocols.execution.set_watchpoint(variable,
                                                  write=write,
                                                  read=read)

    @watch('TargetRemovebreakpoint')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    def remove_breakpoint(self, bkptno):
        """Deletes a breakpoint"""
        return self.protocols.execution.remove_breakpoint(bkptno)

    def update_state(self, state):
        self.log.info("State changed to to %s" % TargetStates(state))
        self.state = state
        #self._no_state_update_pending.set()

    @watch('TargetWait')
    def wait(self, state=TargetStates.STOPPED):
        while self.state & state == 0:
            pass

    def get_status(self):
        """
        Returns useful information about the target as a dict.
        """
        self.status['state'] = self.state
        return self.status

    def _get_unique_name(self, i=0):
        classname = type(self).__name__
        targetname = "{}{}".format(classname, i)
        if self.avatar and self.avatar.targets and targetname in self.avatar.targets:
            return self._get_unique_name(i + 1)
        return targetname

    def _resolve_executable_name(self):
        """
        Resolves the name of the executable for the endpoint.
        Order of operation:
            1: Check if config exists and whether target is installed
            2: Check sys_name from default config
            3: Check apt_name from default config
            4: BailOut
        """
        pass

    # ##generic aliases##
    wr = write_register
    rr = read_register
    rm = read_memory
    wm = write_memory
    bp = set_breakpoint
<commit_msg>bugfix in read_memory<commit_after>import logging
from sys import version_info
from functools import wraps
from threading import Event

from enum import IntEnum

from ..watchmen import watch

def action_valid_decorator_factory(state, protocol):
    """
    This decorator factory  is used to generate decorators which  verify that
    requested actions on a target, such as step(), stop(), read_register(),
    write_register() and so on are actually executable.

    :param state: The required state of the Target
    :type state:  An entry of the Enum TargetStates
    :param protocol: The protocol required to execute the action.
    :type protocol: str
    """

    def decorator(func):
        @wraps(func)
        def check(self, *args, **kwargs):
            if getattr(self.protocols, protocol) is None:
                raise Exception(
                    "%s() requested but %s is undefined." %
                    (func.__name__, protocol))
            if not self.state.value & state.value:
                raise Exception("%s() requested but Target is %s" %
                                (func.__name__, TargetStates(self.state).name))
            return func(self, *args, **kwargs)

        return check

    return decorator

def synchronize_state(*states, **kwargs):
    """
    This decorator can be used to make sure that the target executed a desired
    set of state transitions in an particular order.
    This is useful, when the user explicitly requests the target to change
    it's state and need an update notification on the transition itself.
    Internally, this works by creating an event and using a watchmen to check
    whether it was triggered.
    :param *states: The desired states of the target
    :param transition_optional: Also allow to return if the target is already
                                in the desired states, even if the transition
                                didn't happen
    """
    def decorator(func):
        @wraps(func)
        def state_synchronizer(self, *args, **kwargs):
            state = states[-1]
            transition_optional = kwargs.get('transition_optional', False)

            blocking = kwargs.get('blocking', True)
            avatar = self.avatar
            if blocking is True:
                state_reached = Event()

                def state_synchronize_cb(avatar, message, *args, **kwargs):
                    if message.origin == self:
                        if message.state == state:
                            state_reached.set()
                        elif message.state == TargetStates.EXITED:
                            raise Exception("Target %s exited" % self.name)

                w = avatar.watchmen.add('UpdateState', when='after',
                                        callback=state_synchronize_cb)
            if len(states) == 1:
                ret = func(self, *args, **kwargs)
            else:
                ret = synchronize_state(*states[:-1])(func)(self, *args, **kwargs)
            if blocking is True:
                if not (transition_optional == True and self.state == state):
                    state_reached.wait()
                avatar.watchmen.remove_watchman('UpdateState', w)
            return ret

        return state_synchronizer

    return decorator

class TargetStates(IntEnum):
    """
    A simple Enum for the different states a target can be in.
    """
    CREATED = 0x1
    INITIALIZED = 0x2
    STOPPED = 0x4
    RUNNING = 0x8
    SYNCING = 0x10
    EXITED = 0x20
    NOT_RUNNING = INITIALIZED | STOPPED

class TargetRegs(object):
    def __init__(self, target, register_dict):
        self._target = target
        self.__dict__.update(register_dict)

    def __getattribute__(self, name):
        if name == '_get_names' or name == '__dict__':
            return super(TargetRegs, self).__getattribute__(name)
        elif name in self._get_names():
            return self._target.read_register(name)
        else:
            return super(TargetRegs, self).__getattribute__(name)

    def __setattr__(self, name, value):
        if name == '_target':
            return super(TargetRegs, self).__setattr__(name, value)
        elif name in self._get_names():
            return self._target.write_register(name, value)
        else:
            return super(TargetRegs, self).__setattr__(name, value)

    def _update(self, reg_dict):
        self.__dict__.update(reg_dict)

    def _get_nr_from_name(self, reg_name):
        return self.__dict__[reg_name]

    def _get_names(self):
        names = set(self.__dict__) ^ set(['_target'])
        return names

class TargetProtocolStore(object):
    """This class stores the various protocols associated to one target"""

    DEFAULT_PROTOCOLS = ['memory', 'registers', 'execution']

    def __init__(self, additional_protocols=None):
        self.protocols = set(TargetProtocolStore.DEFAULT_PROTOCOLS)
        self.protocols |= additional_protocols if additional_protocols else set()
        self.unique_protocols = {} #Stores protocol references and their count
        for p in self.protocols:
            setattr(self, p, None)

    def set_all(self, instance, only_defaults=False):
        """
        Sets an instantiated protocol object for either all protocols in this
        store, or only the default ones
        :param instance: the protocol instance
        """
        protocols = (TargetProtocolStore.DEFAULT_PROTOCOLS if only_defaults
                     else self.protocols
                    )
        for p in protocols:
            setattr(self, p, instance)

    def shutdown(self):
        """Shutsdown all the associated protocols"""
        for p in self.protocols:
            setattr(self, p, None)

    def __setattr__(self, name, value):
        if name == 'protocols' or name == 'unique_protocols':
            return super(TargetProtocolStore, self).__setattr__(name, value)

        # Check whether the protocol is already an attribute
        if hasattr(self, name) is False:
            self.protocols.add(name)
            saved_val = None
        else:
            saved_val = getattr(self, name)

        if value is not None and self.unique_protocols.get(value, None) is None:
            self.unique_protocols[value] = 0

        if value is None and saved_val is not None:
            self.unique_protocols[saved_val] -= 1
        elif value is not None and saved_val is None:
            self.unique_protocols[value] += 1
        elif value is not None and saved_val is not None:
            self.unique_protocols[value] += 1
            self.unique_protocols[saved_val] -= 1

        # if there is no reference left, let's shut the prot down
        if saved_val is not None and self.unique_protocols[saved_val] == 0:
            getattr(self, name).shutdown()

        return super(TargetProtocolStore, self).__setattr__(name, value)

class Target(object):
    """The Target object is one of Avatars core concept, as Avatar orchestrate
    different targets.
    While the generic target has no implementation, it provides an insight over
    all the functions a Target MUST implement
    """

    def __init__(self, avatar, name=None):  # type: ('Avatar', str) -> None
        """
        Creates a new instance of a Target.
        :param avatar: The avatar instance this target should be orchestrated by
        :param name: The name of this target, mainly for logging. This is optional and will be autogenerated otherwise.
        """
        super(Target, self).__init__()
        self.state = TargetStates.CREATED

        self.avatar = avatar
        self.name = name if name else self._get_unique_name()

        self.status = {}
        self._arch = avatar.arch
        self.protocols = TargetProtocolStore()

        self.state = TargetStates.CREATED
        self._no_state_update_pending = Event()

        self.log = logging.getLogger('%s.targets.%s' % (avatar.log.name, self.name))
        log_file = logging.FileHandler('%s/%s.log' % (avatar.output_directory, self.name))
        formatter = logging.Formatter('%(asctime)s | %(name)s.%(levelname)s | %(message)s')
        log_file.setFormatter(formatter)
        self.log.addHandler(log_file)

        self.regs = TargetRegs(self, self._arch.registers)

    def dictify(self):
        """
        Returns the memory range as *printable* dictionary for the config
        """

        ignore = ['state', 'status', 'regs', 'protocols', 'log', 'avatar']
        expected_types = (str, bool, int, list)
        if version_info < (3, 0): expected_types += (unicode, )

        t_dict = {'type': self.__class__.__name__,
                  'module': self.__module__}

        for k, v in self.__dict__.items():
            if k in ignore: continue
            if k.startswith('_'): continue
            if v is None: continue
            if not isinstance(v, expected_types):
                raise Exception(
                    "Unsupported type %s for dictifying %s for target %s" %
                    (type(v), k, self.name))
            t_dict[k] = v
        return t_dict

    @watch('TargetInit')
    def init(self):
        """
        Initializes the target to start the analyses
        """
        pass

    @watch('TargetShutdown')
    def shutdown(self):
        """
        Shutdowns the target
        """
        self.protocols.shutdown()

    @watch('TargetCont')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    @synchronize_state(TargetStates.RUNNING)
    def cont(self, blocking=True):
        """
        Continues the execution of the target
        :param blocking: if True, block until the target is RUNNING
        """
        return self.protocols.execution.cont()

    @watch('TargetStop')
    @action_valid_decorator_factory(TargetStates.RUNNING, 'execution')
    @synchronize_state(TargetStates.STOPPED, transition_optional=True)
    def stop(self, blocking=True):
        """
        Stops the execution of the target
        :param blocking: if True, block until the target is STOPPED
        """
        return self.protocols.execution.stop()

    @watch('TargetStep')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    @synchronize_state(TargetStates.RUNNING, TargetStates.STOPPED)
    def step(self, blocking=True):
        """
        Steps one instruction.
        :param blocking: if True, block until the target is STOPPED again
        """
        return self.protocols.execution.step()

    @watch('TargetSetFile')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    def set_file(self, elf):
        """
        Load an ELF file

        :param elf: ELF file to load
        :returns: True on success else False
        """
        if not hasattr(self.protocols.execution, 'set_file'):
            self.log.error('Protocol "' + type(self.protocols.execution).__name__ + '" does not support "set_file"')
            return False

        return self.protocols.execution.set_file(elf)

    @watch('TargetDownload')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    def download(self):
        """
        Download the loaded code to the Target

        :returns: True on success else False
        """
        if not hasattr(self.protocols.execution, 'download'):
            self.log.error('Protocol "' + type(self.protocols.execution).__name__ + '" does not support "download"')
            return False

        return self.protocols.execution.download()

    @watch('TargetGetSymbol')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'memory')
    def get_symbol(self, symbol):
        """
        Get the address of a symbol

        :param symbol:    The name of a symbol whose address is wanted
        :returns:         (True, Address) on success else False
        """
        return self.protocols.memory.get_symbol(symbol)

    @watch('TargetWriteMemory')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'memory')
    def write_memory(self, address, size, value, num_words=1, raw=False):
        """
        Writing to memory of the target

        :param address:   The address from where the memory-write should
                          start
        :param size:      The size of the memory write
        :param value:     The actual value written to memory
        :type val:        int if num_words == 1 and raw == False
                          list if num_words > 1 and raw == False
                          str or byte if raw == True
        :param num_words: The amount of words to read
        :param raw:       Specifies whether to write in raw or word mode
        :returns:         True on success else False
        """
        target_range = self.avatar.get_memory_range(address)
        if target_range.forwarded is True and target_range.forwarded_to != self:
            return target_range.forwarded_to.write_memory(address, size, value,
                                                          num_words, raw)

        return self.protocols.memory.write_memory(address, size, value,
                                                  num_words, raw)

    @watch('TargetReadMemory')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'memory')
    def read_memory(self, address, size, num_words=1, raw=False):
        """
        Reading from memory of the target

        :param address:     The address to read from
        :param size:        The size of a read word
        :param words:       The amount of words to read (default: 1)
        :param raw:         Whether the read memory is returned unprocessed
        :return:          The read memory
        """
        target_range = self.avatar.get_memory_range(address)
        if target_range.forwarded is True and target_range.forwarded_to != self:
            return target_range.forwarded_to.read_memory(address, size,
                                                         num_words, raw)

        return self.protocols.memory.read_memory(address, size, num_words, raw)

    @watch('TargetRegisterWrite')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'registers')
    def write_register(self, register, value):
        """
        Writing a register to the target

        :param register:     The name of the register
        :param value:        The actual value written to the register
        """
        return self.protocols.registers.write_register(register, value)

    @watch('TargetRegisterRead')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'registers')
    def read_register(self, register):
        """
        Reading a register from the target

        :param register:     The name of the register
        :return:             The actual value read from the register
        """
        return self.protocols.registers.read_register(register)

    @watch('TargetSetBreakpoint')
    @action_valid_decorator_factory(TargetStates.NOT_RUNNING, 'execution')
    def set_breakpoint(self, line, hardware=False, temporary=False, regex=False,
                       condition=None, ignore_count=0, thread=0, **kwargs):
        """Inserts a breakpoint

        :param bool hardware: Hardware breakpoint
        :param bool tempory:  Tempory breakpoint
        :param str regex:     If set, inserts breakpoints matching the regex
        :param str condition: If set, inserts a breakpoint with the condition
        :param int ignore_count: Amount of times the bp should be ignored
        :param int thread:    Threadno in which this breakpoints should be added
        """
        return self.protocols.execution.set_breakpoint(line, hardware=hardware,
                                                  temporary=temporary,
                                                  regex=regex,
                                                  condition=condition,
                                                  ignore_count=ignore_count,
                                                  thread=thread, **kwargs)

    @watch('TargetSetWatchPoint')
    @action_valid_decorator_factory(TargetStates.NOT_RUNNING, 'execution')
    def set_watchpoint(self, variable, write=True, read=False):
        """Inserts a watchpoint

        :param      variable: The name of a variable or an address to watch
        :param bool write:    Write watchpoint
        :param bool read:     Read watchpoint
        """
        return self.protocols.execution.set_watchpoint(variable,
                                                  write=write,
                                                  read=read)

    @watch('TargetRemovebreakpoint')
    @action_valid_decorator_factory(TargetStates.STOPPED, 'execution')
    def remove_breakpoint(self, bkptno):
        """Deletes a breakpoint"""
        return self.protocols.execution.remove_breakpoint(bkptno)

    def update_state(self, state):
        self.log.info("State changed to to %s" % TargetStates(state))
        self.state = state
        #self._no_state_update_pending.set()

    @watch('TargetWait')
    def wait(self, state=TargetStates.STOPPED):
        while self.state & state == 0:
            pass

    def get_status(self):
        """
        Returns useful information about the target as a dict.
        """
        self.status['state'] = self.state
        return self.status

    def _get_unique_name(self, i=0):
        classname = type(self).__name__
        targetname = "{}{}".format(classname, i)
        if self.avatar and self.avatar.targets and targetname in self.avatar.targets:
            return self._get_unique_name(i + 1)
        return targetname

    def _resolve_executable_name(self):
        """
        Resolves the name of the executable for the endpoint.
        Order of operation:
            1: Check if config exists and whether target is installed
            2: Check sys_name from default config
            3: Check apt_name from default config
            4: BailOut
        """
        pass

    # ##generic aliases##
    wr = write_register
    rr = read_register
    rm = read_memory
    wm = write_memory
    bp = set_breakpoint
<|endoftext|>


<commit_before># Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.

# Python
import re
import os
from itertools import islice

# Django
from django.utils.encoding import smart_str

__all__ = ['skip_directory', 'could_be_playbook', 'could_be_inventory']

valid_playbook_re = re.compile(r'^\s*?-?\s*?(?:hosts|include):\s*?.*?$')
valid_inventory_re = re.compile(r'^[a-zA-Z0-9_.=\[\]]')

def skip_directory(relative_directory_path):
    path_elements = relative_directory_path.split(os.sep)
    # Exclude files in a roles subdirectory.
    if 'roles' in path_elements:
        return True
    # Filter files in a tasks subdirectory.
    if 'tasks' in path_elements:
        return True
    for element in path_elements:
        # Do not include dot files or dirs
        if element.startswith('.'):
            return True
    # Exclude anything inside of group or host vars directories
    if 'group_vars' in path_elements or 'host_vars' in path_elements:
        return True
    return False

def could_be_playbook(project_path, dir_path, filename):
    if os.path.splitext(filename)[-1] not in ['.yml', '.yaml']:
        return None
    playbook_path = os.path.join(dir_path, filename)
    # Filter files that do not have either hosts or top-level
    # includes. Use regex to allow files with invalid YAML to
    # show up.
    matched = False
    try:
        for n, line in enumerate(file(playbook_path)):
            if valid_playbook_re.match(line):
                matched = True
            # Any YAML file can also be encrypted with vault;
            # allow these to be used as the main playbook.
            elif n == 0 and line.startswith('$ANSIBLE_VAULT;'):
                matched = True
    except IOError:
        return None
    if not matched:
        return None
    return os.path.relpath(playbook_path, smart_str(project_path))

def could_be_inventory(project_path, dir_path, filename):
    # Decisions based exclusively on filename
    inventory_path = os.path.join(dir_path, filename)
    suspected_ext = os.path.splitext(filename)[-1]
    if suspected_ext in ['.yml', '.yaml', '.ini']:
        # Files with any of these extensions are always included
        return os.path.relpath(inventory_path, smart_str(project_path))
    elif '.' in suspected_ext:
        # If not using those extensions, inventory must have _no_ extension
        return None

    # Filter files that do not use a character set consistent with
    # Ansible inventory mainly
    try:
        # only read through first 10 lines for performance
        with open(inventory_path) as inv_file:
            for line in islice(inv_file, 10):
                if not valid_inventory_re.match(line):
                    return None
    except IOError:
        return None
    return os.path.relpath(inventory_path, smart_str(project_path))
<commit_msg>show executable python files in inventory list<commit_after># Copyright (c) 2017 Ansible by Red Hat
# All Rights Reserved.

# Python
import re
import os
from itertools import islice

# Django
from django.utils.encoding import smart_str

__all__ = ['skip_directory', 'could_be_playbook', 'could_be_inventory']

valid_playbook_re = re.compile(r'^\s*?-?\s*?(?:hosts|include):\s*?.*?$')
valid_inventory_re = re.compile(r'^[a-zA-Z0-9_.=\[\]]')

def skip_directory(relative_directory_path):
    path_elements = relative_directory_path.split(os.sep)
    # Exclude files in a roles subdirectory.
    if 'roles' in path_elements:
        return True
    # Filter files in a tasks subdirectory.
    if 'tasks' in path_elements:
        return True
    for element in path_elements:
        # Do not include dot files or dirs
        if element.startswith('.'):
            return True
    # Exclude anything inside of group or host vars directories
    if 'group_vars' in path_elements or 'host_vars' in path_elements:
        return True
    return False

def could_be_playbook(project_path, dir_path, filename):
    if os.path.splitext(filename)[-1] not in ['.yml', '.yaml']:
        return None
    playbook_path = os.path.join(dir_path, filename)
    # Filter files that do not have either hosts or top-level
    # includes. Use regex to allow files with invalid YAML to
    # show up.
    matched = False
    try:
        for n, line in enumerate(file(playbook_path)):
            if valid_playbook_re.match(line):
                matched = True
            # Any YAML file can also be encrypted with vault;
            # allow these to be used as the main playbook.
            elif n == 0 and line.startswith('$ANSIBLE_VAULT;'):
                matched = True
    except IOError:
        return None
    if not matched:
        return None
    return os.path.relpath(playbook_path, smart_str(project_path))

def could_be_inventory(project_path, dir_path, filename):
    # Decisions based exclusively on filename
    inventory_path = os.path.join(dir_path, filename)
    suspected_ext = os.path.splitext(filename)[-1]
    if suspected_ext in ['.yml', '.yaml', '.ini'] or os.access(inventory_path, os.X_OK):
        # Files with any of these extensions are always included
        return os.path.relpath(inventory_path, smart_str(project_path))
    elif '.' in suspected_ext:
        # If not using those extensions, inventory must have _no_ extension
        return None

    # Filter files that do not use a character set consistent with
    # Ansible inventory mainly
    try:
        # only read through first 10 lines for performance
        with open(inventory_path) as inv_file:
            for line in islice(inv_file, 10):
                if not valid_inventory_re.match(line):
                    return None
    except IOError:
        return None
    return os.path.relpath(inventory_path, smart_str(project_path))
<|endoftext|>


<commit_before>from collections import Counter
import os
import sys
import unittest
from unittest.mock import patch, MagicMock

from git_reviewers import reviewers

directory = os.path.dirname(os.path.realpath(__file__))
BASE_DIRECTORY = os.path.normpath(os.path.join(directory, '..', '..'))

class TestFindReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindReviewers()

    def test_get_reviewers(self):
        with self.assertRaises(NotImplementedError):
            self.finder.get_reviewers()

    @patch('subprocess.run')
    def test_run_command(self, mock_run):
        mock_run().stdout = b'asdf'
        data = self.finder.run_command(['ls'])
        self.assertEqual(data, ['asdf'])

    @patch('subprocess.run')
    def test_run_command_empty_response(self, mock_run):
        mock_run().stdout = b''
        data = self.finder.run_command([':'])
        self.assertEqual(data, [])

    def test_extract_username_from_generic_email(self):
        email = 'asdf@gmail.com'
        user = self.finder.extract_username_from_email(email)
        self.assertEqual(user, 'asdf@gmail.com')

    def test_extract_uber_username_from_email(self):
        email = 'asdf@uber.com'
        user = self.finder.extract_username_from_email(email)
        self.assertEqual(user, 'asdf')

class TestFindLogReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindFileLogReviewers()

    def test_gets_generic_emails(self):
        shortlog = '     3\tAlbert Wang <example@gmail.com>\n'
        email = self.finder.extract_username_from_shortlog(shortlog)
        self.assertEqual(email, 'example@gmail.com')

    def test_gets_uber_emails(self):
        shortlog = '     3\tAlbert Wang <albertyw@uber.com>\n'
        email = self.finder.extract_username_from_shortlog(shortlog)
        self.assertEqual(email, 'albertyw')

    def test_get_changed_files(self):
        with self.assertRaises(NotImplementedError):
            self.finder.get_changed_files()

    @patch('subprocess.run')
    def test_gets_reviewers(self, mock_run):
        changed_files = ['README.rst']
        self.finder.get_changed_files = MagicMock(return_value=changed_files)
        process = MagicMock()
        git_shortlog = b'     3\tAlbert Wang <albertyw@uber.com>\n'
        git_shortlog += b'3\tAlbert Wang <example@gmail.com>\n'
        process.stdout = git_shortlog
        mock_run.return_value = process
        users = self.finder.get_reviewers()
        reviewers = Counter({'albertyw': 1, 'example@gmail.com': 1})
        self.assertEqual(users, reviewers)

class TestFindDiffLogReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindDiffLogReviewers()

    @patch('subprocess.run')
    def test_gets_diff_files(self, mock_run):
        process = MagicMock()
        output = b'README.rst\ngit_reviewers/reviewers.py\n'
        process.stdout = output
        mock_run.return_value = process
        diff_files = self.finder.get_changed_files()
        expected = ['README.rst', 'git_reviewers/reviewers.py']
        self.assertEqual(diff_files, expected)

class TestLogReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindLogReviewers()

    def test_get_changed_files(self):
        changed_files = ['README.rst', 'setup.py']
        self.finder.run_command = MagicMock(return_value=changed_files)
        files = self.finder.get_changed_files()
        self.assertEqual(files, ['README.rst', 'setup.py'])

class TestFindArcCommitReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindArcCommitReviewers()

    def test_no_reviewers(self):
        log = ['asdf']
        self.finder.run_command = MagicMock(return_value=log)
        reviewers = self.finder.get_log_reviewers_from_file('file')
        self.assertEqual(reviewers, [])

    def test_reviewers(self):
        log = ['asdf', ' Reviewed By: asdf, qwer']
        self.finder.run_command = MagicMock(return_value=log)
        reviewers = self.finder.get_log_reviewers_from_file('file')
        self.assertEqual(reviewers, ['asdf', 'qwer'])

class TestShowReviewers(unittest.TestCase):
    @patch('builtins.print')
    def test_show_reviewers(self, mock_print):
        usernames = Counter({'albertyw': 1, 'asdf': 2})
        reviewers.show_reviewers(usernames)
        mock_print.assert_called_with('asdf, albertyw')

    @patch('builtins.print')
    def test_limit_show_reviewers(self, mock_print):
        usernames = Counter(
            {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5, 'f': 6, 'g': 7, 'h': 8}
        )
        reviewers.show_reviewers(usernames)
        mock_print.assert_called_with('h, g, f, e, d, c, b')

class TestMain(unittest.TestCase):
    @patch('builtins.print')
    def test_main(self, mock_print):
        with patch.object(sys, 'argv', ['reviewers.py']):
            reviewers.main()
        self.assertTrue(mock_print.called)

    @patch('argparse.ArgumentParser._print_message')
    def test_version(self, mock_print):
        with patch.object(sys, 'argv', ['reviewers.py', '-v']):
            with self.assertRaises(SystemExit):
                reviewers.main()
        self.assertTrue(mock_print.called)
        version = reviewers.__version__ + "\n"
        self.assertEqual(mock_print.call_args[0][0], version)

    @patch('builtins.print')
    def test_ignore_reviewers(self, mock_print):
        counter = Counter({'asdf':1, 'qwer': 1})
        with patch.object(sys, 'argv', ['reviewers.py', '-i', 'asdf']):
            with patch('git_reviewers.reviewers.FindFileLogReviewers.get_reviewers') as mock_get_reviewers:
                mock_get_reviewers.return_value = counter
                reviewers.main()
        self.assertEqual(mock_print.call_args[0][0], 'qwer')
<commit_msg>Flake8<commit_after>from collections import Counter
import os
import sys
import unittest
from unittest.mock import patch, MagicMock

from git_reviewers import reviewers

directory = os.path.dirname(os.path.realpath(__file__))
BASE_DIRECTORY = os.path.normpath(os.path.join(directory, '..', '..'))

class TestFindReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindReviewers()

    def test_get_reviewers(self):
        with self.assertRaises(NotImplementedError):
            self.finder.get_reviewers()

    @patch('subprocess.run')
    def test_run_command(self, mock_run):
        mock_run().stdout = b'asdf'
        data = self.finder.run_command(['ls'])
        self.assertEqual(data, ['asdf'])

    @patch('subprocess.run')
    def test_run_command_empty_response(self, mock_run):
        mock_run().stdout = b''
        data = self.finder.run_command([':'])
        self.assertEqual(data, [])

    def test_extract_username_from_generic_email(self):
        email = 'asdf@gmail.com'
        user = self.finder.extract_username_from_email(email)
        self.assertEqual(user, 'asdf@gmail.com')

    def test_extract_uber_username_from_email(self):
        email = 'asdf@uber.com'
        user = self.finder.extract_username_from_email(email)
        self.assertEqual(user, 'asdf')

class TestFindLogReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindFileLogReviewers()

    def test_gets_generic_emails(self):
        shortlog = '     3\tAlbert Wang <example@gmail.com>\n'
        email = self.finder.extract_username_from_shortlog(shortlog)
        self.assertEqual(email, 'example@gmail.com')

    def test_gets_uber_emails(self):
        shortlog = '     3\tAlbert Wang <albertyw@uber.com>\n'
        email = self.finder.extract_username_from_shortlog(shortlog)
        self.assertEqual(email, 'albertyw')

    def test_get_changed_files(self):
        with self.assertRaises(NotImplementedError):
            self.finder.get_changed_files()

    @patch('subprocess.run')
    def test_gets_reviewers(self, mock_run):
        changed_files = ['README.rst']
        self.finder.get_changed_files = MagicMock(return_value=changed_files)
        process = MagicMock()
        git_shortlog = b'     3\tAlbert Wang <albertyw@uber.com>\n'
        git_shortlog += b'3\tAlbert Wang <example@gmail.com>\n'
        process.stdout = git_shortlog
        mock_run.return_value = process
        users = self.finder.get_reviewers()
        reviewers = Counter({'albertyw': 1, 'example@gmail.com': 1})
        self.assertEqual(users, reviewers)

class TestFindDiffLogReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindDiffLogReviewers()

    @patch('subprocess.run')
    def test_gets_diff_files(self, mock_run):
        process = MagicMock()
        output = b'README.rst\ngit_reviewers/reviewers.py\n'
        process.stdout = output
        mock_run.return_value = process
        diff_files = self.finder.get_changed_files()
        expected = ['README.rst', 'git_reviewers/reviewers.py']
        self.assertEqual(diff_files, expected)

class TestLogReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindLogReviewers()

    def test_get_changed_files(self):
        changed_files = ['README.rst', 'setup.py']
        self.finder.run_command = MagicMock(return_value=changed_files)
        files = self.finder.get_changed_files()
        self.assertEqual(files, ['README.rst', 'setup.py'])

class TestFindArcCommitReviewers(unittest.TestCase):
    def setUp(self):
        self.finder = reviewers.FindArcCommitReviewers()

    def test_no_reviewers(self):
        log = ['asdf']
        self.finder.run_command = MagicMock(return_value=log)
        reviewers = self.finder.get_log_reviewers_from_file('file')
        self.assertEqual(reviewers, [])

    def test_reviewers(self):
        log = ['asdf', ' Reviewed By: asdf, qwer']
        self.finder.run_command = MagicMock(return_value=log)
        reviewers = self.finder.get_log_reviewers_from_file('file')
        self.assertEqual(reviewers, ['asdf', 'qwer'])

class TestShowReviewers(unittest.TestCase):
    @patch('builtins.print')
    def test_show_reviewers(self, mock_print):
        usernames = Counter({'albertyw': 1, 'asdf': 2})
        reviewers.show_reviewers(usernames)
        mock_print.assert_called_with('asdf, albertyw')

    @patch('builtins.print')
    def test_limit_show_reviewers(self, mock_print):
        usernames = Counter(
            {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5, 'f': 6, 'g': 7, 'h': 8}
        )
        reviewers.show_reviewers(usernames)
        mock_print.assert_called_with('h, g, f, e, d, c, b')

class TestMain(unittest.TestCase):
    @patch('builtins.print')
    def test_main(self, mock_print):
        with patch.object(sys, 'argv', ['reviewers.py']):
            reviewers.main()
        self.assertTrue(mock_print.called)

    @patch('argparse.ArgumentParser._print_message')
    def test_version(self, mock_print):
        with patch.object(sys, 'argv', ['reviewers.py', '-v']):
            with self.assertRaises(SystemExit):
                reviewers.main()
        self.assertTrue(mock_print.called)
        version = reviewers.__version__ + "\n"
        self.assertEqual(mock_print.call_args[0][0], version)

    @patch('builtins.print')
    def test_ignore_reviewers(self, mock_print):
        counter = Counter({'asdf': 1, 'qwer': 1})
        get_reviewers = (
            'git_reviewers.reviewers.'
            'FindFileLogReviewers.get_reviewers'
        )
        with patch.object(sys, 'argv', ['reviewers.py', '-i', 'asdf']):
            with patch(get_reviewers) as mock_get_reviewers:
                mock_get_reviewers.return_value = counter
                reviewers.main()
        self.assertEqual(mock_print.call_args[0][0], 'qwer')
<|endoftext|>


<commit_before>
from __future__ import print_function

import os
from six.moves import shlex_quote
import tempfile

import sh

def _multi_output(commands):
    '''
        Shows output of multiple git commands

        .. make sure the output is shell-safe
    '''
    tname = None
    try:
        with tempfile.NamedTemporaryFile(mode='w', delete=False) as fp:
            tname = fp.name
            for command in commands:
                fp.write(command)
                fp.write('\n')

        os.system('bash %s | less -FRX' % tname)
    finally:
        if tname:
            os.unlink(tname)

def _get_commits(fname, rev_range):
    # Ask for the commit, timestamp, and the filename in case it changed
    # .. the timestamp doesn't guarantee ordering, but good enough

    args = ['log', '--follow', '--pretty=%ct %H', '--name-only']
    if rev_range:
        args.append(rev_range)

    args.append(fname)
    it = sh.git(*args, _tty_out=False, _iter=True)

    while True:
        l1 = next(it).strip().split()
        next(it)
        l2 = next(it).strip()

        yield int(l1[0]), l1[1], l2

def git_diff(filenames, rev_range):
    commands = ['git diff --follow -w --color %s %s' % (rev_range, shlex_quote(f)) for f in filenames]
    _multi_output(commands)

def git_log(cfg, filenames, rev_range=None):
    '''
        A git log implementation that allows more flexibility:

        - Follow multiple files
        - Exclude commits we don't want to see
    '''

    if len(filenames) == 0:
        print("Specify at least one file to log")

    elif len(filenames) == 1 and not cfg.excluded_commits:
        if rev_range:
            cmd = 'git log --follow -p %s %s' % (rev_range, filenames[0])
        else:
            cmd = 'git log --follow -p %s' % (filenames[0])

        os.system(cmd)
    else:
        # To show renames properly, we have to switch to the root
        # of the repository and then specify the potentially renamed file
        # for each commit
        oldcwd = os.getcwd()
        git_toplevel = str(sh.git('rev-parse', '--show-toplevel')).strip()

        # Use git log to generate lists of commits for each file, sort
        commit_data = []
        for fname in filenames:
            commit_data += _get_commits(fname, rev_range)

        if not len(commit_data):
            return

        # Create an index of filenames per commit id
        fname_by_commit = {}
        for _, commit, fname in commit_data:
            fname_by_commit.setdefault(commit, []).append(fname)

        # Uniquify (http://www.peterbe.com/plog/uniqifiers-benchmark)
        seen = set()
        seen_add = seen.add
        commits = [c for _, c, _ in commit_data if not (c in seen or seen_add(c))]

        # Finally, display them
        try:
            os.chdir(git_toplevel)

            commands = []

            for commit in commits:
                if not cfg.is_commit_excluded(commit):
                    for i, fname in enumerate(sorted(fname_by_commit[commit])):
                        fname = shlex_quote(fname)
                        # git log --follow only allows a single filename, so
                        # merge the outputs together using separate commands
                        if i == 0:
                            commands.append('git log -p -1 --follow --color %s -- %s' % (commit, fname))
                        else:
                            commands.append('git log -p -1 --follow --color --pretty='' %s -- %s' % (commit, fname))

            _multi_output(commands)

        finally:
            os.chdir(oldcwd)

<commit_msg>Turns out that sort was needed<commit_after>
from __future__ import print_function

import os
from six.moves import shlex_quote
import tempfile

import sh

def _multi_output(commands):
    '''
        Shows output of multiple git commands

        .. make sure the output is shell-safe
    '''
    tname = None
    try:
        with tempfile.NamedTemporaryFile(mode='w', delete=False) as fp:
            tname = fp.name
            for command in commands:
                fp.write(command)
                fp.write('\n')

        os.system('bash %s | less -FRX' % tname)
    finally:
        if tname:
            os.unlink(tname)

def _get_commits(fname, rev_range):
    # Ask for the commit, timestamp, and the filename in case it changed
    # .. the timestamp doesn't guarantee ordering, but good enough

    args = ['log', '--follow', '--pretty=%ct %H', '--name-only']
    if rev_range:
        args.append(rev_range)

    args.append(fname)
    it = sh.git(*args, _tty_out=False, _iter=True)

    while True:
        l1 = next(it).strip().split()
        next(it)
        l2 = next(it).strip()

        yield int(l1[0]), l1[1], l2

def git_diff(filenames, rev_range):
    commands = ['git diff --follow -w --color %s %s' % (rev_range, shlex_quote(f)) for f in filenames]
    _multi_output(commands)

def git_log(cfg, filenames, rev_range=None):
    '''
        A git log implementation that allows more flexibility:

        - Follow multiple files
        - Exclude commits we don't want to see
    '''

    if len(filenames) == 0:
        print("Specify at least one file to log")

    elif len(filenames) == 1 and not cfg.excluded_commits:
        if rev_range:
            cmd = 'git log --follow -p %s %s' % (rev_range, filenames[0])
        else:
            cmd = 'git log --follow -p %s' % (filenames[0])

        os.system(cmd)
    else:
        # To show renames properly, we have to switch to the root
        # of the repository and then specify the potentially renamed file
        # for each commit
        oldcwd = os.getcwd()
        git_toplevel = str(sh.git('rev-parse', '--show-toplevel')).strip()

        # Use git log to generate lists of commits for each file, sort
        commit_data = []
        for fname in filenames:
            commit_data += _get_commits(fname, rev_range)

        if not len(commit_data):
            return

        # Sort the results by timestamp
        if len(filenames) > 1:
            commit_data.sort(reverse=True)

        # Create an index of filenames per commit id
        fname_by_commit = {}
        for _, commit, fname in commit_data:
            fname_by_commit.setdefault(commit, []).append(fname)

        # Uniquify (http://www.peterbe.com/plog/uniqifiers-benchmark)
        seen = set()
        seen_add = seen.add
        commits = [c for _, c, _ in commit_data if not (c in seen or seen_add(c))]

        # Finally, display them
        try:
            os.chdir(git_toplevel)

            commands = []

            for commit in commits:
                if not cfg.is_commit_excluded(commit):
                    for i, fname in enumerate(sorted(fname_by_commit[commit])):
                        fname = shlex_quote(fname)
                        # git log --follow only allows a single filename, so
                        # merge the outputs together using separate commands
                        if i == 0:
                            commands.append('git log -p -1 --follow --color %s -- %s' % (commit, fname))
                        else:
                            commands.append('git log -p -1 --follow --color --pretty='' %s -- %s' % (commit, fname))

            _multi_output(commands)

        finally:
            os.chdir(oldcwd)

<|endoftext|>