117 lines
4.6 KiB
Python
117 lines
4.6 KiB
Python
# tractor: structured concurrent "actors".
|
|
# Copyright 2018-eternity Tyler Goodlet.
|
|
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
"""
|
|
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
|
|
|
|
|
|
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
|