Skip to content

Commit 7a9e4b2

Browse files
thesayynMatt Mackay
and
Matt Mackay
authored
feat: implement --build_python_zip pex (#324)
### Type of change - New feature or functionality (#236) ### Test plan - New test cases added --------- Co-authored-by: Matt Mackay <[email protected]>
1 parent 5142191 commit 7a9e4b2

File tree

13 files changed

+475
-1
lines changed

13 files changed

+475
-1
lines changed

MODULE.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ python.toolchain(
2323
tools = use_extension("//py:extensions.bzl", "py_tools")
2424
tools.rules_py_tools()
2525
use_repo(tools, "rules_py_tools")
26+
use_repo(tools, "rules_py_pex_2_3_1")
2627

2728
register_toolchains(
2829
"@rules_py_tools//:all",

docs/rules.md

Lines changed: 22 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

examples/py_pex_binary/BUILD.bazel

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
load("//py:defs.bzl", "py_binary", "py_pex_binary")
2+
3+
py_binary(
4+
name = "binary",
5+
srcs = ["say.py"],
6+
data = ["data.txt"],
7+
env = {
8+
"TEST": "1"
9+
},
10+
deps = [
11+
"@pypi_cowsay//:pkg",
12+
"@bazel_tools//tools/python/runfiles",
13+
],
14+
)
15+
16+
py_pex_binary(
17+
name = "py_pex_binary",
18+
binary = ":binary",
19+
inject_env = {
20+
"TEST": "1"
21+
}
22+
)

examples/py_pex_binary/data.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Mooo!

examples/py_pex_binary/say.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import cowsay
2+
import sys
3+
import os
4+
from bazel_tools.tools.python.runfiles import runfiles
5+
6+
print("sys.path entries:")
7+
for p in sys.path:
8+
print(" ", p)
9+
10+
print("")
11+
print("os.environ entries:")
12+
print(" runfiles dir:", os.environ.get("RUNFILES_DIR"))
13+
print(" injected env:", os.environ.get("TEST"))
14+
15+
print("")
16+
print("dir info: ")
17+
print(" current dir:", os.curdir)
18+
print(" current dir (absolute):", os.path.abspath(os.curdir))
19+
20+
21+
r = runfiles.Create()
22+
data_path = r.Rlocation("aspect_rules_py/examples/py_pex_binary/data.txt")
23+
24+
print("")
25+
print("runfiles lookup:")
26+
print(" data.txt:", data_path)
27+
28+
cowsay.cow(open(data_path).read())

py/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ bzl_library(
3737
"//py/private:py_venv",
3838
"//py/private:py_wheel",
3939
"//py/private:virtual",
40+
"//py/private:py_pex_binary",
4041
"@aspect_bazel_lib//lib:utils",
4142
],
4243
)

py/defs.bzl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@ load("@aspect_bazel_lib//lib:utils.bzl", "propagate_common_rule_attributes")
44
load("//py/private:py_binary.bzl", _py_binary = "py_binary", _py_test = "py_test")
55
load("//py/private:py_executable.bzl", "determine_main")
66
load("//py/private:py_library.bzl", _py_library = "py_library")
7+
load("//py/private:py_pex_binary.bzl", _py_pex_binary = "py_pex_binary")
78
load("//py/private:py_pytest_main.bzl", _py_pytest_main = "py_pytest_main")
89
load("//py/private:py_unpacked_wheel.bzl", _py_unpacked_wheel = "py_unpacked_wheel")
910
load("//py/private:virtual.bzl", _resolutions = "resolutions")
1011
load("//py/private:py_venv.bzl", _py_venv = "py_venv")
1112

13+
py_pex_binary = _py_pex_binary
1214
py_pytest_main = _py_pytest_main
1315

1416
py_venv = _py_venv

py/private/BUILD.bazel

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,16 @@ bzl_library(
106106
visibility = ["//py:__subpackages__"],
107107
)
108108

109+
bzl_library(
110+
name = "py_pex_binary",
111+
srcs = ["py_pex_binary.bzl"],
112+
visibility = ["//py:__subpackages__"],
113+
deps = [
114+
":py_semantics",
115+
"//py/private/toolchain:types",
116+
],
117+
)
118+
109119
bzl_library(
110120
name = "virtual",
111121
srcs = ["virtual.bzl"],

py/private/py_pex_binary.bzl

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
"Create python zip file https://peps.python.org/pep-0441/ (PEX)"
2+
3+
load("@rules_python//python:defs.bzl", "PyInfo")
4+
load("//py/private:py_semantics.bzl", _py_semantics = "semantics")
5+
load("//py/private/toolchain:types.bzl", "PY_TOOLCHAIN")
6+
7+
def _runfiles_path(file, workspace):
8+
if file.short_path.startswith("../"):
9+
return file.short_path[3:]
10+
else:
11+
return workspace + "/" + file.short_path
12+
13+
exclude_paths = [
14+
# following two lines will match paths we want to exclude in non-bzlmod setup
15+
"toolchain",
16+
"aspect_rules_py/py/tools/",
17+
# these will match in bzlmod setup
18+
"rules_python~~python~",
19+
"aspect_rules_py~/py/tools/",
20+
# these will match in bzlmod setup with --incompatible_use_plus_in_repo_names flag flipped.
21+
"rules_python++python+",
22+
"aspect_rules_py+/py/tools/"
23+
]
24+
25+
# determines if the given file is a `distinfo`, `dep` or a `source`
26+
# this required to allow PEX to put files into different places.
27+
#
28+
# --dep: into `<PEX_UNPACK_ROOT>/.deps/<name_of_the_package>`
29+
# --distinfo: is only used for determining package metadata
30+
# --source: into `<PEX_UNPACK_ROOT>/<relative_path_to_workspace_root>/<file_name>`
31+
def _map_srcs(f, workspace):
32+
dest_path = _runfiles_path(f, workspace)
33+
34+
# We exclude files from hermetic python toolchain.
35+
for exclude in exclude_paths:
36+
if dest_path.find(exclude) != -1:
37+
return []
38+
39+
site_packages_i = f.path.find("site-packages")
40+
41+
# if path contains `site-packages` and there is only two path segments
42+
# after it, it will be treated as third party dep.
43+
# Here are some examples of path we expect and use and ones we ignore.
44+
#
45+
# Match: `external/rules_python~~pip~pypi_39_rtoml/site-packages/rtoml-0.11.0.dist-info/INSTALLER`
46+
# Reason: It has two `/` after first `site-packages` substring.
47+
#
48+
# No Match: `external/rules_python~~pip~pypi_39_rtoml/site-packages/rtoml-0.11.0/src/mod/parse.py`
49+
# Reason: It has three `/` after first `site-packages` substring.
50+
if site_packages_i != -1 and f.path.count("/", site_packages_i) == 2:
51+
if f.path.find("dist-info", site_packages_i) != -1:
52+
return ["--distinfo={}".format(f.dirname)]
53+
return ["--dep={}".format(f.dirname)]
54+
55+
# If the path does not have a `site-packages` in it, then put it into
56+
# the standard runfiles tree.
57+
elif site_packages_i == -1:
58+
return ["--source={}={}".format(f.path, dest_path)]
59+
60+
return []
61+
62+
def _py_python_pex_impl(ctx):
63+
py_toolchain = _py_semantics.resolve_toolchain(ctx)
64+
65+
binary = ctx.attr.binary
66+
runfiles = binary[DefaultInfo].data_runfiles
67+
68+
output = ctx.actions.declare_file(ctx.attr.name + ".pex")
69+
70+
args = ctx.actions.args()
71+
72+
# Copy workspace name here to prevent ctx
73+
# being transferred to the execution phase.
74+
workspace_name = str(ctx.workspace_name)
75+
76+
args.add_all(
77+
ctx.attr.inject_env.items(),
78+
map_each = lambda e: "--inject-env={}={}".format(e[0], e[1]),
79+
# this is needed to allow passing a lambda to map_each
80+
allow_closure = True,
81+
)
82+
83+
args.add_all(
84+
binary[PyInfo].imports,
85+
format_each = "--sys-path=%s"
86+
)
87+
88+
args.add_all(
89+
runfiles.files,
90+
map_each = lambda f: _map_srcs(f, workspace_name),
91+
uniquify = True,
92+
# this is needed to allow passing a lambda (with workspace_name) to map_each
93+
allow_closure = True,
94+
)
95+
args.add(binary[DefaultInfo].files_to_run.executable, format = "--executable=%s")
96+
args.add(ctx.attr.python_shebang, format = "--python-shebang=%s")
97+
args.add(py_toolchain.python, format = "--python=%s")
98+
99+
py_version = py_toolchain.interpreter_version_info
100+
args.add_all(
101+
[
102+
constraint.format(major = py_version.major, minor = py_version.minor, patch = py_version.micro)
103+
for constraint in ctx.attr.python_interpreter_constraints
104+
],
105+
format_each = "--python-version-constraint=%s"
106+
)
107+
args.add(output, format = "--output-file=%s")
108+
109+
ctx.actions.run(
110+
executable = ctx.executable._pex,
111+
inputs = runfiles.files,
112+
arguments = [args],
113+
outputs = [output],
114+
mnemonic = "PyPex",
115+
progress_message = "Building PEX binary %{label}",
116+
)
117+
118+
return [
119+
DefaultInfo(files = depset([output]), executable = output)
120+
]
121+
122+
123+
_attrs = dict({
124+
"binary": attr.label(executable = True, cfg = "target", mandatory = True, doc = "A py_binary target"),
125+
"inject_env": attr.string_dict(
126+
doc = "Environment variables to set when running the pex binary.",
127+
default = {},
128+
),
129+
"python_shebang": attr.string(default = "#!/usr/bin/env python3"),
130+
"python_interpreter_constraints": attr.string_list(
131+
default = ["CPython=={major}.{minor}.*"],
132+
doc = """\
133+
Python interpreter versions this PEX binary is compatible with. A list of semver strings.
134+
The placeholder strings `{major}`, `{minor}`, `{patch}` can be used for gathering version
135+
information from the hermetic python toolchain.
136+
137+
For example, to enforce same interpreter version that Bazel uses, following can be used.
138+
139+
```starlark
140+
py_pex_binary
141+
python_interpreter_constraints = [
142+
"CPython=={major}.{minor}.{patch}"
143+
]
144+
)
145+
```
146+
"""),
147+
"_pex": attr.label(executable = True, cfg = "exec", default = "//py/tools/pex")
148+
})
149+
150+
151+
py_pex_binary = rule(
152+
doc = "Build a pex executable from a py_binary",
153+
implementation = _py_python_pex_impl,
154+
attrs = _attrs,
155+
toolchains = [
156+
PY_TOOLCHAIN
157+
],
158+
executable = True,
159+
)

py/private/run.tmpl.sh

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
# NB: we don't use a path from @bazel_tools//tools/sh:toolchain_type because that's configured for the exec
33
# configuration, while this script executes in the target configuration at runtime.
44

5+
# This is a special comment for py_pex_binary to find the python entrypoint.
6+
# __PEX_PY_BINARY_ENTRYPOINT__ {{ENTRYPOINT}}
7+
58
{{BASH_RLOCATION_FN}}
69
runfiles_export_envvars
710

@@ -55,4 +58,4 @@ if [ -n "${BASH:-}" -o -n "${ZSH_VERSION:-}" ] ; then
5558
hash -r 2> /dev/null
5659
fi
5760

58-
exec "{{EXEC_PYTHON_BIN}}" {{INTERPRETER_FLAGS}} "$(rlocation {{ENTRYPOINT}})" "$@"
61+
exec "{{EXEC_PYTHON_BIN}}" {{INTERPRETER_FLAGS}} "$(rlocation {{ENTRYPOINT}})" "$@"

py/toolchains.bzl

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Declare toolchains"""
22

3+
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_file")
34
load("//py/private/toolchain:autodetecting.bzl", _register_autodetecting_python_toolchain = "register_autodetecting_python_toolchain")
45
load("//py/private/toolchain:repo.bzl", "prerelease_toolchains_repo", "toolchains_repo")
56
load("//py/private/toolchain:tools.bzl", "TOOLCHAIN_PLATFORMS", "prebuilt_tool_repo")
@@ -32,3 +33,11 @@ def rules_py_toolchains(name = DEFAULT_TOOLS_REPOSITORY, register = True, is_pre
3233

3334
if register:
3435
native.register_toolchains("@{}//:all".format(name))
36+
37+
38+
http_file(
39+
name = "rules_py_pex_2_3_1",
40+
urls = ["https://files.pythonhosted.org/packages/e7/d0/fbda2a4d41d62d86ce53f5ae4fbaaee8c34070f75bb7ca009090510ae874/pex-2.3.1-py2.py3-none-any.whl"],
41+
sha256 = "64692a5bf6f298403aab930d22f0d836ae4736c5bc820e262e9092fe8c56f830",
42+
downloaded_file_path = "pex-2.3.1-py2.py3-none-any.whl",
43+
)

py/tools/pex/BUILD.bazel

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
load("//py:defs.bzl", "py_binary", "py_unpacked_wheel")
2+
3+
py_unpacked_wheel(
4+
name = "pex_unpacked",
5+
src = "@rules_py_pex_2_3_1//file",
6+
py_package_name = "pex"
7+
)
8+
9+
py_binary(
10+
name = "pex",
11+
srcs = ["main.py"],
12+
main = "main.py",
13+
deps = [":pex_unpacked"],
14+
visibility = ["//visibility:public"]
15+
)

0 commit comments

Comments
 (0)