Make __mp_main__ an alias for __main__ in all processes to simplify

shibturn · shibturn · commit 5046e9741cd6 · 2012-10-08T13:07:00.000+01:00
pickling of classes defined in main module.
diff --git a/Lib/multiprocessing/__init__.py b/Lib/multiprocessing/__init__.py
@@ -39,6 +39,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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,13 @@`
`39`	`39`	`from multiprocessing.process import Process, current_process, active_children`
`40`	`40`	`from multiprocessing.util import SUBDEBUG, SUBWARNING`
`41`	`41`
	`42`	`+#`
	`43`	`+# Alias for main module -- will be reset by bootstrapping child processes`
	`44`	`+#`
	`45`	`+`
	`46`	`+if '__main__' in sys.modules:`
	`47`	`+ sys.modules['__mp_main__'] = sys.modules['__main__']`
	`48`	`+`
`42`	`49`	`#`
`43`	`50`	`# Exceptions`
`44`	`51`	`#`