squash merge wickman/pull/71

.gitignore

@@ -8,6 +8,7 @@
 *.swp
 *.*.swp
 .*.*.swo
+*~
 bin/
 build/
 develop-eggs/

Makefile

@@ -37,7 +37,7 @@ $(PYTHON):
 build: $(PYTHON)
 ifeq ($(GEVENT_SUPPORTED),yes)
 	$(INSTALL) -U -r requirements_gevent.txt
-	$(INSTALL) -f https://github.com/SiteSupport/gevent/downloads gevent==$(GEVENT_VERSION)
+	$(INSTALL) -f https://github.com/surfly/gevent/downloads gevent==$(GEVENT_VERSION)
 endif
 ifneq ($(TRAVIS), true)
 	$(INSTALL) -U -r requirements_sphinx.txt

kazoo/handlers/threading.py

@@ -13,6 +13,7 @@ environments that use threads.
 from __future__ import absolute_import
 
 import atexit
+import functools
 import logging
 import select
 import socket

kazoo/handlers/utils.py

@@ -5,6 +5,7 @@ try:
     import fcntl
 except ImportError:  # pragma: nocover
     HAS_FNCTL = False
+import functools
 import os
 
 
@@ -27,3 +28,40 @@ def create_tcp_socket(module):
         flags = fcntl.fcntl(sock, fcntl.F_GETFD)
         fcntl.fcntl(sock, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
     return sock
+
+
+def capture_exceptions(async_result):
+    """Return a new decorated function that propagates the exceptions of the
+    wrapped function to an async_result.
+
+    :param async_result: An async result implementing :class:`IAsyncResult`
+
+    """
+    def capture(function):
+        @functools.wraps(function)
+        def captured_function(*args, **kwargs):
+            try:
+                return function(*args, **kwargs)
+            except Exception as exc:
+                async_result.set_exception(exc)
+        return captured_function
+    return capture
+
+
+def wrap(async_result):
+    """Return a new decorated function that propagates the return value or
+    exception of wrapped function to an async_result.  NOTE: Only propagates a
+    non-None return value.
+
+    :param async_result: An async result implementing :class:`IAsyncResult`
+
+    """
+    def capture(function):
+        @capture_exceptions(async_result)
+        def captured_function(*args, **kwargs):
+            value = function(*args, **kwargs)
+            if value is not None:
+                async_result.set(value)
+            return value
+        return captured_function
+    return capture

kazoo/tests/test_threading_handler.py

@@ -2,6 +2,7 @@ import threading
 import unittest
 
 import mock
+from nose.tools import assert_raises
 from nose.tools import eq_
 from nose.tools import raises
 
@@ -261,3 +262,59 @@ class TestThreadingAsync(unittest.TestCase):
         async.unlink(add_on)
         async.set('fred')
         assert not mock_handler.completion_queue.put.called
+
+    def test_captured_exception(self):
+        from kazoo.handlers.utils import capture_exceptions
+
+        mock_handler = mock.Mock()
+        async = self._makeOne(mock_handler)
+
+        @capture_exceptions(async)
+        def exceptional_function():
+            return 1/0
+
+        exceptional_function()
+
+        assert_raises(ZeroDivisionError, async.get)
+
+    def test_no_capture_exceptions(self):
+        from kazoo.handlers.utils import capture_exceptions
+
+        mock_handler = mock.Mock()
+        async = self._makeOne(mock_handler)
+
+        lst = []
+
+        def add_on():
+            lst.append(True)
+
+        async.rawlink(add_on)
+
+        @capture_exceptions(async)
+        def regular_function():
+            return True
+
+        regular_function()
+
+        assert not mock_handler.completion_queue.put.called
+
+    def test_wraps(self):
+        from kazoo.handlers.utils import wrap
+
+        mock_handler = mock.Mock()
+        async = self._makeOne(mock_handler)
+
+        lst = []
+
+        def add_on(result):
+            lst.append(result.get())
+
+        async.rawlink(add_on)
+
+        @wrap(async)
+        def regular_function():
+            return 'hello'
+
+        assert regular_function() == 'hello'
+        assert mock_handler.completion_queue.put.called
+        assert async.get() == 'hello'