xref: /aosp_15_r20/external/bazelbuild-rules_python/python/private/zip_main_template.py (revision 60517a1edbc8ecf509223e9af94a7adec7d736b8)
1# Template for the __main__.py file inserted into zip files
2#
3# NOTE: This file is a "stage 1" bootstrap, so it's responsible for locating the
4# desired runtime and having it run the stage 2 bootstrap. This means it can't
5# assume much about the current runtime and environment. e.g., the current
6# runtime may not be the correct one, the zip may not have been extract, the
7# runfiles env vars may not be set, etc.
8#
9# NOTE: This program must retain compatibility with a wide variety of Python
10# versions since it is run by an unknown Python interpreter.
11
12import sys
13
14# The Python interpreter unconditionally prepends the directory containing this
15# script (following symlinks) to the import path. This is the cause of #9239,
16# and is a special case of #7091. We therefore explicitly delete that entry.
17# TODO(#7091): Remove this hack when no longer necessary.
18del sys.path[0]
19
20import os
21import shutil
22import subprocess
23import tempfile
24import zipfile
25
26_STAGE2_BOOTSTRAP = "%stage2_bootstrap%"
27_PYTHON_BINARY = "%python_binary%"
28_WORKSPACE_NAME = "%workspace_name%"
29
30
31# Return True if running on Windows
32def is_windows():
33    return os.name == "nt"
34
35
36def get_windows_path_with_unc_prefix(path):
37    """Adds UNC prefix after getting a normalized absolute Windows path.
38
39    No-op for non-Windows platforms or if running under python2.
40    """
41    path = path.strip()
42
43    # No need to add prefix for non-Windows platforms.
44    # And \\?\ doesn't work in python 2 or on mingw
45    if not is_windows() or sys.version_info[0] < 3:
46        return path
47
48    # Starting in Windows 10, version 1607(OS build 14393), MAX_PATH limitations have been
49    # removed from common Win32 file and directory functions.
50    # Related doc: https://docs.microsoft.com/en-us/windows/win32/fileio/maximum-file-path-limitation?tabs=cmd#enable-long-paths-in-windows-10-version-1607-and-later
51    import platform
52
53    if platform.win32_ver()[1] >= "10.0.14393":
54        return path
55
56    # import sysconfig only now to maintain python 2.6 compatibility
57    import sysconfig
58
59    if sysconfig.get_platform() == "mingw":
60        return path
61
62    # Lets start the unicode fun
63    unicode_prefix = "\\\\?\\"
64    if path.startswith(unicode_prefix):
65        return path
66
67    # os.path.abspath returns a normalized absolute path
68    return unicode_prefix + os.path.abspath(path)
69
70
71def has_windows_executable_extension(path):
72    return path.endswith(".exe") or path.endswith(".com") or path.endswith(".bat")
73
74
75if is_windows() and not has_windows_executable_extension(_PYTHON_BINARY):
76    _PYTHON_BINARY = _PYTHON_BINARY + ".exe"
77
78
79def search_path(name):
80    """Finds a file in a given search path."""
81    search_path = os.getenv("PATH", os.defpath).split(os.pathsep)
82    for directory in search_path:
83        if directory:
84            path = os.path.join(directory, name)
85            if os.path.isfile(path) and os.access(path, os.X_OK):
86                return path
87    return None
88
89
90def find_python_binary(module_space):
91    """Finds the real Python binary if it's not a normal absolute path."""
92    return find_binary(module_space, _PYTHON_BINARY)
93
94
95def print_verbose(*args, mapping=None, values=None):
96    if bool(os.environ.get("RULES_PYTHON_BOOTSTRAP_VERBOSE")):
97        if mapping is not None:
98            for key, value in sorted((mapping or {}).items()):
99                print(
100                    "bootstrap: stage 1:",
101                    *args,
102                    f"{key}={value!r}",
103                    file=sys.stderr,
104                    flush=True,
105                )
106        elif values is not None:
107            for i, v in enumerate(values):
108                print(
109                    "bootstrap: stage 1:",
110                    *args,
111                    f"[{i}] {v!r}",
112                    file=sys.stderr,
113                    flush=True,
114                )
115        else:
116            print("bootstrap: stage 1:", *args, file=sys.stderr, flush=True)
117
118
119def find_binary(module_space, bin_name):
120    """Finds the real binary if it's not a normal absolute path."""
121    if not bin_name:
122        return None
123    if bin_name.startswith("//"):
124        # Case 1: Path is a label. Not supported yet.
125        raise AssertionError(
126            "Bazel does not support execution of Python interpreters via labels yet"
127        )
128    elif os.path.isabs(bin_name):
129        # Case 2: Absolute path.
130        return bin_name
131    # Use normpath() to convert slashes to os.sep on Windows.
132    elif os.sep in os.path.normpath(bin_name):
133        # Case 3: Path is relative to the repo root.
134        return os.path.join(module_space, bin_name)
135    else:
136        # Case 4: Path has to be looked up in the search path.
137        return search_path(bin_name)
138
139
140def extract_zip(zip_path, dest_dir):
141    """Extracts the contents of a zip file, preserving the unix file mode bits.
142
143    These include the permission bits, and in particular, the executable bit.
144
145    Ideally the zipfile module should set these bits, but it doesn't. See:
146    https://bugs.python.org/issue15795.
147
148    Args:
149        zip_path: The path to the zip file to extract
150        dest_dir: The path to the destination directory
151    """
152    zip_path = get_windows_path_with_unc_prefix(zip_path)
153    dest_dir = get_windows_path_with_unc_prefix(dest_dir)
154    with zipfile.ZipFile(zip_path) as zf:
155        for info in zf.infolist():
156            zf.extract(info, dest_dir)
157            # UNC-prefixed paths must be absolute/normalized. See
158            # https://docs.microsoft.com/en-us/windows/desktop/fileio/naming-a-file#maximum-path-length-limitation
159            file_path = os.path.abspath(os.path.join(dest_dir, info.filename))
160            # The Unix st_mode bits (see "man 7 inode") are stored in the upper 16
161            # bits of external_attr. Of those, we set the lower 12 bits, which are the
162            # file mode bits (since the file type bits can't be set by chmod anyway).
163            attrs = info.external_attr >> 16
164            if attrs != 0:  # Rumor has it these can be 0 for zips created on Windows.
165                os.chmod(file_path, attrs & 0o7777)
166
167
168# Create the runfiles tree by extracting the zip file
169def create_module_space():
170    temp_dir = tempfile.mkdtemp("", "Bazel.runfiles_")
171    extract_zip(os.path.dirname(__file__), temp_dir)
172    # IMPORTANT: Later code does `rm -fr` on dirname(module_space) -- it's
173    # important that deletion code be in sync with this directory structure
174    return os.path.join(temp_dir, "runfiles")
175
176
177def execute_file(
178    python_program,
179    main_filename,
180    args,
181    env,
182    module_space,
183    workspace,
184):
185    # type: (str, str, list[str], dict[str, str], str, str|None, str|None) -> ...
186    """Executes the given Python file using the various environment settings.
187
188    This will not return, and acts much like os.execv, except is much
189    more restricted, and handles Bazel-related edge cases.
190
191    Args:
192      python_program: (str) Path to the Python binary to use for execution
193      main_filename: (str) The Python file to execute
194      args: (list[str]) Additional args to pass to the Python file
195      env: (dict[str, str]) A dict of environment variables to set for the execution
196      module_space: (str) Path to the module space/runfiles tree directory
197      workspace: (str|None) Name of the workspace to execute in. This is expected to be a
198          directory under the runfiles tree.
199    """
200    # We want to use os.execv instead of subprocess.call, which causes
201    # problems with signal passing (making it difficult to kill
202    # Bazel). However, these conditions force us to run via
203    # subprocess.call instead:
204    #
205    # - On Windows, os.execv doesn't handle arguments with spaces
206    #   correctly, and it actually starts a subprocess just like
207    #   subprocess.call.
208    # - When running in a workspace or zip file, we need to clean up the
209    #   workspace after the process finishes so control must return here.
210    try:
211        subprocess_argv = [python_program, main_filename] + args
212        print_verbose("subprocess argv:", values=subprocess_argv)
213        print_verbose("subprocess env:", mapping=env)
214        print_verbose("subprocess cwd:", workspace)
215        ret_code = subprocess.call(subprocess_argv, env=env, cwd=workspace)
216        sys.exit(ret_code)
217    finally:
218        # NOTE: dirname() is called because create_module_space() creates a
219        # sub-directory within a temporary directory, and we want to remove the
220        # whole temporary directory.
221        shutil.rmtree(os.path.dirname(module_space), True)
222
223
224def main():
225    print_verbose("running zip main bootstrap")
226    print_verbose("initial argv:", values=sys.argv)
227    print_verbose("initial environ:", mapping=os.environ)
228    print_verbose("initial sys.executable", sys.executable)
229    print_verbose("initial sys.version", sys.version)
230
231    args = sys.argv[1:]
232
233    new_env = {}
234
235    # The main Python source file.
236    # The magic string percent-main-percent is replaced with the runfiles-relative
237    # filename of the main file of the Python binary in BazelPythonSemantics.java.
238    main_rel_path = _STAGE2_BOOTSTRAP
239    if is_windows():
240        main_rel_path = main_rel_path.replace("/", os.sep)
241
242    module_space = create_module_space()
243    print_verbose("extracted runfiles to:", module_space)
244
245    new_env["RUNFILES_DIR"] = module_space
246
247    # Don't prepend a potentially unsafe path to sys.path
248    # See: https://docs.python.org/3.11/using/cmdline.html#envvar-PYTHONSAFEPATH
249    new_env["PYTHONSAFEPATH"] = "1"
250
251    main_filename = os.path.join(module_space, main_rel_path)
252    main_filename = get_windows_path_with_unc_prefix(main_filename)
253    assert os.path.exists(main_filename), (
254        "Cannot exec() %r: file not found." % main_filename
255    )
256    assert os.access(main_filename, os.R_OK), (
257        "Cannot exec() %r: file not readable." % main_filename
258    )
259
260    program = python_program = find_python_binary(module_space)
261    if python_program is None:
262        raise AssertionError("Could not find python binary: " + _PYTHON_BINARY)
263
264    # Some older Python versions on macOS (namely Python 3.7) may unintentionally
265    # leave this environment variable set after starting the interpreter, which
266    # causes problems with Python subprocesses correctly locating sys.executable,
267    # which subsequently causes failure to launch on Python 3.11 and later.
268    if "__PYVENV_LAUNCHER__" in os.environ:
269        del os.environ["__PYVENV_LAUNCHER__"]
270
271    new_env.update((key, val) for key, val in os.environ.items() if key not in new_env)
272
273    workspace = None
274    # If RUN_UNDER_RUNFILES equals 1, it means we need to
275    # change directory to the right runfiles directory.
276    # (So that the data files are accessible)
277    if os.environ.get("RUN_UNDER_RUNFILES") == "1":
278        workspace = os.path.join(module_space, _WORKSPACE_NAME)
279
280    sys.stdout.flush()
281    execute_file(
282        python_program,
283        main_filename,
284        args,
285        new_env,
286        module_space,
287        workspace,
288    )
289
290
291if __name__ == "__main__":
292    main()
293