core: support takes_econtext and takes_router decorators.

docs/api.rst

@@ -22,6 +22,28 @@ mitogen.core
 
 .. automodule:: mitogen.core
 
+.. function:: mitogen.core.takes_econtext
+
+    Decorator that marks a function or class method to automatically receive a
+    kwarg named `econtext`, referencing the
+    :py:class:`econtext.core.ExternalContext` active in the context in which
+    the function is being invoked in. The decorator is only meaningful when the
+    function is invoked via :py:data:`econtext.core.CALL_FUNCTION`.
+
+    When the function is invoked directly, `econtext` must still be passed to it
+    explicitly.
+
+.. function:: mitogen.core.takes_router
+
+    Decorator that marks a function or class method to automatically receive a
+    kwarg named `router`, referencing the :py:class:`econtext.core.Router`
+    active in the context in which the function is being invoked in. The
+    decorator is only meaningful when the function is invoked via
+    :py:data:`econtext.core.CALL_FUNCTION`.
+
+    When the function is invoked directly, `router` must still be passed to it
+    explicitly.
+
 
 mitogen.master
 --------------

mitogen/core.py

@@ -117,6 +117,16 @@ def fire(obj, name, *args, **kwargs):
     return [func(*args, **kwargs) for func in signals.get(name, ())]
 
 
+def takes_econtext(func):
+    func.mitogen_takes_econtext = True
+    return func
+
+
+def takes_router(func):
+    func.mitogen_takes_router = True
+    return func
+
+
 def set_cloexec(fd):
     flags = fcntl.fcntl(fd, fcntl.F_GETFD)
     fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
@@ -1185,6 +1195,10 @@ class ExternalContext(object):
                 if klass:
                     obj = getattr(obj, klass)
                 fn = getattr(obj, func)
+                if getattr(fn, 'mitogen_takes_econtext', None):
+                    kwargs.setdefault('econtext', self)
+                if getattr(fn, 'mitogen_takes_router', None):
+                    kwargs.setdefault('router', self.router)
                 ret = fn(*args, **kwargs)
                 self.router.route(
                     Message.pickled(ret, dst_id=msg.src_id, handle=msg.reply_to)

mitogen/fakessh.py

@@ -311,6 +311,7 @@ def _fakessh_main(mitogen_, dest_context_id):
 # Public API.
 #
 
+@mitogen.core.takes_router
 def run(dest, router, args, deadline=None):
     """
     Run the command specified by the argument vector `args` such that ``PATH``