Add support for callbacks that take more than a single positional argument.

bdarnell · bdarnell · commit 5872db2e660a · 2011-09-11T15:50:34.000-07:00
Closes tornadoweb#351.
diff --git a/tornado/gen.py b/tornado/gen.py
@@ -27,10 +27,10 @@ def get(self):
             do_something_with_response(response)
             self.render("template.html")
 
-`Task` works with any function that takes a ``callback`` keyword argument
-(and runs that callback with zero or one arguments).  You can also yield
-a list of ``Tasks``, which will be started at the same time and run in parallel;
-a list of results will be returned when they are all finished::
+`Task` works with any function that takes a ``callback`` keyword
+argument.  You can also yield a list of ``Tasks``, which will be
+started at the same time and run in parallel; a list of results will
+be returned when they are all finished::
 
     def get(self):
         http_client = AsyncHTTPClient()
@@ -55,9 +55,16 @@ def get(self):
 asynchronous operations to be started at different times and proceed
 in parallel: yield several callbacks with different keys, then wait
 for them once all the async operations have started.
+
+The result of a `Wait` or `Task` yield expression depends on how the callback
+was run.  If it was called with no arguments, the result is ``None``.  If
+it was called with one argument, the result is that argument.  If it was
+called with more than one argument or any keyword arguments, the result
+is an `Arguments` object, which is a named tuple ``(args, kwargs)``.
 """
 
 import functools
+import operator
 import sys
 import types
 
@@ -134,10 +141,7 @@ def is_ready(self):
         return True
 
     def get_result(self):
-        return self.callback
-
-    def callback(self, arg=None):
-        self.runner.set_result(self.key, arg)
+        return self.runner.result_callback(self.key)
 
 class Wait(YieldPoint):
     """Returns the argument passed to the result of a previous `Callback`."""
@@ -191,24 +195,23 @@ class Task(YieldPoint):
     """
     def __init__(self, func, *args, **kwargs):
         assert "callback" not in kwargs
-        kwargs["callback"] = self.callback
-        self.func = functools.partial(func, *args, **kwargs)
+        self.args = args
+        self.kwargs = kwargs
+        self.func = func
 
     def start(self, runner):
         self.runner = runner
         self.key = object()
         runner.register_callback(self.key)
-        self.func()
+        self.kwargs["callback"] = runner.result_callback(self.key)
+        self.func(*self.args, **self.kwargs)
     
     def is_ready(self):
         return self.runner.is_ready(self.key)
 
     def get_result(self):
         return self.runner.pop_result(self.key)
 
-    def callback(self, arg=None):
-        self.runner.set_result(self.key, arg)
-
 class Multi(YieldPoint):
     """Runs multiple asynchronous operations in parallel.
 
@@ -318,3 +321,29 @@ def run(self):
         finally:
             self.running = False
 
+    def result_callback(self, key):
+        def inner(*args, **kwargs):
+            if kwargs or len(args) > 1:
+                result = Arguments(args, kwargs)
+            elif args:
+                result = args[0]
+            else:
+                result = None
+            self.set_result(key, result)
+        return inner
+
+# in python 2.6+ this could be a collections.namedtuple
+class Arguments(tuple):
+    """The result of a yield expression whose callback had more than one
+    argument (or keyword arguments).
+
+    The `Arguments` object can be used as a tuple ``(args, kwargs)``
+    or an object with attributes ``args`` and ``kwargs``.
+    """
+    __slots__ = ()
+
+    def __new__(cls, args, kwargs):
+        return tuple.__new__(cls, (args, kwargs))
+
+    args = property(operator.itemgetter(0))
+    kwargs = property(operator.itemgetter(1))
diff --git a/tornado/test/gen_test.py b/tornado/test/gen_test.py
@@ -191,6 +191,35 @@ def f():
             self.stop()
         self.run_gen(f)
 
+    def test_arguments(self):
+        @gen.engine
+        def f():
+            (yield gen.Callback("noargs"))()
+            self.assertEqual((yield gen.Wait("noargs")), None)
+            (yield gen.Callback("1arg"))(42)
+            self.assertEqual((yield gen.Wait("1arg")), 42)
+
+            (yield gen.Callback("kwargs"))(value=42)
+            result = yield gen.Wait("kwargs")
+            self.assertTrue(isinstance(result, gen.Arguments))
+            self.assertEqual(((), dict(value=42)), result)
+            self.assertEqual(dict(value=42), result.kwargs)
+
+            (yield gen.Callback("2args"))(42, 43)
+            result = yield gen.Wait("2args")
+            self.assertTrue(isinstance(result, gen.Arguments))
+            self.assertEqual(((42, 43), {}), result)
+            self.assertEqual((42, 43), result.args)
+
+            def task_func(callback):
+                callback(None, error="foo")
+            result = yield gen.Task(task_func)
+            self.assertTrue(isinstance(result, gen.Arguments))
+            self.assertEqual(((None,), dict(error="foo")), result)
+
+            self.stop()
+        self.run_gen(f)
+
 
 class GenSequenceHandler(RequestHandler):
     @asynchronous
diff --git a/website/sphinx/gen.rst b/website/sphinx/gen.rst
@@ -21,3 +21,8 @@
    .. autoclass:: Wait
 
    .. autoclass:: WaitAll
+
+   Other classes
+   -------------
+
+   .. autoclass:: Arguments
