Make __mp_main__ an alias for __main__ in all processes to simplify

pickling of classes defined in main module.

2 files changed, 13 insertions, 16 deletions

diff --git a/Lib/multiprocessing/__init__.py b/Lib/multiprocessing/__init__.py
index 1f3e67c..efad532 100644
--- a/Lib/multiprocessing/__init__.py
+++ b/Lib/multiprocessing/__init__.py
@@ -40,6 +40,13 @@ from multiprocessing.process import Process, current_process, active_children
 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
index af6580d..fe4ee33 100644
--- a/Lib/multiprocessing/forking.py
+++ b/Lib/multiprocessing/forking.py
@@ -441,27 +441,17 @@ def prepare(data):
                 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