The pip_lock rule now allows for custom BUILD files in the local whee…

…ls package

src/bin/lock_pip_requirements.py

@@ -4,7 +4,7 @@
 import sys
 import textwrap
 
-from piprules import lockfile, pipcompat, requirements, resolve
+from piprules import localwheels, lockfile, pipcompat, requirements, resolve
 
 
 LOG = logging.getLogger()
@@ -42,22 +42,26 @@ def main():
 
     lock_file.update_requirements_for_current_environment(resolved_requirements)
 
-    new_local_wheels_package = "@{workspace}//{package}".format(
+    new_local_wheels_package_label = "@{workspace}//{package}".format(
         workspace=args.workspace_name,
         package=args.wheel_dir,
     )
-    if new_local_wheels_package != lock_file.local_wheels_package:
+    if new_local_wheels_package_label != lock_file.local_wheels_package:
         LOG.warning(
             "Local wheels package is changing from '%s' to '%s'",
             lock_file.local_wheels_package,
-            new_local_wheels_package,
+            new_local_wheels_package_label,
         )
-        lock_file.local_wheels_package = new_local_wheels_package
+        lock_file.local_wheels_package = new_local_wheels_package_label
 
     if os.path.isdir(wheel_directory):
-        create_local_wheel_package_build_file(wheel_directory)
-
-    purge_unused_local_wheels(lock_file, wheel_directory)
+        local_wheels_package = localwheels.Package(wheel_directory)
+        local_wheels_package.ensure_build_file_exists()
+        local_wheels_package.purge_wheels(keep={
+            source.file
+            for source in lock_file.sources.values()
+            if source.file
+        })
 
     if lock_file_path:
         lock_file.dump(lock_file_path)
@@ -135,36 +139,5 @@ def get_workspace_directory():
         sys.exit("This tool must by executed via 'bazel run'")
 
 
-def create_local_wheel_package_build_file(wheel_directory):
-    LOG.info("Creating local wheel package in %s", wheel_directory)
-
-    path = os.path.join(wheel_directory, "BUILD")
-
-    with open(path, mode="w") as build_file:
-        build_file.write(textwrap.dedent("""
-            package(default_visibility = ["//visibility:public"])
-
-            exports_files(glob(["*.whl"]))
-
-            filegroup(
-                name = "wheels",
-                srcs = glob(["*.whl"]),
-            )
-        """).lstrip())
-
-
-def purge_unused_local_wheels(lock_file, wheel_directory):
-    local_wheels = {
-        source.file
-        for source in lock_file.sources.values()
-        if source.file
-    }
-
-    for filename in os.listdir(wheel_directory):
-        if filename.endswith(".whl") and filename not in local_wheels:
-            LOG.info("Removing unused local wheel %s", filename)
-            os.remove(os.path.join(wheel_directory, filename))
-
-
 if __name__ == "__main__":
     main()

src/piprules/localwheels.py

@@ -0,0 +1,52 @@
+import logging
+import os
+
+
+LOG = logging.getLogger(__name__)
+
+
+class Package(object):
+
+    _VALID_BUILD_FILE_NAMES = {"BUILD", "BUILD.bazel"}
+    _STUB_BUILD_FILE_COMMENT = (
+        "# This is a generated file which may be overwritten with a custom BUILD file"
+    )
+
+    def __init__(self, directory):
+        self.directory = directory
+
+    @property
+    def _default_build_file_path(self):
+        return os.path.join(self.directory, "BUILD")
+
+    def ensure_build_file_exists(self):
+        if not self._does_build_file_exist():
+            self._generate_stub_build_file()
+
+    def _does_build_file_exist(self):
+        return any(
+            filename in self._VALID_BUILD_FILE_NAMES
+            for filename in self._list_files()
+        )
+
+    def _list_files(self):
+        return os.listdir(self.directory)
+
+    def _generate_stub_build_file(self):
+        LOG.info("Generating stub BUILD file for local wheels package")
+        with open(self._default_build_file_path, mode="w") as build_file:
+            build_file.write(self._STUB_BUILD_FILE_COMMENT)
+
+    def purge_wheels(self, keep=None):
+        if keep is None:
+            keep = set()
+
+        for filename in self._iterate_wheels():
+            if filename not in keep:
+                LOG.info("Removing unused local wheel %s", filename)
+                os.remove(os.path.join(self.directory, filename))
+
+    def _iterate_wheels(self):
+        for filename in self._list_files():
+            if filename.endswith(".whl"):
+                yield filename

thirdparty/pip/wheels/BUILD

@@ -1 +1 @@
-exports_files(["*"])
+# This is a generated file which may be overwritten with a custom BUILD file