feat: allow setting custom environment variables on pip_repository and whl_library (#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 #490
* test for pip_data_exclude in arguments parsing test.

* Also update docs in repository rule attr definition

Co-authored-by: Jonathon Belotti <[email protected]>
Co-authored-by: Henry Fuller <[email protected]>

Author: 3 people

examples/pip_install/WORKSPACE

@@ -29,6 +29,12 @@ pip_install(
     # (Optional) You can set quiet to False if you want to see pip output.
     #quiet = False,
 
+    # (Optional) You can set an environment in the pip process to control its
+    # behavior. Note that pip is run in "isolated" mode so no PIP_<VAR>_<NAME>
+    # style env vars are read, but env vars that control requests and urllib3
+    # can be passed
+    #environment = {"HTTP_PROXY": "http://my.proxy.fun/"},
+
     # Uses the default repository name "pip"
     requirements = "//:requirements.txt",
 )

examples/pip_parse/WORKSPACE

@@ -29,6 +29,12 @@ pip_parse(
     # (Optional) You can set quiet to False if you want to see pip output.
     #quiet = False,
 
+    # (Optional) You can set an environment in the pip process to control its
+    # behavior. Note that pip is run in "isolated" mode so no PIP_<VAR>_<NAME>
+    # style env vars are read, but env vars that control requests and urllib3
+    # can be passed
+    # environment = {"HTTPS_PROXY": "http://my.proxy.fun/"},
+
     # Uses the default repository name "pip_parsed_deps"
     requirements_lock = "//:requirements_lock.txt",
 )

python/pip_install/extract_wheels/__init__.py

@@ -60,29 +60,29 @@ def main() -> None:
     )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
+    deserialized_args = dict(vars(args))
+    arguments.deserialize_structured_args(deserialized_args)
 
-    pip_args = [sys.executable, "-m", "pip", "--isolated", "wheel", "-r", args.requirements]
-    if args.extra_pip_args:
-        pip_args += json.loads(args.extra_pip_args)["args"]
+    pip_args = (
+        [sys.executable, "-m", "pip", "--isolated", "wheel", "-r", args.requirements] +
+        deserialized_args["extra_pip_args"]
+    )
 
+    env = os.environ.copy()
+    env.update(deserialized_args["environment"])
     # Assumes any errors are logged by pip so do nothing. This command will fail if pip fails
-    subprocess.run(pip_args, check=True)
+    subprocess.run(pip_args, check=True, env=env)
 
     extras = requirements.parse_extras(args.requirements)
 
-    if args.pip_data_exclude:
-        pip_data_exclude = json.loads(args.pip_data_exclude)["exclude"]
-    else:
-        pip_data_exclude = []
-
     repo_label = "@%s" % args.repo
 
     targets = [
         '"%s%s"'
         % (
             repo_label,
             bazel.extract_wheel(
-                whl, extras, pip_data_exclude, args.enable_implicit_namespace_pkgs
+                whl, extras, deserialized_args["pip_data_exclude"], args.enable_implicit_namespace_pkgs
             ),
         )
         for whl in glob.glob("*.whl")

python/pip_install/extract_wheels/lib/arguments.py

@@ -1,3 +1,4 @@
+import json
 from argparse import ArgumentParser
 
 
@@ -21,4 +22,24 @@ def parse_common_args(parser: ArgumentParser) -> ArgumentParser:
         action="store_true",
         help="Disables conversion of implicit namespace packages into pkg-util style packages.",
     )
+    parser.add_argument(
+        "--environment",
+        action="store",
+        help="Extra environment variables to set on the pip environment.",
+    )
     return parser
+
+
+def deserialize_structured_args(args):
+    """Deserialize structured arguments passed from the starlark rules.
+    Args:
+        args: dict of parsed command line arguments
+    """
+    structured_args = ("extra_pip_args", "pip_data_exclude", "environment")
+    for arg_name in structured_args:
+        if args.get(arg_name) is not None:
+            args[arg_name] = json.loads(args[arg_name])["arg"]
+        else:
+            args[arg_name] = []
+    return args
+

python/pip_install/extract_wheels/lib/arguments_test.py

@@ -3,7 +3,6 @@
 import unittest
 
 from python.pip_install.extract_wheels.lib import arguments
-from python.pip_install.parse_requirements_to_bzl import deserialize_structured_args
 
 
 class ArgumentsTestCase(unittest.TestCase):
@@ -12,15 +11,26 @@ def test_arguments(self) -> None:
         parser = arguments.parse_common_args(parser)
         repo_name = "foo"
         index_url = "--index_url=pypi.org/simple"
+        extra_pip_args = [index_url]
         args_dict = vars(parser.parse_args(
-            args=["--repo", repo_name, "--extra_pip_args={index_url}".format(index_url=json.dumps({"args": index_url}))]))
-        args_dict = deserialize_structured_args(args_dict)
+            args=["--repo", repo_name, f"--extra_pip_args={json.dumps({'arg': extra_pip_args})}"]))
+        args_dict = arguments.deserialize_structured_args(args_dict)
         self.assertIn("repo", args_dict)
         self.assertIn("extra_pip_args", args_dict)
         self.assertEqual(args_dict["pip_data_exclude"], [])
         self.assertEqual(args_dict["enable_implicit_namespace_pkgs"], False)
         self.assertEqual(args_dict["repo"], repo_name)
-        self.assertEqual(args_dict["extra_pip_args"], index_url)
+        self.assertEqual(args_dict["extra_pip_args"], extra_pip_args)
+
+    def test_deserialize_structured_args(self) -> None:
+        serialized_args = {
+            "pip_data_exclude": json.dumps({"arg": ["**.foo"]}),
+            "environment": json.dumps({"arg": {"PIP_DO_SOMETHING": "True"}}),
+        }
+        args = arguments.deserialize_structured_args(serialized_args)
+        self.assertEqual(args["pip_data_exclude"], ["**.foo"])
+        self.assertEqual(args["environment"], {"PIP_DO_SOMETHING": "True"})
+        self.assertEqual(args["extra_pip_args"], [])
 
 
 if __name__ == "__main__":

python/pip_install/parse_requirements_to_bzl/__init__.py

@@ -45,19 +45,6 @@ def repo_names_and_requirements(install_reqs: List[Tuple[InstallRequirement, str
         for ir, line in install_reqs
     ]
 
-def deserialize_structured_args(args):
-    """Deserialize structured arguments passed from the starlark rules.
-        Args:
-            args: dict of parsed command line arguments
-    """
-    structured_args = ("extra_pip_args", "pip_data_exclude")
-    for arg_name in structured_args:
-        if args.get(arg_name) is not None:
-            args[arg_name] = json.loads(args[arg_name])["args"]
-        else:
-            args[arg_name] = []
-    return args
-
 
 def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str:
     """
@@ -69,7 +56,7 @@ def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str:
     """
 
     args = dict(vars(all_args))
-    args = deserialize_structured_args(args)
+    args = arguments.deserialize_structured_args(args)
     args.setdefault("python_interpreter", sys.executable)
     # Pop this off because it wont be used as a config argument to the whl_library rule.
     requirements_lock = args.pop("requirements_lock")

python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py

@@ -23,12 +23,15 @@ def main() -> None:
     )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
+    deserialized_args = dict(vars(args))
+    arguments.deserialize_structured_args(deserialized_args)
 
     configure_reproducible_wheels()
 
-    pip_args = [sys.executable, "-m", "pip", "--isolated", "wheel", "--no-deps"]
-    if args.extra_pip_args:
-        pip_args += json.loads(args.extra_pip_args)["args"]
+    pip_args = (
+        [sys.executable, "-m", "pip", "--isolated", "wheel", "--no-deps"] +
+        deserialized_args["extra_pip_args"]
+    )
 
     requirement_file = NamedTemporaryFile(mode='wb', delete=False)
     try:
@@ -41,8 +44,10 @@ def main() -> None:
         # so write our single requirement into a temp file in case it has any of those flags.
         pip_args.extend(["-r", requirement_file.name])
 
+        env = os.environ.copy()
+        env.update(deserialized_args["environment"])
         # Assumes any errors are logged by pip so do nothing. This command will fail if pip fails
-        subprocess.run(pip_args, check=True)
+        subprocess.run(pip_args, check=True, env=env)
     finally:
         try:
             os.unlink(requirement_file.name)
@@ -53,16 +58,11 @@ def main() -> None:
     name, extras_for_pkg = requirements._parse_requirement_for_extra(args.requirement)
     extras = {name: extras_for_pkg} if extras_for_pkg and name else dict()
 
-    if args.pip_data_exclude:
-        pip_data_exclude = json.loads(args.pip_data_exclude)["exclude"]
-    else:
-        pip_data_exclude = []
-
     whl = next(iter(glob.glob("*.whl")))
     bazel.extract_wheel(
         whl,
         extras,
-        pip_data_exclude,
+        deserialized_args["pip_data_exclude"],
         args.enable_implicit_namespace_pkgs,
         incremental=True,
         incremental_repo_prefix=bazel.whl_library_repo_prefix(args.repo)

python/pip_install/parse_requirements_to_bzl/parse_requirements_to_bzl_test.py

@@ -23,7 +23,10 @@ def test_generated_requirements_bzl(self) -> None:
             args.requirements_lock = requirements_lock.name
             args.repo = "pip_parsed_deps"
             extra_pip_args = ["--index-url=pypi.org/simple"]
-            args.extra_pip_args = json.dumps({"args": extra_pip_args})
+            pip_data_exclude = ["**.foo"]
+            args.extra_pip_args = json.dumps({"arg": extra_pip_args})
+            args.pip_data_exclude= json.dumps({"arg": pip_data_exclude})
+            args.environment= json.dumps({"arg": {}})
             contents = generate_parsed_requirements_contents(args)
             library_target = "@pip_parsed_deps_pypi__foo//:pkg"
             whl_target = "@pip_parsed_deps_pypi__foo//:whl"
@@ -32,9 +35,11 @@ def test_generated_requirements_bzl(self) -> None:
             self.assertIn(all_requirements, contents, contents)
             self.assertIn(all_whl_requirements, contents, contents)
             self.assertIn(requirement_string, contents, contents)
-            self.assertIn(requirement_string, contents, contents)
             all_flags = extra_pip_args + ["--require-hashes", "True"]
             self.assertIn("'extra_pip_args': {}".format(repr(all_flags)), contents, contents)
+            self.assertIn("'pip_data_exclude': {}".format(repr(pip_data_exclude)), contents, contents)
+            # Assert it gets set to an empty dict by default.
+            self.assertIn("'environment': {}", contents, contents)
 
 
 if __name__ == "__main__":

python/pip_install/pip_repository.bzl

@@ -27,26 +27,38 @@ def _construct_pypath(rctx):
 def _parse_optional_attrs(rctx, args):
     """Helper function to parse common attributes of pip_repository and whl_library repository rules.
 
+    This function also serializes the structured arguments as JSON
+    so they can be passed on the command line to subprocesses.
+
     Args:
         rctx: Handle to the rule repository context.
         args: A list of parsed args for the rule.
     Returns: Augmented args list.
     """
-    if rctx.attr.extra_pip_args:
+
+    # Check for None so we use empty default types from our attrs.
+    # Some args want to be list, and some want to be dict.
+    if rctx.attr.extra_pip_args != None:
         args += [
             "--extra_pip_args",
-            struct(args = rctx.attr.extra_pip_args).to_json(),
+            struct(arg = rctx.attr.extra_pip_args).to_json(),
         ]
 
-    if rctx.attr.pip_data_exclude:
+    if rctx.attr.pip_data_exclude != None:
         args += [
             "--pip_data_exclude",
-            struct(exclude = rctx.attr.pip_data_exclude).to_json(),
+            struct(arg = rctx.attr.pip_data_exclude).to_json(),
         ]
 
     if rctx.attr.enable_implicit_namespace_pkgs:
         args.append("--enable_implicit_namespace_pkgs")
 
+    if rctx.attr.environment != None:
+        args += [
+            "--environment",
+            struct(arg = rctx.attr.environment).to_json(),
+        ]
+
     return args
 
 _BUILD_FILE_CONTENTS = """\
@@ -102,10 +114,8 @@ def _pip_repository_impl(rctx):
 
     result = rctx.execute(
         args,
-        environment = {
-            # Manually construct the PYTHONPATH since we cannot use the toolchain here
-            "PYTHONPATH": pypath,
-        },
+        # Manually construct the PYTHONPATH since we cannot use the toolchain here
+        environment = {"PYTHONPATH": _construct_pypath(rctx)},
         timeout = rctx.attr.timeout,
         quiet = rctx.attr.quiet,
     )
@@ -126,6 +136,16 @@ and py_test targets must specify either `legacy_create_init=False` or the global
 This option is required to support some packages which cannot handle the conversion to pkg-util style.
             """,
     ),
+    "environment": attr.string_dict(
+        doc = """
+Environment variables to set in the pip subprocess.
+Can be used to set common variables such as `http_proxy`, `https_proxy` and `no_proxy`
+Note that pip is run with "--isolated" on the CLI so PIP_<VAR>_<NAME>
+style env vars are ignored, but env vars that control requests and urllib3
+can be passed.
+        """,
+        default = {},
+    ),
     "extra_pip_args": attr.string_list(
         doc = "Extra arguments to pass on to pip. Must not contain spaces.",
     ),
@@ -221,7 +241,6 @@ py_binary(
 def _impl_whl_library(rctx):
     # pointer to parent repo so these rules rerun if the definitions in requirements.bzl change.
     _parent_repo_label = Label("@{parent}//:requirements.bzl".format(parent = rctx.attr.repo))
-    pypath = _construct_pypath(rctx)
     args = [
         rctx.attr.python_interpreter,
         "-m",
@@ -232,12 +251,11 @@ def _impl_whl_library(rctx):
         rctx.attr.repo,
     ]
     args = _parse_optional_attrs(rctx, args)
+
     result = rctx.execute(
         args,
-        environment = {
-            # Manually construct the PYTHONPATH since we cannot use the toolchain here
-            "PYTHONPATH": pypath,
-        },
+        # Manually construct the PYTHONPATH since we cannot use the toolchain here
+        environment = {"PYTHONPATH": _construct_pypath(rctx)},
         quiet = rctx.attr.quiet,
         timeout = rctx.attr.timeout,
     )