diff --git a/docs/howitworks.rst b/docs/howitworks.rst index 8ae45c96..f23ac9b4 100644 --- a/docs/howitworks.rst +++ b/docs/howitworks.rst @@ -456,6 +456,24 @@ source is fetched, the method builds a new module object using the best practice documented in PEP-302. +Neutralizing ``__main__`` +######################### + +To avoid accidental execution of the ``__main__`` module's code in a slave +context, when serving the source of the main module, Mitogen removes any code +occurring after the first conditional that looks like a standard ``__main__`` +execution guard: + +.. code-block:: python + + # Code that looks like this is stripped from __main__. + if __name__ == '__main__': + run_some_code() + +This is a hack, but it's the least annoying hack I've found for the problem +yet. + + Avoiding Negative Imports ######################### diff --git a/mitogen/master.py b/mitogen/master.py index a1a959be..0a5eeca9 100644 --- a/mitogen/master.py +++ b/mitogen/master.py @@ -566,6 +566,17 @@ class ModuleResponder(object): def __repr__(self): return 'ModuleResponder(%r)' % (self._router,) + MAIN_RE = re.compile(r'^if\s+__name__\s*==\s*.__main__.\s*:', re.M) + + def neutralize_main(self, src): + """Given the source for the __main__ module, try to find where it + begins conditional execution based on a "if __name__ == '__main__'" + guard, and remove any code after that point.""" + match = self.MAIN_RE.search(src) + if match: + return src[:match.start()] + return src + def _on_get_module(self, msg): LOG.debug('%r.get_module(%r)', self, msg) if msg == mitogen.core._DEAD: @@ -584,6 +595,8 @@ class ModuleResponder(object): else: pkg_present = None + if fullname == '__main__': + source = self.neutralize_main(source) compressed = zlib.compress(source) related = list(self._finder.find_related(fullname)) self._router.route(