Skip to content

Commit 9d8c6b4

Browse files
committed
pybind_library_test: fix PYTHONHOME setup
build_defs.bzl: * Added a new pybind_py_env_test rule. This rule resolves the Python toolchain at analysis time, extracts the runtime path, and generates a shell wrapper script. * The wrapper script dynamically identifies the absolute path to the Python runtime within the Bazel runfiles and sets the PYTHONHOME environment variable before executing the actual C++ test binary. * Updated the pybind_library_test macro to build the C++ code as a private cc_binary and then wrap it with pybind_py_env_test.
1 parent 6a29558 commit 9d8c6b4

1 file changed

Lines changed: 95 additions & 3 deletions

File tree

build_defs.bzl

Lines changed: 95 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,80 @@ load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
1010
load("@rules_cc//cc:cc_library.bzl", "cc_library")
1111
load("@rules_cc//cc:cc_test.bzl", "cc_test")
1212

13+
def _pybind_py_env_test_impl(ctx):
14+
toolchain = ctx.toolchains["@rules_python//python:toolchain_type"]
15+
py3_runtime = toolchain.py3_runtime
16+
if not py3_runtime:
17+
fail("No python3 runtime found in toolchain")
18+
19+
# On Windows, we cannot use the shell script wrapper.
20+
if ctx.target_platform_has_constraint(ctx.attr._windows_constraint[platform_common.ConstraintValueInfo]):
21+
# On Windows, we need to return an executable created by this rule.
22+
# We create a symlink to the actual binary.
23+
# We use the same extension as the original binary (usually .exe).
24+
extension = ctx.executable.binary.extension
25+
executable = ctx.actions.declare_file(ctx.label.name + ("." + extension if extension else ""))
26+
ctx.actions.symlink(output = executable, target_file = ctx.executable.binary, is_executable = True)
27+
return [
28+
DefaultInfo(
29+
executable = executable,
30+
runfiles = ctx.runfiles(files = [executable])
31+
.merge(ctx.attr.binary[DefaultInfo].default_runfiles)
32+
.merge(ctx.runfiles(transitive_files = py3_runtime.files)),
33+
),
34+
]
35+
36+
interpreter = py3_runtime.interpreter
37+
38+
# Generate a wrapper script that sets PYTHONHOME and runs the C++ binary.
39+
script = ctx.actions.declare_file(ctx.label.name + ".sh")
40+
41+
content = "#!/bin/bash\n"
42+
content += "if [ -z \"$RUNFILES_DIR\" ]; then\n"
43+
content += " if [ -d \"$0.runfiles\" ]; then\n"
44+
content += " RUNFILES_DIR=\"$0.runfiles\"\n"
45+
content += " else\n"
46+
content += " RUNFILES_DIR=\"$(dirname \"$0\")/../..\"\n"
47+
content += " fi\n"
48+
content += "fi\n"
49+
content += "INTERPRETER_PATH=\"$RUNFILES_DIR/" + ctx.workspace_name + "/" + interpreter.short_path + "\"\n"
50+
content += "if [ ! -f \"$INTERPRETER_PATH\" ]; then\n"
51+
content += " INTERPRETER_PATH=$(find \"$RUNFILES_DIR\" -path \"*/" + interpreter.short_path + "\" | head -n 1)\n"
52+
content += "fi\n"
53+
content += "export PYTHONHOME=$(dirname $(dirname $(readlink -f \"$INTERPRETER_PATH\")))\n"
54+
content += "BINARY_PATH=\"$RUNFILES_DIR/" + ctx.workspace_name + "/" + ctx.executable.binary.short_path + "\"\n"
55+
content += "if [ ! -f \"$BINARY_PATH\" ]; then\n"
56+
content += " BINARY_PATH=$(find \"$RUNFILES_DIR\" -path \"*/" + ctx.executable.binary.short_path + "\" | head -n 1)\n"
57+
content += "fi\n"
58+
content += "exec \"$BINARY_PATH\" \"$@\"\n"
59+
60+
ctx.actions.write(script, content, is_executable = True)
61+
62+
runfiles = ctx.runfiles(files = [script, ctx.executable.binary])
63+
runfiles = runfiles.merge(ctx.attr.binary[DefaultInfo].default_runfiles)
64+
runfiles = runfiles.merge(ctx.runfiles(transitive_files = py3_runtime.files))
65+
66+
return [
67+
DefaultInfo(
68+
executable = script,
69+
runfiles = runfiles,
70+
),
71+
]
72+
73+
pybind_py_env_test = rule(
74+
implementation = _pybind_py_env_test_impl,
75+
test = True,
76+
attrs = {
77+
"binary": attr.label(
78+
executable = True,
79+
cfg = "target",
80+
mandatory = True,
81+
),
82+
"_windows_constraint": attr.label(default = "@platforms//os:windows"),
83+
},
84+
toolchains = ["@rules_python//python:toolchain_type"],
85+
)
86+
1387
def register_extension_info(**kwargs):
1488
pass
1589

@@ -148,17 +222,35 @@ def pybind_library_test(
148222
# Mark common dependencies as required for build_cleaner.
149223
tags = tags + ["req_dep=%s" % dep for dep in PYBIND_DEPS]
150224

151-
cc_test(
152-
name = name,
225+
# Pop test-only attributes that cc_binary doesn't support.
226+
test_kwargs = {}
227+
for attr in ["size", "timeout", "flaky", "shard_count", "local"]:
228+
if attr in kwargs:
229+
test_kwargs[attr] = kwargs.pop(attr)
230+
231+
# Build the actual C++ binary.
232+
cc_binary(
233+
name = name + "_bin",
153234
copts = copts + PYBIND_COPTS,
154235
features = features + PYBIND_FEATURES,
155-
tags = tags,
236+
testonly = True,
237+
visibility = ["//visibility:private"],
156238
deps = deps + PYBIND_DEPS + [
157239
"@rules_python//python/cc:current_py_cc_libs",
158240
],
159241
**kwargs
160242
)
161243

244+
# Use a wrapper rule to set PYTHONHOME and run the binary.
245+
pybind_py_env_test(
246+
name = name,
247+
binary = ":" + name + "_bin",
248+
testonly = True,
249+
tags = tags,
250+
visibility = kwargs.get("visibility"),
251+
**test_kwargs
252+
)
253+
162254
# Register extension with build_cleaner.
163255
register_extension_info(
164256
extension = pybind_extension,

0 commit comments

Comments
 (0)