[Python-checkins] cpython: Make __mp_main__ an alias for __main__ in all processes to simplify

[richard.oudkerk]

http://hg.python.org/cpython/rev/22417bcd57f3
changeset:   79594:22417bcd57f3
user:        Richard Oudkerk <shibturn at gmail.com>
date:        Mon Oct 08 13:07:00 2012 +0100
summary:
  Make __mp_main__ an alias for __main__ in all processes to simplify
pickling of classes defined in main module.

files:
  Lib/multiprocessing/__init__.py |   7 ++++++
  Lib/multiprocessing/forking.py  |  22 +++++---------------
  2 files changed, 13 insertions(+), 16 deletions(-)


diff --git a/Lib/multiprocessing/__init__.py b/Lib/multiprocessing/__init__.py
--- a/Lib/multiprocessing/__init__.py
+++ b/Lib/multiprocessing/__init__.py
@@ -40,6 +40,13 @@
 from multiprocessing.util import SUBDEBUG, SUBWARNING
 
 #
+# Alias for main module -- will be reset by bootstrapping child processes
+#
+
+if '__main__' in sys.modules:
+    sys.modules['__mp_main__'] = sys.modules['__main__']
+
+#
 # Exceptions
 #
 
diff --git a/Lib/multiprocessing/forking.py b/Lib/multiprocessing/forking.py
--- a/Lib/multiprocessing/forking.py
+++ b/Lib/multiprocessing/forking.py
@@ -441,27 +441,17 @@
                 dirs = [os.path.dirname(main_path)]
 
             assert main_name not in sys.modules, main_name
+            sys.modules.pop('__mp_main__', None)
             file, path_name, etc = imp.find_module(main_name, dirs)
             try:
-                # We would like to do "imp.load_module('__main__', ...)"
-                # here.  However, that would cause 'if __name__ ==
-                # "__main__"' clauses to be executed.
+                # We should not do 'imp.load_module("__main__", ...)'
+                # since that would execute 'if __name__ == "__main__"'
+                # clauses, potentially causing a psuedo fork bomb.
                 main_module = imp.load_module(
-                    '__parents_main__', file, path_name, etc
+                    '__mp_main__', file, path_name, etc
                     )
             finally:
                 if file:
                     file.close()
 
-            sys.modules['__main__'] = main_module
-            main_module.__name__ = '__main__'
-
-            # Try to make the potentially picklable objects in
-            # sys.modules['__main__'] realize they are in the main
-            # module -- somewhat ugly.
-            for obj in list(main_module.__dict__.values()):
-                try:
-                    if obj.__module__ == '__parents_main__':
-                        obj.__module__ = '__main__'
-                except Exception:
-                    pass
+            sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module

-- 
Repository URL: http://hg.python.org/cpython