Skip to content

Commit 304b85c

Browse files
mattemJonathon Belottihrfuller
authored
feat: allow setting custom environment variables on pip_repository and whl_library (bazel-contrib#460)
* feat: allow setting custom environment variables on pip_repository and whl_library * Serialize and deserialize environment dict in python process instead of starlark. * Refactor shared functions between extract_wheel and extract_single_wheel. * Every structured arg now has the same key when serialized. fixes bazel-contrib#490 * test for pip_data_exclude in arguments parsing test. * Also update docs in repository rule attr definition Co-authored-by: Jonathon Belotti <jonathon@canva.com> Co-authored-by: Henry Fuller <hrofuller@gmail.com>
1 parent cd64466 commit 304b85c

9 files changed

Lines changed: 106 additions & 53 deletions

File tree

‎examples/pip_install/WORKSPACE‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@ pip_install(
2929
# (Optional) You can set quiet to False if you want to see pip output.
3030
#quiet = False,
3131

32+
# (Optional) You can set an environment in the pip process to control its
33+
# behavior. Note that pip is run in "isolated" mode so no PIP_<VAR>_<NAME>
34+
# style env vars are read, but env vars that control requests and urllib3
35+
# can be passed
36+
#environment = {"HTTP_PROXY": "http://my.proxy.fun/"},
37+
3238
# Uses the default repository name "pip"
3339
requirements = "//:requirements.txt",
3440
)

‎examples/pip_parse/WORKSPACE‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@ pip_parse(
2929
# (Optional) You can set quiet to False if you want to see pip output.
3030
#quiet = False,
3131

32+
# (Optional) You can set an environment in the pip process to control its
33+
# behavior. Note that pip is run in "isolated" mode so no PIP_<VAR>_<NAME>
34+
# style env vars are read, but env vars that control requests and urllib3
35+
# can be passed
36+
# environment = {"HTTPS_PROXY": "http://my.proxy.fun/"},
37+
3238
# Uses the default repository name "pip_parsed_deps"
3339
requirements_lock = "//:requirements_lock.txt",
3440
)

‎python/pip_install/extract_wheels/__init__.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -60,29 +60,29 @@ def main() -> None:
6060
)
6161
arguments.parse_common_args(parser)
6262
args = parser.parse_args()
63+
deserialized_args = dict(vars(args))
64+
arguments.deserialize_structured_args(deserialized_args)
6365

64-
pip_args = [sys.executable, "-m", "pip", "--isolated", "wheel", "-r", args.requirements]
65-
if args.extra_pip_args:
66-
pip_args += json.loads(args.extra_pip_args)["args"]
66+
pip_args = (
67+
[sys.executable, "-m", "pip", "--isolated", "wheel", "-r", args.requirements] +
68+
deserialized_args["extra_pip_args"]
69+
)
6770

71+
env = os.environ.copy()
72+
env.update(deserialized_args["environment"])
6873
# Assumes any errors are logged by pip so do nothing. This command will fail if pip fails
69-
subprocess.run(pip_args, check=True)
74+
subprocess.run(pip_args, check=True, env=env)
7075

7176
extras = requirements.parse_extras(args.requirements)
7277

73-
if args.pip_data_exclude:
74-
pip_data_exclude = json.loads(args.pip_data_exclude)["exclude"]
75-
else:
76-
pip_data_exclude = []
77-
7878
repo_label = "@%s" % args.repo
7979

8080
targets = [
8181
'"%s%s"'
8282
% (
8383
repo_label,
8484
bazel.extract_wheel(
85-
whl, extras, pip_data_exclude, args.enable_implicit_namespace_pkgs
85+
whl, extras, deserialized_args["pip_data_exclude"], args.enable_implicit_namespace_pkgs
8686
),
8787
)
8888
for whl in glob.glob("*.whl")

‎python/pip_install/extract_wheels/lib/arguments.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import json
12
from argparse import ArgumentParser
23

34

@@ -21,4 +22,24 @@ def parse_common_args(parser: ArgumentParser) -> ArgumentParser:
2122
action="store_true",
2223
help="Disables conversion of implicit namespace packages into pkg-util style packages.",
2324
)
25+
parser.add_argument(
26+
"--environment",
27+
action="store",
28+
help="Extra environment variables to set on the pip environment.",
29+
)
2430
return parser
31+
32+
33+
def deserialize_structured_args(args):
34+
"""Deserialize structured arguments passed from the starlark rules.
35+
Args:
36+
args: dict of parsed command line arguments
37+
"""
38+
structured_args = ("extra_pip_args", "pip_data_exclude", "environment")
39+
for arg_name in structured_args:
40+
if args.get(arg_name) is not None:
41+
args[arg_name] = json.loads(args[arg_name])["arg"]
42+
else:
43+
args[arg_name] = []
44+
return args
45+

‎python/pip_install/extract_wheels/lib/arguments_test.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import unittest
44

55
from python.pip_install.extract_wheels.lib import arguments
6-
from python.pip_install.parse_requirements_to_bzl import deserialize_structured_args
76

87

98
class ArgumentsTestCase(unittest.TestCase):
@@ -12,15 +11,26 @@ def test_arguments(self) -> None:
1211
parser = arguments.parse_common_args(parser)
1312
repo_name = "foo"
1413
index_url = "--index_url=pypi.org/simple"
14+
extra_pip_args = [index_url]
1515
args_dict = vars(parser.parse_args(
16-
args=["--repo", repo_name, "--extra_pip_args={index_url}".format(index_url=json.dumps({"args": index_url}))]))
17-
args_dict = deserialize_structured_args(args_dict)
16+
args=["--repo", repo_name, f"--extra_pip_args={json.dumps({'arg': extra_pip_args})}"]))
17+
args_dict = arguments.deserialize_structured_args(args_dict)
1818
self.assertIn("repo", args_dict)
1919
self.assertIn("extra_pip_args", args_dict)
2020
self.assertEqual(args_dict["pip_data_exclude"], [])
2121
self.assertEqual(args_dict["enable_implicit_namespace_pkgs"], False)
2222
self.assertEqual(args_dict["repo"], repo_name)
23-
self.assertEqual(args_dict["extra_pip_args"], index_url)
23+
self.assertEqual(args_dict["extra_pip_args"], extra_pip_args)
24+
25+
def test_deserialize_structured_args(self) -> None:
26+
serialized_args = {
27+
"pip_data_exclude": json.dumps({"arg": ["**.foo"]}),
28+
"environment": json.dumps({"arg": {"PIP_DO_SOMETHING": "True"}}),
29+
}
30+
args = arguments.deserialize_structured_args(serialized_args)
31+
self.assertEqual(args["pip_data_exclude"], ["**.foo"])
32+
self.assertEqual(args["environment"], {"PIP_DO_SOMETHING": "True"})
33+
self.assertEqual(args["extra_pip_args"], [])
2434

2535

2636
if __name__ == "__main__":

‎python/pip_install/parse_requirements_to_bzl/__init__.py‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -45,19 +45,6 @@ def repo_names_and_requirements(install_reqs: List[Tuple[InstallRequirement, str
4545
for ir, line in install_reqs
4646
]
4747

48-
def deserialize_structured_args(args):
49-
"""Deserialize structured arguments passed from the starlark rules.
50-
Args:
51-
args: dict of parsed command line arguments
52-
"""
53-
structured_args = ("extra_pip_args", "pip_data_exclude")
54-
for arg_name in structured_args:
55-
if args.get(arg_name) is not None:
56-
args[arg_name] = json.loads(args[arg_name])["args"]
57-
else:
58-
args[arg_name] = []
59-
return args
60-
6148

6249
def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str:
6350
"""
@@ -69,7 +56,7 @@ def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str:
6956
"""
7057

7158
args = dict(vars(all_args))
72-
args = deserialize_structured_args(args)
59+
args = arguments.deserialize_structured_args(args)
7360
args.setdefault("python_interpreter", sys.executable)
7461
# Pop this off because it wont be used as a config argument to the whl_library rule.
7562
requirements_lock = args.pop("requirements_lock")

‎python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,15 @@ def main() -> None:
2323
)
2424
arguments.parse_common_args(parser)
2525
args = parser.parse_args()
26+
deserialized_args = dict(vars(args))
27+
arguments.deserialize_structured_args(deserialized_args)
2628

2729
configure_reproducible_wheels()
2830

29-
pip_args = [sys.executable, "-m", "pip", "--isolated", "wheel", "--no-deps"]
30-
if args.extra_pip_args:
31-
pip_args += json.loads(args.extra_pip_args)["args"]
31+
pip_args = (
32+
[sys.executable, "-m", "pip", "--isolated", "wheel", "--no-deps"] +
33+
deserialized_args["extra_pip_args"]
34+
)
3235

3336
requirement_file = NamedTemporaryFile(mode='wb', delete=False)
3437
try:
@@ -41,8 +44,10 @@ def main() -> None:
4144
# so write our single requirement into a temp file in case it has any of those flags.
4245
pip_args.extend(["-r", requirement_file.name])
4346

47+
env = os.environ.copy()
48+
env.update(deserialized_args["environment"])
4449
# Assumes any errors are logged by pip so do nothing. This command will fail if pip fails
45-
subprocess.run(pip_args, check=True)
50+
subprocess.run(pip_args, check=True, env=env)
4651
finally:
4752
try:
4853
os.unlink(requirement_file.name)
@@ -53,16 +58,11 @@ def main() -> None:
5358
name, extras_for_pkg = requirements._parse_requirement_for_extra(args.requirement)
5459
extras = {name: extras_for_pkg} if extras_for_pkg and name else dict()
5560

56-
if args.pip_data_exclude:
57-
pip_data_exclude = json.loads(args.pip_data_exclude)["exclude"]
58-
else:
59-
pip_data_exclude = []
60-
6161
whl = next(iter(glob.glob("*.whl")))
6262
bazel.extract_wheel(
6363
whl,
6464
extras,
65-
pip_data_exclude,
65+
deserialized_args["pip_data_exclude"],
6666
args.enable_implicit_namespace_pkgs,
6767
incremental=True,
6868
incremental_repo_prefix=bazel.whl_library_repo_prefix(args.repo)

‎python/pip_install/parse_requirements_to_bzl/parse_requirements_to_bzl_test.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,10 @@ def test_generated_requirements_bzl(self) -> None:
2323
args.requirements_lock = requirements_lock.name
2424
args.repo = "pip_parsed_deps"
2525
extra_pip_args = ["--index-url=pypi.org/simple"]
26-
args.extra_pip_args = json.dumps({"args": extra_pip_args})
26+
pip_data_exclude = ["**.foo"]
27+
args.extra_pip_args = json.dumps({"arg": extra_pip_args})
28+
args.pip_data_exclude= json.dumps({"arg": pip_data_exclude})
29+
args.environment= json.dumps({"arg": {}})
2730
contents = generate_parsed_requirements_contents(args)
2831
library_target = "@pip_parsed_deps_pypi__foo//:pkg"
2932
whl_target = "@pip_parsed_deps_pypi__foo//:whl"
@@ -32,9 +35,11 @@ def test_generated_requirements_bzl(self) -> None:
3235
self.assertIn(all_requirements, contents, contents)
3336
self.assertIn(all_whl_requirements, contents, contents)
3437
self.assertIn(requirement_string, contents, contents)
35-
self.assertIn(requirement_string, contents, contents)
3638
all_flags = extra_pip_args + ["--require-hashes", "True"]
3739
self.assertIn("'extra_pip_args': {}".format(repr(all_flags)), contents, contents)
40+
self.assertIn("'pip_data_exclude': {}".format(repr(pip_data_exclude)), contents, contents)
41+
# Assert it gets set to an empty dict by default.
42+
self.assertIn("'environment': {}", contents, contents)
3843

3944

4045
if __name__ == "__main__":

‎python/pip_install/pip_repository.bzl‎

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,26 +27,38 @@ def _construct_pypath(rctx):
2727
def _parse_optional_attrs(rctx, args):
2828
"""Helper function to parse common attributes of pip_repository and whl_library repository rules.
2929
30+
This function also serializes the structured arguments as JSON
31+
so they can be passed on the command line to subprocesses.
32+
3033
Args:
3134
rctx: Handle to the rule repository context.
3235
args: A list of parsed args for the rule.
3336
Returns: Augmented args list.
3437
"""
35-
if rctx.attr.extra_pip_args:
38+
39+
# Check for None so we use empty default types from our attrs.
40+
# Some args want to be list, and some want to be dict.
41+
if rctx.attr.extra_pip_args != None:
3642
args += [
3743
"--extra_pip_args",
38-
struct(args = rctx.attr.extra_pip_args).to_json(),
44+
struct(arg = rctx.attr.extra_pip_args).to_json(),
3945
]
4046

41-
if rctx.attr.pip_data_exclude:
47+
if rctx.attr.pip_data_exclude != None:
4248
args += [
4349
"--pip_data_exclude",
44-
struct(exclude = rctx.attr.pip_data_exclude).to_json(),
50+
struct(arg = rctx.attr.pip_data_exclude).to_json(),
4551
]
4652

4753
if rctx.attr.enable_implicit_namespace_pkgs:
4854
args.append("--enable_implicit_namespace_pkgs")
4955

56+
if rctx.attr.environment != None:
57+
args += [
58+
"--environment",
59+
struct(arg = rctx.attr.environment).to_json(),
60+
]
61+
5062
return args
5163

5264
_BUILD_FILE_CONTENTS = """\
@@ -102,10 +114,8 @@ def _pip_repository_impl(rctx):
102114

103115
result = rctx.execute(
104116
args,
105-
environment = {
106-
# Manually construct the PYTHONPATH since we cannot use the toolchain here
107-
"PYTHONPATH": pypath,
108-
},
117+
# Manually construct the PYTHONPATH since we cannot use the toolchain here
118+
environment = {"PYTHONPATH": _construct_pypath(rctx)},
109119
timeout = rctx.attr.timeout,
110120
quiet = rctx.attr.quiet,
111121
)
@@ -126,6 +136,16 @@ and py_test targets must specify either `legacy_create_init=False` or the global
126136
This option is required to support some packages which cannot handle the conversion to pkg-util style.
127137
""",
128138
),
139+
"environment": attr.string_dict(
140+
doc = """
141+
Environment variables to set in the pip subprocess.
142+
Can be used to set common variables such as `http_proxy`, `https_proxy` and `no_proxy`
143+
Note that pip is run with "--isolated" on the CLI so PIP_<VAR>_<NAME>
144+
style env vars are ignored, but env vars that control requests and urllib3
145+
can be passed.
146+
""",
147+
default = {},
148+
),
129149
"extra_pip_args": attr.string_list(
130150
doc = "Extra arguments to pass on to pip. Must not contain spaces.",
131151
),
@@ -221,7 +241,6 @@ py_binary(
221241
def _impl_whl_library(rctx):
222242
# pointer to parent repo so these rules rerun if the definitions in requirements.bzl change.
223243
_parent_repo_label = Label("@{parent}//:requirements.bzl".format(parent = rctx.attr.repo))
224-
pypath = _construct_pypath(rctx)
225244
args = [
226245
rctx.attr.python_interpreter,
227246
"-m",
@@ -232,12 +251,11 @@ def _impl_whl_library(rctx):
232251
rctx.attr.repo,
233252
]
234253
args = _parse_optional_attrs(rctx, args)
254+
235255
result = rctx.execute(
236256
args,
237-
environment = {
238-
# Manually construct the PYTHONPATH since we cannot use the toolchain here
239-
"PYTHONPATH": pypath,
240-
},
257+
# Manually construct the PYTHONPATH since we cannot use the toolchain here
258+
environment = {"PYTHONPATH": _construct_pypath(rctx)},
241259
quiet = rctx.attr.quiet,
242260
timeout = rctx.attr.timeout,
243261
)

0 commit comments

Comments
 (0)