Install Python packages data files to venv (cedarai#12)

* Install Python packages data files to venv

Author: kevinpark1217

.bazelignore

@@ -0,0 +1 @@
+example

build_env.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import collections
+from enum import Enum
 import itertools
 import json
 import os
@@ -21,12 +21,20 @@
 import sys
 import textwrap
 import venv
-from typing import Dict, List, Optional
+from typing import Dict, List, NamedTuple
 
 import importlib_metadata
 
 
-EnvFile = collections.namedtuple("EnvFile", ["path", "site_packages_path"])
+class EnvPathType(Enum):
+    SITE_PACKAGES = 1
+    DATA = 2
+
+
+class EnvFile(NamedTuple):
+    path: pathlib.Path
+    env_path: pathlib.Path
+    type_: EnvPathType = EnvPathType.SITE_PACKAGES
 
 
 def console_script(env_path: pathlib.Path, module: str, func: str) -> str:
@@ -48,9 +56,9 @@ def path_starts_with(path: pathlib.Path, prefix: pathlib.Path) -> bool:
     return path.parts[: len(prefix.parts)] == prefix.parts
 
 
-def get_site_packages_path(
+def get_env_path(
     workspace: str, path: pathlib.Path, imports: List[pathlib.Path]
-) -> pathlib.Path:
+) -> EnvFile:
     # Import prefixes start with the workspace name, which might be the local workspace.
     # We first normalize the given path so that it starts with its workspace name.
     if path.parts[0] == "..":
@@ -62,20 +70,24 @@ def get_site_packages_path(
 
     for imp in imports:
         if path_starts_with(wspath, imp):
-            return wspath.relative_to(imp)
+            return EnvFile(path, wspath.relative_to(imp))
+
+        imp_data = imp.parent / "data"
+        if path_starts_with(wspath, imp_data):
+            return EnvFile(path, wspath.relative_to(imp_data), EnvPathType.DATA)
 
     if not is_external:
         # If the input wasn't an external path and it didn't match any import prefixes,
         # just return it as given.
-        return path
+        return EnvFile(path, path)
 
     # External file that didn't match imports. Include but warn.
     # We include it as relative to its workspace directory, so strip the first component
     # off wspath.
     include_path = wspath.relative_to(wspath.parts[0])
     print(f"Warning: [{path}] didn't match any imports. Including as [{include_path}]")
 
-    return include_path
+    return EnvFile(path, include_path)
 
 
 def is_external(file_: pathlib.Path) -> bool:
@@ -123,27 +135,32 @@ def get_files(build_env_input: Dict) -> List[EnvFile]:
             paths = [input_path]
 
         for path in paths:
-            site_packages_path = get_site_packages_path(workspace, path, imports)
-            if site_packages_path != path or always_link:
-                files.append(EnvFile(path, site_packages_path))
+            env_file = get_env_path(workspace, path, imports)
+            if env_file.env_path != path or always_link:
+                files.append(env_file)
 
     return files
 
 
 def is_data_file(file: EnvFile) -> bool:
-    return file.site_packages_path.parts[0].endswith(".data")
+    return (
+        file.type_ == EnvPathType.DATA
+        or file.env_path.parts[0].endswith(".data")
+    )
 
 
 def install_data_file(env_path: pathlib.Path, file: EnvFile) -> None:
     if (
-        len(file.site_packages_path.parts) > 2
-        and file.site_packages_path.parts[1] == "scripts"
+        len(file.env_path.parts) > 2
+        and file.env_path.parts[1] == "scripts"
     ):
         install_included_script(env_path, file.path)
+    elif file.type_ == EnvPathType.DATA:
+        install_site_file(env_path, file)
 
 
 def install_site_file(site_packages_path: pathlib.Path, file: EnvFile) -> None:
-    site_path = site_packages_path / file.site_packages_path
+    site_path = site_packages_path / file.env_path
     if not site_path.exists():
         site_path.parent.mkdir(parents=True, exist_ok=True)
         site_path.symlink_to(file.path.resolve())

example/BUILD.bazel

@@ -15,16 +15,19 @@
 package(default_visibility = ["//visibility:public"])
 
 load("@rules_pyvenv//:venv.bzl", "py_venv")
-load("@example_deps//:requirements.bzl", "requirement")
+load("@example_deps//:requirements.bzl", "requirement", "data_requirement")
 load("@rules_python//python/pip_install:requirements.bzl", "compile_pip_requirements")
 
 py_venv(
     name = "venv",
     deps = [
-	"//libraries/liba",
+        "//libraries/liba",
         requirement("black"),
         requirement("numpy"),
     ],
+    data = [
+        data_requirement("black"),
+    ],
     extra_pip_commands = [
 #        "install ipython jupyter-console",
     ]

example/WORKSPACE

@@ -19,12 +19,14 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 # Required rules_python
 http_archive(
     name = "rules_python",
-    sha256 = "8c8fe44ef0a9afc256d1e75ad5f448bb59b81aba149b8958f02f7b3a98f5d9b4",
-    strip_prefix = "rules_python-0.13.0",
-    url = "https://github.com/bazelbuild/rules_python/archive/refs/tags/0.13.0.tar.gz",
+    sha256 = "0a8003b044294d7840ac7d9d73eef05d6ceb682d7516781a4ec62eeb34702578",
+    strip_prefix = "rules_python-0.24.0",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.24.0/rules_python-0.24.0.tar.gz",
 )
 
-load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
+
+py_repositories()
 
 python_register_toolchains(
     name = "python3_10",

example/requirements.txt

@@ -1,6 +1,6 @@
 #
-# This file is autogenerated by pip-compile with python 3.10
-# To update, run:
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
 #
 #    bazel run //:requirements.update
 #
@@ -26,7 +26,7 @@ black==22.10.0 \
     --hash=sha256:e41a86c6c650bcecc6633ee3180d80a025db041a8e2398dcc059b3afa8382cd4 \
     --hash=sha256:f513588da599943e0cde4e32cc9879e825d58720d6557062d1098c5ad80080e1 \
     --hash=sha256:fba8a281e570adafb79f7755ac8721b6cf1bbf691186a287e990c7929c7692ff
-    # via -r ./requirements.in
+    # via -r requirements.in
 click==8.1.3 \
     --hash=sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e \
     --hash=sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48
@@ -64,7 +64,7 @@ numpy==1.23.3 \
     --hash=sha256:efd9d3abe5774404becdb0748178b48a218f1d8c44e0375475732211ea47c67e \
     --hash=sha256:f8c02ec3c4c4fcb718fdf89a6c6f709b14949408e8cf2a2be5bfa9c49548fd85 \
     --hash=sha256:ffcf105ecdd9396e05a8e58e81faaaf34d3f9875f137c7372450baa5d77c9a54
-    # via -r ./requirements.in
+    # via -r requirements.in
 pathspec==0.10.1 \
     --hash=sha256:46846318467efc4556ccfd27816e004270a9eeeeb4d062ce5e6fc7a87c573f93 \
     --hash=sha256:7ace6161b621d31e7902eb6b5ae148d12cfd23f4a249b9ffb6b9fee12084323d

venv.bzl

@@ -27,10 +27,11 @@ def _py_venv_deps_impl(ctx):
         imports.extend([i for i in dep[PyInfo].imports.to_list() if i not in imports])
 
     deps = depset(transitive = [dep[DefaultInfo].default_runfiles.files for dep in ctx.attr.deps])
+    data = depset(transitive = [data[DefaultInfo].files for data in ctx.attr.data])
     out = ctx.outputs.output
 
     files = []
-    for dep in deps.to_list():
+    for dep in deps.to_list() + data.to_list():
         # Skip files that are provided by the python toolchain.
         # They don't need to be in the venv.
         if dep in toolchain_files:
@@ -54,15 +55,17 @@ _py_venv_deps = rule(
     implementation = _py_venv_deps_impl,
     attrs = {
         "deps": attr.label_list(),
+        "data": attr.label_list(),
         "commands": attr.string_list(),
         "always_link": attr.bool(),
         "output": attr.output(),
     },
     toolchains = [PYTHON_TOOLCHAIN_TYPE],
 )
 
-def py_venv(name, deps = None, extra_pip_commands = None, always_link = False, venv_location = None, **kwargs):
+def py_venv(name, deps = None, data = None, extra_pip_commands = None, always_link = False, venv_location = None, **kwargs):
     deps = deps or []
+    data = data or []
     extra_pip_commands = extra_pip_commands or []
 
     deps_name = "_" + name + "_deps"
@@ -71,6 +74,7 @@ def py_venv(name, deps = None, extra_pip_commands = None, always_link = False, v
     _py_venv_deps(
         name = deps_name,
         deps = deps,
+        data = data,
         commands = extra_pip_commands,
         always_link = always_link,
         output = out_name,
@@ -88,7 +92,7 @@ def py_venv(name, deps = None, extra_pip_commands = None, always_link = False, v
         name = name,
         srcs = ["@rules_pyvenv//:build_env.py"],
         deps = ["@rules_pyvenv//vendor/importlib_metadata"],
-        data = [out_label] + deps,
+        data = [out_label] + deps + data,
         main = "@rules_pyvenv//:build_env.py",
         env = env,
         **kwargs,