forked from goodboy/tractor
101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
"""
|
|
Helpers pulled mostly verbatim from ``multiprocessing.spawn``
|
|
to aid with "fixing up" the ``__main__`` module in subprocesses.
|
|
|
|
These helpers are needed for any spawing backend that doesn't already handle this.
|
|
For example when using ``trio_run_in_process`` it is needed but obviously not when
|
|
we're already using ``multiprocessing``.
|
|
"""
|
|
import os
|
|
import sys
|
|
import platform
|
|
import types
|
|
import runpy
|
|
from typing import Dict
|
|
|
|
|
|
ORIGINAL_DIR = os.path.abspath(os.getcwd())
|
|
|
|
|
|
def _mp_figure_out_main() -> Dict[str, str]:
|
|
"""Taken from ``multiprocessing.spawn.get_preparation_data()``.
|
|
|
|
Retrieve parent actor `__main__` module data.
|
|
"""
|
|
d = {}
|
|
# Figure out whether to initialise main in the subprocess as a module
|
|
# or through direct execution (or to leave it alone entirely)
|
|
main_module = sys.modules['__main__']
|
|
main_mod_name = getattr(main_module.__spec__, "name", None)
|
|
if main_mod_name is not None:
|
|
d['init_main_from_name'] = main_mod_name
|
|
# elif sys.platform != 'win32' or (not WINEXE and not WINSERVICE):
|
|
elif platform.system() != 'Windows':
|
|
main_path = getattr(main_module, '__file__', None)
|
|
if main_path is not None:
|
|
if (
|
|
not os.path.isabs(main_path) and (
|
|
ORIGINAL_DIR is not None)
|
|
):
|
|
# process.ORIGINAL_DIR is not None):
|
|
# main_path = os.path.join(process.ORIGINAL_DIR, main_path)
|
|
main_path = os.path.join(ORIGINAL_DIR, main_path)
|
|
d['init_main_from_path'] = os.path.normpath(main_path)
|
|
|
|
return d
|
|
|
|
|
|
# Multiprocessing module helpers to fix up the main module in
|
|
# spawned subprocesses
|
|
def _fixup_main_from_name(mod_name: str) -> None:
|
|
# __main__.py files for packages, directories, zip archives, etc, run
|
|
# their "main only" code unconditionally, so we don't even try to
|
|
# populate anything in __main__, nor do we make any changes to
|
|
# __main__ attributes
|
|
current_main = sys.modules['__main__']
|
|
if mod_name == "__main__" or mod_name.endswith(".__main__"):
|
|
return
|
|
|
|
# If this process was forked, __main__ may already be populated
|
|
if getattr(current_main.__spec__, "name", None) == mod_name:
|
|
return
|
|
|
|
# Otherwise, __main__ may contain some non-main code where we need to
|
|
# support unpickling it properly. We rerun it as __mp_main__ and make
|
|
# the normal __main__ an alias to that
|
|
# old_main_modules.append(current_main)
|
|
main_module = types.ModuleType("__mp_main__")
|
|
main_content = runpy.run_module(mod_name,
|
|
run_name="__mp_main__",
|
|
alter_sys=True) # type: ignore
|
|
main_module.__dict__.update(main_content)
|
|
sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
|
|
|
|
|
|
def _fixup_main_from_path(main_path: str) -> None:
|
|
# If this process was forked, __main__ may already be populated
|
|
current_main = sys.modules['__main__']
|
|
|
|
# Unfortunately, the main ipython launch script historically had no
|
|
# "if __name__ == '__main__'" guard, so we work around that
|
|
# by treating it like a __main__.py file
|
|
# See https://github.com/ipython/ipython/issues/4698
|
|
main_name = os.path.splitext(os.path.basename(main_path))[0]
|
|
if main_name == 'ipython':
|
|
return
|
|
|
|
# Otherwise, if __file__ already has the setting we expect,
|
|
# there's nothing more to do
|
|
if getattr(current_main, '__file__', None) == main_path:
|
|
return
|
|
|
|
# If the parent process has sent a path through rather than a module
|
|
# name we assume it is an executable script that may contain
|
|
# non-main code that needs to be executed
|
|
# old_main_modules.append(current_main)
|
|
main_module = types.ModuleType("__mp_main__")
|
|
main_content = runpy.run_path(main_path,
|
|
run_name="__mp_main__") # type: ignore
|
|
main_module.__dict__.update(main_content)
|
|
sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
|