From 07d4d799f1b82657ed23d911eb1c8d486439e32a Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 13 Feb 2018 19:59:26 +0545
Subject: [PATCH] Add mitogen.main() decorator mainly for docs and demo use.

---
 docs/api.rst        |  1 +
 docs/internals.rst  |  6 +++---
 mitogen/__init__.py | 40 ++++++++++++++++++++++++++++++++++++++++
 mitogen/core.py     |  1 +
 4 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/docs/api.rst b/docs/api.rst
index 11f31354..58b6b500 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -16,6 +16,7 @@ mitogen Package
 .. autodata:: mitogen.context_id
 .. autodata:: mitogen.parent_id
 .. autodata:: mitogen.parent_ids
+.. autofunction:: mitogen.main
 
 
 mitogen.core
diff --git a/docs/internals.rst b/docs/internals.rst
index 247d6a41..e2281a47 100644
--- a/docs/internals.rst
+++ b/docs/internals.rst
@@ -134,7 +134,7 @@ Stream Classes
 .. autoclass:: Stream
    :members:
 
-.. currentmodule:: mitogen.master
+.. currentmodule:: mitogen.parent
 
 .. autoclass:: Stream
    :members:
@@ -181,7 +181,7 @@ Responder Class
 Forwarder Class
 ---------------
 
-.. currentmodule:: mitogen.master
+.. currentmodule:: mitogen.parent
 .. autoclass:: ModuleForwarder
    :members:
 
@@ -369,7 +369,7 @@ Helper Functions
         List of canonical submodule names.
 
 
-.. currentmodule:: mitogen.master
+.. currentmodule:: mitogen.parent
 
 .. autofunction:: minimize_source (source)
 
diff --git a/mitogen/__init__.py b/mitogen/__init__.py
index d6c94782..2c7a1c5d 100644
--- a/mitogen/__init__.py
+++ b/mitogen/__init__.py
@@ -63,3 +63,43 @@ parent_id = None
 #: This is an empty list in a master, otherwise it is a list of parent context
 #: IDs ordered from most direct to least direct.
 parent_ids = []
+
+
+def main(log_level='INFO'):
+    """
+    Convenience decorator primarily useful for writing discardable test scripts.
+
+    In the master process, when `func` is defined in the ``__main__`` module,
+    arranges for `func(router)` to be invoked immediately, with
+    :py:class:`mitogen.master.Router` construction and destruction handled just
+    as in :py:func:`mitogen.utils.run_with_router`. In slaves, this function
+    does nothing.
+
+    :param str log_level:
+        Logging package level to configure via
+        :py:func:`mitogen.utils.log_to_file`.
+
+    Example:
+
+    ::
+
+        import mitogen
+        import requests
+
+        def get_url(url):
+            return requests.get(url).text
+
+        @mitogen.main()
+        def main(router):
+            z = router.ssh(hostname='k3')
+            print z.call(get_url, 'http://www.google.com/')
+
+    """
+
+    def wrapper(func):
+        if func.__module__ != '__main__':
+            return func
+        from . import utils
+        utils.log_to_file(level=log_level)
+        return utils.run_with_router(func)
+    return wrapper
diff --git a/mitogen/core.py b/mitogen/core.py
index 31b8e9a2..e84c6794 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -1213,6 +1213,7 @@ class ExternalContext(object):
         mitogen.context_id = context_id
         mitogen.parent_ids = parent_ids
         mitogen.parent_id = parent_ids[0]
+        mitogen.main = lambda *args, **kwargs: (lambda func: None)
         mitogen.core = sys.modules['__main__']
         mitogen.core.__file__ = 'x/mitogen/core.py'  # For inspect.getsource()
         mitogen.core.__loader__ = self.importer