feat(venv): support data, include, and scripts schemes

MODULE.bazel

@@ -246,6 +246,7 @@ use_repo(
     "somepkg_with_build_files",
     "whl_library_extras_direct_dep",
     "whl_with_build_files",
+    "whl_with_data",
 )
 
 dev_rules_python_config = use_extension(

python/config_settings/BUILD.bazel

@@ -229,7 +229,7 @@ string_flag(
 )
 
 config_setting(
-    name = "is_venvs_site_packages",
+    name = "_is_venvs_site_packages_yes",
     flag_values = {
         ":venvs_site_packages": VenvsSitePackages.YES,
     },

python/private/internal_dev_deps.bzl

@@ -91,6 +91,17 @@ def _internal_dev_deps_impl(mctx):
         enable_implicit_namespace_pkgs = False,
     )
 
+    whl_from_dir_repo(
+        name = "whl_with_data_whl",
+        root = "//tests/repos/whl_with_data:BUILD.bazel",
+        output = "whl_with_data-1.0-any-none-any.whl",
+    )
+    whl_library(
+        name = "whl_with_data",
+        whl_file = "@whl_with_data_whl//:whl_with_data-1.0-any-none-any.whl",
+        requirement = "whl-with-data",
+    )
+
     _whl_library_from_dir(
         name = "whl_library_extras_direct_dep",
         root = "//tests/pypi/whl_library/testdata/pkg:BUILD.bazel",

python/private/py_executable.bzl

@@ -576,8 +576,10 @@ def _create_venv(ctx, output_prefix, imports, runtime_details, add_runfiles_root
     )
 
     venv_dir_map = {
-        VenvSymlinkKind.BIN: venv_details.bin_dir,
+        VenvSymlinkKind.BIN: "{}/{}".format(venv_ctx_rel_root, venv_details.bin_dir),
         VenvSymlinkKind.LIB: site_packages,
+        VenvSymlinkKind.INCLUDE: "{}/{}".format(venv_ctx_rel_root, venv_details.include_dir),
+        VenvSymlinkKind.DATA: "{}/data".format(venv_ctx_rel_root),
     }
     venv_app_files = create_venv_app_files(
         ctx,
@@ -659,7 +661,7 @@ def _create_venv_unixy(ctx, *, venv_ctx_rel_root, runtime, interpreter_actual_pa
 
     recreate_venv_at_runtime = False
 
-    bin_dir = "{}/bin".format(venv_ctx_rel_root)
+    venv_bin_ctx_rel_path = "{}/bin".format(venv_ctx_rel_root)
     if create_full_venv:
         # Some wrappers around the interpreter (e.g. pyenv) use the program
         # name to decide what to do, so preserve the name.
@@ -671,15 +673,15 @@ def _create_venv_unixy(ctx, *, venv_ctx_rel_root, runtime, interpreter_actual_pa
             # When the venv symlinks are disabled, the $venv/bin/python3 file isn't
             # needed or used at runtime. However, the zip code uses the interpreter
             # File object to figure out some paths.
-            interpreter = ctx.actions.declare_file("{}/{}".format(bin_dir, py_exe_basename))
+            interpreter = ctx.actions.declare_file("{}/{}".format(venv_bin_ctx_rel_path, py_exe_basename))
             ctx.actions.write(interpreter, "actual:{}".format(interpreter_actual_path))
 
         elif runtime.interpreter:
             # Even though ctx.actions.symlink() is used, using
             # declare_symlink() is required to ensure that the resulting file
             # in runfiles is always a symlink. An RBE implementation, for example,
             # may choose to write what symlink() points to instead.
-            interpreter = ctx.actions.declare_symlink("{}/{}".format(bin_dir, py_exe_basename))
+            interpreter = ctx.actions.declare_symlink("{}/{}".format(venv_bin_ctx_rel_path, py_exe_basename))
             interpreter_runfiles.add(interpreter)
 
             rel_path = relative_path(
@@ -690,7 +692,7 @@ def _create_venv_unixy(ctx, *, venv_ctx_rel_root, runtime, interpreter_actual_pa
             )
             ctx.actions.symlink(output = interpreter, target_path = rel_path)
         else:
-            interpreter = ctx.actions.declare_symlink("{}/{}".format(bin_dir, py_exe_basename))
+            interpreter = ctx.actions.declare_symlink("{}/{}".format(venv_bin_ctx_rel_path, py_exe_basename))
             interpreter_runfiles.add(interpreter)
             ctx.actions.symlink(output = interpreter, target_path = runtime.interpreter_path)
     else:
@@ -715,7 +717,8 @@ def _create_venv_unixy(ctx, *, venv_ctx_rel_root, runtime, interpreter_actual_pa
         interpreter = interpreter,
         pyvenv_cfg = pyvenv_cfg,
         site_packages = site_packages,
-        bin_dir = bin_dir,
+        bin_dir = "bin",
+        include_dir = "include",
         recreate_venv_at_runtime = recreate_venv_at_runtime,
         interpreter_runfiles = interpreter_runfiles.build(ctx),
         interpreter_symlinks = depset(),
@@ -777,7 +780,8 @@ def _create_venv_windows(ctx, *, venv_ctx_rel_root, runtime, interpreter_actual_
         interpreter = interpreter,
         pyvenv_cfg = None,
         site_packages = site_packages,
-        bin_dir = venv_bin_ctx_rel_path,
+        bin_dir = venv_bin_rel_path,
+        include_dir = "Include",
         recreate_venv_at_runtime = True,
         interpreter_runfiles = interpreter_runfiles.build(ctx),
         interpreter_symlinks = interpreter_symlinks.build(),
@@ -789,6 +793,7 @@ def _venv_details(
         pyvenv_cfg,
         site_packages,
         bin_dir,
+        include_dir,
         recreate_venv_at_runtime,
         interpreter_runfiles,
         interpreter_symlinks):
@@ -801,8 +806,10 @@ def _venv_details(
         pyvenv_cfg = pyvenv_cfg,
         # str; venv-relative path to the site-packages directory
         site_packages = site_packages,
-        # str; ctx-relative path to the venv's bin directory.
+        # str; venv-relative path to the venv's bin directory.
         bin_dir = bin_dir,
+        # str; venv-relative-path to the venv's include directory.
+        include_dir = include_dir,
         # bool; True if the venv needs to be recreated at runtime (because the
         # build-time construction isn't sufficient). False if the build-time
         # constructed venv is sufficient.

python/private/py_info.bzl

@@ -45,6 +45,7 @@ VenvSymlinkKind = struct(
     BIN = "BIN",
     LIB = "LIB",
     INCLUDE = "INCLUDE",
+    DATA = "DATA",
 )
 
 def _VenvSymlinkEntry_init(**kwargs):

python/private/pypi/whl_library_targets.bzl

@@ -43,6 +43,8 @@ _BAZEL_REPO_FILE_GLOBS = [
     "WORKSPACE.bazel",
 ]
 
+_IS_VENV_SITE_PACKAGES_YES = Label("//python/config_settings:_is_venvs_site_packages_yes")
+
 def whl_library_targets_from_requires(
         *,
         name,
@@ -194,8 +196,10 @@ def whl_library_targets(
             DIST_INFO_LABEL: dict(
                 include = ["site-packages/*.dist-info/**"],
             ),
+
+            ## TO CHECK: should bin/ and include/ be part of the data target?
             DATA_LABEL: dict(
-                include = ["data/**"],
+                include = ["data/**", "bin/**", "include/**"],
             ),
         }
 
@@ -356,11 +360,15 @@ def whl_library_targets(
             if item not in _data_exclude:
                 _data_exclude.append(item)
 
-        data = data + native.glob(
+        data = data + [":" + DATA_LABEL] + native.glob(
             ["site-packages/**/*"],
             exclude = _data_exclude,
             allow_empty = True,
         )
+        data = data + select({
+            _IS_VENV_SITE_PACKAGES_YES: [DATA_LABEL],
+            "//conditions:default": [],
+        })
 
         pyi_srcs = native.glob(
             ["site-packages/**/*.pyi"],
@@ -369,7 +377,7 @@ def whl_library_targets(
 
         if not enable_implicit_namespace_pkgs:
             generated_namespace_package_files = select({
-                Label("//python/config_settings:is_venvs_site_packages"): [],
+                _IS_VENV_SITE_PACKAGES_YES: [],
                 "//conditions:default": rules.create_inits(
                     srcs = srcs + data + pyi_srcs,
                     ignored_dirnames = [],  # If you need to ignore certain folders, you can patch rules_python here to do so.

python/private/venv_runfiles.bzl

@@ -275,10 +275,16 @@ def _get_file_venv_path(ctx, f, site_packages_root):
 
     Returns:
         A tuple `(venv_path, rf_root_path)` if the file is under
-        `site_packages_root`, otherwise `(None, None)`.
+        `site_packages_root` or data/, bin/, include/ otherwise `(None, None)`.
     """
     rf_root_path = runfiles_root_path(ctx, f.short_path)
     _, _, repo_rel_path = rf_root_path.partition("/")
+
+    # Check for wheel data/bin/include folders first
+    for prefix in ["data/", "bin/", "include/"]:
+        if repo_rel_path.startswith(prefix):
+            return (repo_rel_path, rf_root_path)
+
     head, found_sp_root, venv_path = repo_rel_path.partition(site_packages_root)
     if head or not found_sp_root:
         # If head is set, then the path didn't start with site_packages_root
@@ -452,19 +458,36 @@ def get_venv_symlinks(
 
     # Finally, for each group, we create the VenvSymlinkEntry objects
     for venv_path, files in optimized_groups.items():
+        if venv_path.startswith("data/"):
+            out_venv_path = venv_path[len("data/"):]
+            kind = VenvSymlinkKind.DATA
+            prefix = ""
+        elif venv_path.startswith("include/"):
+            out_venv_path = venv_path[len("include/"):]
+            kind = VenvSymlinkKind.INCLUDE
+            prefix = ""
+        elif venv_path.startswith("bin/"):
+            out_venv_path = venv_path[len("bin/"):]
+            kind = VenvSymlinkKind.BIN
+            prefix = ""
+        else:
+            out_venv_path = venv_path
+            kind = VenvSymlinkKind.LIB
+            prefix = site_packages_root
+
         link_to_path = (
             _get_label_runfiles_repo(ctx, files[0].owner) +
             "/" +
-            site_packages_root +
+            prefix +
             venv_path
         )
         venv_symlinks[venv_path] = VenvSymlinkEntry(
-            kind = VenvSymlinkKind.LIB,
+            kind = kind,
             link_to_path = link_to_path,
             link_to_file = None,
             package = package,
             version = version_str,
-            venv_path = venv_path,
+            venv_path = out_venv_path,
             files = depset(files),
         )
 

tests/pypi/whl_library_targets/whl_library_targets_tests.bzl

@@ -245,7 +245,7 @@ def _test_whl_and_library_deps_from_requires(env):
     env.expect.that_dict(py_library_call).contains_exactly({
         "name": "pkg",
         "srcs": ["site-packages/foo/SRCS.py"] + select({
-            Label("//python/config_settings:is_venvs_site_packages"): [],
+            Label("//python/config_settings:_is_venvs_site_packages_yes"): [],
             "//conditions:default": ["_create_inits_target"],
         }),
         "pyi_srcs": ["site-packages/foo/PYI.pyi"],
@@ -259,7 +259,7 @@ def _test_whl_and_library_deps_from_requires(env):
         "visibility": ["//visibility:public"],
         "experimental_venvs_site_packages": Label("//python/config_settings:venvs_site_packages"),
         "namespace_package_files": [] + select({
-            Label("//python/config_settings:is_venvs_site_packages"): [],
+            Label("//python/config_settings:_is_venvs_site_packages_yes"): [],
             "//conditions:default": ["_create_inits_target"],
         }),
     })  # buildifier: @unsorted-dict-items
@@ -361,7 +361,7 @@ def _test_whl_and_library_deps(env):
     env.expect.that_dict(py_library_calls[0]).contains_exactly({
         "name": "pkg",
         "srcs": ["site-packages/foo/SRCS.py"] + select({
-            Label("//python/config_settings:is_venvs_site_packages"): [],
+            Label("//python/config_settings:_is_venvs_site_packages_yes"): [],
             "//conditions:default": ["_create_inits_target"],
         }),
         "pyi_srcs": ["site-packages/foo/PYI.pyi"],
@@ -386,7 +386,7 @@ def _test_whl_and_library_deps(env):
         "visibility": ["//visibility:public"],
         "experimental_venvs_site_packages": Label("//python/config_settings:venvs_site_packages"),
         "namespace_package_files": [] + select({
-            Label("//python/config_settings:is_venvs_site_packages"): [],
+            Label("//python/config_settings:_is_venvs_site_packages_yes"): [],
             "//conditions:default": ["_create_inits_target"],
         }),
     })  # buildifier: @unsorted-dict-items
@@ -444,7 +444,7 @@ def _test_group(env):
     ).contains_exactly({
         "name": "_pkg",
         "srcs": ["site-packages/foo/srcs.py"] + select({
-            Label("//python/config_settings:is_venvs_site_packages"): [],
+            Label("//python/config_settings:_is_venvs_site_packages_yes"): [],
             "//conditions:default": ["_create_inits_target"],
         }),
         "pyi_srcs": ["site-packages/foo/pyi.pyi"],
@@ -459,7 +459,7 @@ def _test_group(env):
         "visibility": ["@pypi__config//_groups:__pkg__"],
         "experimental_venvs_site_packages": Label("//python/config_settings:venvs_site_packages"),
         "namespace_package_files": [] + select({
-            Label("//python/config_settings:is_venvs_site_packages"): [],
+            Label("//python/config_settings:_is_venvs_site_packages_yes"): [],
             "//conditions:default": ["_create_inits_target"],
         }),
     })  # buildifier: @unsorted-dict-items

tests/repos/whl_with_data/BUILD.bazel

@@ -0,0 +1 @@
+exports_files(glob(["*"]))

tests/repos/whl_with_data/whl_with_data-1.0.data/data/whl_with_data/data_data_file.txt

@@ -0,0 +1 @@
+from .data/data

tests/repos/whl_with_data/whl_with_data-1.0.data/headers/whl_with_data/header_file.h

@@ -0,0 +1 @@
+from .data/headers

tests/repos/whl_with_data/whl_with_data-1.0.data/platlib/whl_with_data/platlib_file.txt

@@ -0,0 +1 @@
+from .data/platlib

tests/repos/whl_with_data/whl_with_data-1.0.data/purelib/whl_with_data/data_file.txt

@@ -0,0 +1 @@
+from .data

tests/repos/whl_with_data/whl_with_data-1.0.data/scripts/whl_script.sh

@@ -0,0 +1 @@
+#!/bin/sh

tests/repos/whl_with_data/whl_with_data-1.0.dist-info/METADATA

@@ -0,0 +1,3 @@
+Metadata-Version: 2.1
+Name: whl-with-data
+Version: 1.0

tests/repos/whl_with_data/whl_with_data-1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+whl_with_data-1.0.data/platlib/whl_with_data/platlib_file.txt,sha256=123,123
+whl_with_data-1.0.data/scripts/whl_script.sh,sha256=123,123
+whl_with_data-1.0.data/headers/whl_with_data/header_file.h,sha256=123,123
+whl_with_data-1.0.data/purelib/whl_with_data/data_file.txt,sha256=123,123
+whl_with_data-1.0.data/data/whl_with_data/data_data_file.txt,sha256=123,123
+whl_with_data-1.0.data/data/whl_with_data/data_data_file.txt,sha256=123,123

tests/repos/whl_with_data/whl_with_data-1.0.dist-info/WHEEL

@@ -0,0 +1 @@
+Wheel-Version: 1.0

tests/venv_site_packages_libs/BUILD.bazel

@@ -45,6 +45,7 @@ py_reconfig_test(
         "@other//nspkg_gamma",
         "@other//nspkg_single",
         "@other//with_external_data",
+        "@whl_with_data//:pkg",
     ],
 )
 

tests/venv_site_packages_libs/bin.py

@@ -70,7 +70,7 @@ def test_dirs_from_replaced_package_are_not_present(self):
         module_path = Path(module.__file__)
 
         site_packages = module_path.parent.parent
-        dist_info_dirs = [p.name for p in site_packages.glob("*.dist-info")]
+        dist_info_dirs = [p.name for p in site_packages.glob("simple*.dist-info")]
         self.assertEqual(
             ["simple-1.0.0.dist-info"],
             dist_info_dirs,
@@ -91,6 +91,33 @@ def test_data_from_another_pkg_is_included_via_copy_file(self):
         files = [p.name for p in d.glob("*")]
         self.assertIn("another_module_data.txt", files)
 
+    def test_whl_with_data_included(self):
+        module = self.assert_imported_from_venv("whl_with_data")
+        module_path = Path(module.__file__)
+        site_packages = module_path.parent.parent
+
+        # purelib
+        data_file = site_packages / "whl_with_data" / "data_file.txt"
+        self.assertTrue(data_file.exists(), f"Expected {data_file} to exist")
+
+        # platlib
+        platlib_file = site_packages / "whl_with_data" / "platlib_file.txt"
+        self.assertTrue(platlib_file.exists(), f"Expected {platlib_file} to exist")
+
+        venv_root = Path(self.venv)
+
+        # data
+        data_data_file = venv_root / "data" / "whl_with_data" / "data_data_file.txt"
+        self.assertTrue(data_data_file.exists(), f"Expected {data_data_file} to exist")
+
+        # scripts
+        script_file = venv_root / "bin" / "whl_script.sh"
+        self.assertTrue(script_file.exists(), f"Expected {script_file} to exist")
+
+        # headers
+        header_file = venv_root / "include" / "whl_with_data" / "header_file.h"
+        self.assertTrue(header_file.exists(), f"Expected {header_file} to exist")
+
 
 if __name__ == "__main__":
     unittest.main()