Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 304b85c

Browse filesBrowse 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
Copy full SHA for 304b85c

File tree

9 files changed

+106
-53
lines changed
Filter options

9 files changed

+106
-53
lines changed

‎examples/pip_install/WORKSPACE

Copy file name to clipboardExpand all lines: examples/pip_install/WORKSPACE
+6Lines 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

Copy file name to clipboardExpand all lines: examples/pip_parse/WORKSPACE
+6Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/extract_wheels/__init__.py
+10-10Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/extract_wheels/lib/arguments.py
+21Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/extract_wheels/lib/arguments_test.py
+14-4Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/parse_requirements_to_bzl/__init__.py
+1-14Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py
+10-10Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/parse_requirements_to_bzl/parse_requirements_to_bzl_test.py
+7-2Lines 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

Copy file name to clipboardExpand all lines: python/pip_install/pip_repository.bzl
+31-13Lines 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 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.