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