@@ -42,6 +42,7 @@ def __init__(self):
4242 self ._inherited_fds = None
4343 self ._lock = threading .Lock ()
4444 self ._preload_modules = ['__main__' ]
45+ self ._preload_on_error = 'ignore'
4546
4647 def _stop (self ):
4748 # Method used by unit tests to stop the server
@@ -64,11 +65,22 @@ def _stop_unlocked(self):
6465 self ._forkserver_address = None
6566 self ._forkserver_authkey = None
6667
67- def set_forkserver_preload (self , modules_names ):
68- '''Set list of module names to try to load in forkserver process.'''
68+ def set_forkserver_preload (self , modules_names , * , on_error = 'ignore' ):
69+ '''Set list of module names to try to load in forkserver process.
70+
71+ The on_error parameter controls how import failures are handled:
72+ "ignore" (default) silently ignores failures, "warn" emits warnings,
73+ and "fail" raises exceptions breaking the forkserver context.
74+ '''
6975 if not all (type (mod ) is str for mod in modules_names ):
7076 raise TypeError ('module_names must be a list of strings' )
77+ if on_error not in ('ignore' , 'warn' , 'fail' ):
78+ raise ValueError (
79+ f"on_error must be 'ignore', 'warn', or 'fail', "
80+ f"not { on_error !r} "
81+ )
7182 self ._preload_modules = modules_names
83+ self ._preload_on_error = on_error
7284
7385 def get_inherited_fds (self ):
7486 '''Return list of fds inherited from parent process.
@@ -107,6 +119,14 @@ def connect_to_new_process(self, fds):
107119 wrapped_client , self ._forkserver_authkey )
108120 connection .deliver_challenge (
109121 wrapped_client , self ._forkserver_authkey )
122+ except (EOFError , ConnectionError , BrokenPipeError ) as exc :
123+ if (self ._preload_modules and
124+ self ._preload_on_error == 'fail' ):
125+ exc .add_note (
126+ "Forkserver process may have crashed during module "
127+ "preloading. Check stderr."
128+ )
129+ raise
110130 finally :
111131 wrapped_client ._detach ()
112132 del wrapped_client
@@ -154,6 +174,8 @@ def ensure_running(self):
154174 main_kws ['main_path' ] = data ['init_main_from_path' ]
155175 if 'sys_argv' in data :
156176 main_kws ['sys_argv' ] = data ['sys_argv' ]
177+ if self ._preload_on_error != 'ignore' :
178+ main_kws ['on_error' ] = self ._preload_on_error
157179
158180 with socket .socket (socket .AF_UNIX ) as listener :
159181 address = connection .arbitrary_address ('AF_UNIX' )
@@ -198,8 +220,69 @@ def ensure_running(self):
198220#
199221#
200222
223+ def _handle_import_error (on_error , modinfo , exc , * , warn_stacklevel ):
224+ """Handle an import error according to the on_error policy."""
225+ match on_error :
226+ case 'fail' :
227+ raise
228+ case 'warn' :
229+ warnings .warn (
230+ f"Failed to preload { modinfo } : { exc } " ,
231+ ImportWarning ,
232+ stacklevel = warn_stacklevel + 1
233+ )
234+ case 'ignore' :
235+ pass
236+
237+
238+ def _handle_preload (preload , main_path = None , sys_path = None , sys_argv = None ,
239+ on_error = 'ignore' ):
240+ """Handle module preloading with configurable error handling.
241+
242+ Args:
243+ preload: List of module names to preload.
244+ main_path: Path to __main__ module if '__main__' is in preload.
245+ sys_path: sys.path to use for imports (None means use current).
246+ sys_argv: sys.argv to use (None means use current).
247+ on_error: How to handle import errors ("ignore", "warn", or "fail").
248+ """
249+ if not preload :
250+ return
251+
252+ if sys_argv is not None :
253+ sys .argv [:] = sys_argv
254+ if sys_path is not None :
255+ sys .path [:] = sys_path
256+
257+ if '__main__' in preload and main_path is not None :
258+ process .current_process ()._inheriting = True
259+ try :
260+ spawn .import_main_path (main_path )
261+ except Exception as e :
262+ # Catch broad Exception because import_main_path() uses
263+ # runpy.run_path() which executes the script and can raise
264+ # any exception, not just ImportError
265+ _handle_import_error (
266+ on_error , f"__main__ from { main_path !r} " , e , warn_stacklevel = 2
267+ )
268+ finally :
269+ del process .current_process ()._inheriting
270+
271+ for modname in preload :
272+ try :
273+ __import__ (modname )
274+ except ImportError as e :
275+ _handle_import_error (
276+ on_error , f"module { modname !r} " , e , warn_stacklevel = 2
277+ )
278+
279+ # gh-135335: flush stdout/stderr in case any of the preloaded modules
280+ # wrote to them, otherwise children might inherit buffered data
281+ util ._flush_std_streams ()
282+
283+
201284def main (listener_fd , alive_r , preload , main_path = None , sys_path = None ,
202- * , sys_argv = None , authkey_r = None ):
285+ * , sys_argv = None , authkey_r = None , on_error = 'ignore' ):
203286 """Run forkserver."""
204287 if authkey_r is not None :
205288 try :
@@ -210,26 +293,7 @@ def main(listener_fd, alive_r, preload, main_path=None, sys_path=None,
210293 else :
211294 authkey = b''
212295
213- if preload :
214- if sys_argv is not None :
215- sys .argv [:] = sys_argv
216- if sys_path is not None :
217- sys .path [:] = sys_path
218- if '__main__' in preload and main_path is not None :
219- process .current_process ()._inheriting = True
220- try :
221- spawn .import_main_path (main_path )
222- finally :
223- del process .current_process ()._inheriting
224- for modname in preload :
225- try :
226- __import__ (modname )
227- except ImportError :
228- pass
229-
230- # gh-135335: flush stdout/stderr in case any of the preloaded modules
231- # wrote to them, otherwise children might inherit buffered data
232- util ._flush_std_streams ()
296+ _handle_preload (preload , main_path , sys_path , sys_argv , on_error )
233297
234298 util ._close_stdin ()
235299
0 commit comments