11
11
import os
12
12
import pickle
13
13
import sys
14
+ import runpy
15
+ import types
14
16
15
17
from . import get_start_method , set_start_method
16
18
from . import process
@@ -157,15 +159,19 @@ def get_preparation_data(name):
157
159
start_method = get_start_method (),
158
160
)
159
161
160
- if sys .platform != 'win32' or (not WINEXE and not WINSERVICE ):
161
- main_path = getattr (sys .modules ['__main__' ], '__file__' , None )
162
- if not main_path and sys .argv [0 ] not in ('' , '-c' ):
163
- main_path = sys .argv [0 ]
162
+ # Figure out whether to initialise main in the subprocess as a module
163
+ # or through direct execution (or to leave it alone entirely)
164
+ main_module = sys .modules ['__main__' ]
165
+ main_mod_name = getattr (main_module .__spec__ , "name" , None )
166
+ if main_mod_name is not None :
167
+ d ['init_main_from_name' ] = main_mod_name
168
+ elif sys .platform != 'win32' or (not WINEXE and not WINSERVICE ):
169
+ main_path = getattr (main_module , '__file__' , None )
164
170
if main_path is not None :
165
171
if (not os .path .isabs (main_path ) and
166
172
process .ORIGINAL_DIR is not None ):
167
173
main_path = os .path .join (process .ORIGINAL_DIR , main_path )
168
- d ['main_path ' ] = os .path .normpath (main_path )
174
+ d ['init_main_from_path ' ] = os .path .normpath (main_path )
169
175
170
176
return d
171
177
@@ -206,55 +212,68 @@ def prepare(data):
206
212
if 'start_method' in data :
207
213
set_start_method (data ['start_method' ])
208
214
209
- if 'main_path' in data :
210
- import_main_path (data ['main_path' ])
215
+ if 'init_main_from_name' in data :
216
+ _fixup_main_from_name (data ['init_main_from_name' ])
217
+ elif 'init_main_from_path' in data :
218
+ _fixup_main_from_path (data ['init_main_from_path' ])
219
+
220
+ # Multiprocessing module helpers to fix up the main module in
221
+ # spawned subprocesses
222
+ def _fixup_main_from_name (mod_name ):
223
+ # __main__.py files for packages, directories, zip archives, etc, run
224
+ # their "main only" code unconditionally, so we don't even try to
225
+ # populate anything in __main__, nor do we make any changes to
226
+ # __main__ attributes
227
+ current_main = sys .modules ['__main__' ]
228
+ if mod_name == "__main__" or mod_name .endswith (".__main__" ):
229
+ return
230
+
231
+ # If this process was forked, __main__ may already be populated
232
+ if getattr (current_main .__spec__ , "name" , None ) == mod_name :
233
+ return
234
+
235
+ # Otherwise, __main__ may contain some non-main code where we need to
236
+ # support unpickling it properly. We rerun it as __mp_main__ and make
237
+ # the normal __main__ an alias to that
238
+ old_main_modules .append (current_main )
239
+ main_module = types .ModuleType ("__mp_main__" )
240
+ main_content = runpy .run_module (mod_name ,
241
+ run_name = "__mp_main__" ,
242
+ alter_sys = True )
243
+ main_module .__dict__ .update (main_content )
244
+ sys .modules ['__main__' ] = sys .modules ['__mp_main__' ] = main_module
245
+
246
+
247
+ def _fixup_main_from_path (main_path ):
248
+ # If this process was forked, __main__ may already be populated
249
+ current_main = sys .modules ['__main__' ]
250
+
251
+ # Unfortunately, the main ipython launch script historically had no
252
+ # "if __name__ == '__main__'" guard, so we work around that
253
+ # by treating it like a __main__.py file
254
+ # See https://github.com/ipython/ipython/issues/4698
255
+ main_name = os .path .splitext (os .path .basename (main_path ))[0 ]
256
+ if main_name == 'ipython' :
257
+ return
258
+
259
+ # Otherwise, if __file__ already has the setting we expect,
260
+ # there's nothing more to do
261
+ if getattr (current_main , '__file__' , None ) == main_path :
262
+ return
263
+
264
+ # If the parent process has sent a path through rather than a module
265
+ # name we assume it is an executable script that may contain
266
+ # non-main code that needs to be executed
267
+ old_main_modules .append (current_main )
268
+ main_module = types .ModuleType ("__mp_main__" )
269
+ main_content = runpy .run_path (main_path ,
270
+ run_name = "__mp_main__" )
271
+ main_module .__dict__ .update (main_content )
272
+ sys .modules ['__main__' ] = sys .modules ['__mp_main__' ] = main_module
211
273
212
274
213
275
def import_main_path (main_path ):
214
276
'''
215
277
Set sys.modules['__main__'] to module at main_path
216
278
'''
217
- # XXX (ncoghlan): The following code makes several bogus
218
- # assumptions regarding the relationship between __file__
219
- # and a module's real name. See PEP 302 and issue #10845
220
- if getattr (sys .modules ['__main__' ], '__file__' , None ) == main_path :
221
- return
222
-
223
- main_name = os .path .splitext (os .path .basename (main_path ))[0 ]
224
- if main_name == '__init__' :
225
- main_name = os .path .basename (os .path .dirname (main_path ))
226
-
227
- if main_name == '__main__' :
228
- main_module = sys .modules ['__main__' ]
229
- main_module .__file__ = main_path
230
- elif main_name != 'ipython' :
231
- # Main modules not actually called __main__.py may
232
- # contain additional code that should still be executed
233
- import importlib
234
- import types
235
-
236
- if main_path is None :
237
- dirs = None
238
- elif os .path .basename (main_path ).startswith ('__init__.py' ):
239
- dirs = [os .path .dirname (os .path .dirname (main_path ))]
240
- else :
241
- dirs = [os .path .dirname (main_path )]
242
-
243
- assert main_name not in sys .modules , main_name
244
- sys .modules .pop ('__mp_main__' , None )
245
- # We should not try to load __main__
246
- # since that would execute 'if __name__ == "__main__"'
247
- # clauses, potentially causing a psuedo fork bomb.
248
- main_module = types .ModuleType (main_name )
249
- # XXX Use a target of main_module?
250
- spec = importlib .find_spec (main_name , path = dirs )
251
- if spec is None :
252
- raise ImportError (name = main_name )
253
- methods = importlib ._bootstrap ._SpecMethods (spec )
254
- methods .init_module_attrs (main_module )
255
- main_module .__name__ = '__mp_main__'
256
- code = spec .loader .get_code (main_name )
257
- exec (code , main_module .__dict__ )
258
-
259
- old_main_modules .append (sys .modules ['__main__' ])
260
- sys .modules ['__main__' ] = sys .modules ['__mp_main__' ] = main_module
279
+ _fixup_main_from_path (main_path )
0 commit comments