Have mitogen.main(profiling=True) effect the master process too

6 years ago · e7c98e3bda
parent a868498469
commit e7c98e3bda
1 changed files with 7 additions and 2 deletions
--- a/mitogen/init.py
+++ b/mitogen/init.py
@ -108,9 +108,14 @@ def main(log_level='INFO', profiling=False):
    def wrapper(func):
        if func.__module__ != '__main__':
            return func
-        import mitogen.master
+        import mitogen.parent
        import mitogen.utils
+        mitogen.core.enable_profiling()
        mitogen.master.Router.profiling = profiling
        utils.log_to_file(level=log_level)
-        return utils.run_with_router(func)
+        return mitogen.core._profile_hook(
+            'main',
+            utils.run_with_router,
+            func,
+        )
    return wrapper