Move whl.sh into a Python script that has unit testing.

mattmoor · mattmoor · commit 56be8e80d774 · 2017-09-08T18:24:41.000Z
diff --git a/WORKSPACE b/WORKSPACE
@@ -34,6 +34,33 @@ load("@io_bazel_skydoc//skylark:skylark.bzl", "skydoc_repositories")
 
 skydoc_repositories()
 
+# Test data for WHL tool testing.
+http_file(
+    name = "grpc_whl",
+    sha256 = "c232d6d168cb582e5eba8e1c0da8d64b54b041dd5ea194895a2fe76050916561",
+    # From https://pypi.python.org/pypi/grpcio/1.6.0
+    url = ("https://pypi.python.org/packages/c6/28/" +
+           "67651b4eabe616b27472c5518f9b2aa3f63beab8f62100b26f05ac428639/" +
+           "grpcio-1.6.0-cp27-cp27m-manylinux1_i686.whl"),
+)
+
+http_file(
+    name = "futures_whl",
+    sha256 = "c4884a65654a7c45435063e14ae85280eb1f111d94e542396717ba9828c4337f",
+    # From https://pypi.python.org/pypi/futures
+    url = ("https://pypi.python.org/packages/a6/1c/" +
+           "72a18c8c7502ee1b38a604a5c5243aa8c2a64f4bba4e6631b1b8972235dd/" +
+           "futures-3.1.1-py2-none-any.whl"),
+)
+
+http_file(
+    name = "mock_whl",
+    sha256 = "5ce3c71c5545b472da17b72268978914d0252980348636840bd34a00b5cc96c1",
+    # From https://pypi.python.org/pypi/mock
+    url = ("https://pypi.python.org/packages/e6/35/" +
+           "f187bdf23be87092bd0f1200d43d23076cee4d0dec109f195173fd3ebc79/" +
+           "mock-2.0.0-py2.py3-none-any.whl"),
+)
 
 # Imports for examples
 load("//python:pip.bzl", "pip_import")
diff --git a/python/BUILD b/python/BUILD
@@ -22,3 +22,21 @@ exports_files([
     "whl.bzl",
     "whl.sh",
 ])
+
+load(":python.bzl", "py_library", "py_test")
+
+py_library(
+    name = "whl",
+    srcs = ["whl.py"],
+)
+
+py_test(
+    name = "whl_test",
+    srcs = ["whl_test.py"],
+    data = [
+        "@futures_whl//file",
+        "@grpc_whl//file",
+        "@mock_whl//file",
+    ],
+    deps = [":whl"],
+)
diff --git a/python/whl.bzl b/python/whl.bzl
@@ -17,9 +17,10 @@ def _whl_impl(repository_ctx):
   """Core implementation of whl_library."""
 
   result = repository_ctx.execute([
+    "python",
     repository_ctx.path(repository_ctx.attr._script),
-    repository_ctx.path(repository_ctx.attr.whl),
-    repository_ctx.attr.requirements,
+    "--whl", repository_ctx.path(repository_ctx.attr.whl),
+    "--requirements", repository_ctx.attr.requirements,
   ])
   if result.return_code:
     fail("whl_library failed: %s (%s)" % (result.stdout, result.stderr))
@@ -34,7 +35,7 @@ whl_library = repository_rule(
         "requirements": attr.string(),
         "_script": attr.label(
             executable = True,
-            default = Label("//python:whl.sh"),
+            default = Label("//python:whl.py"),
             cfg = "host",
         ),
     },
diff --git a/python/whl.py b/python/whl.py
@@ -0,0 +1,114 @@
+# Copyright 2017 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The packages modules defines classes for interacting with Python packages."""
+
+import argparse
+import json
+import os
+import zipfile
+
+
+class Wheel(object):
+
+  def __init__(self, path):
+    self._path = path
+
+  def _basename(self):
+      return os.path.basename(self._path)
+
+  def distribution(self):
+    # See https://www.python.org/dev/peps/pep-0427/#file-name-convention
+    parts = self._basename().split('-')
+    return parts[0]
+
+  def version(self):
+    # See https://www.python.org/dev/peps/pep-0427/#file-name-convention
+    parts = self._basename().split('-')
+    return parts[1]
+
+  def _dist_info(self):
+    # Return the name of the dist-info directory within the .whl file.
+    return '{}-{}.dist-info'.format(self.distribution(), self.version())
+
+  def metadata(self):
+    # Extract the structured data from metadata.json in the WHL's dist-info
+    # directory.
+    with zipfile.ZipFile(self._path, 'r') as whl:
+      with whl.open(os.path.join(self._dist_info(), 'metadata.json')) as f:
+        return json.loads(f.read())
+
+  def name(self):
+    return self.metadata().get('name')
+
+  def dependencies(self):
+    # TODO(mattmoor): Is there a schema to follow for this?
+    run_requires = self.metadata().get('run_requires', [])
+    for requirement in run_requires:
+      if 'extra' in requirement:
+        # TODO(mattmoor): What's the best way to support "extras"?
+        # https://packaging.python.org/tutorials/installing-packages/#installing-setuptools-extras
+        continue
+      requires = requirement.get('requires', [])
+      for entry in requires:
+        # Strip off any trailing versioning data.
+        parts = entry.split(' ', 1)
+        yield parts[0]
+
+  def expand(self, directory):
+    with zipfile.ZipFile(self._path, 'r') as whl:
+      whl.extractall(directory)
+
+
+parser = argparse.ArgumentParser(
+    description='Unpack a WHL file as a py_library.')
+
+parser.add_argument('--whl', action='store',
+                    help=('The .whl file we are expanding.'))
+
+parser.add_argument('--requirements', action='store',
+                    help='The pip_import from which to draw dependencies.')
+
+parser.add_argument('--directory', action='store', default='.',
+                    help='The directory into which to expand things.')
+
+def main():
+  args = parser.parse_args()
+  whl = Wheel(args.whl)
+
+  # Extract the files into the current directory
+  whl.expand(args.directory)
+
+  with open(os.path.join(args.directory, 'BUILD'), 'w') as f:
+    f.write("""
+package(default_visibility = ["//visibility:public"])
+
+load("{requirements}", "packages")
+
+py_library(
+  name = "pkg",
+  srcs = glob(["**/*.py"]),
+  data = glob(["**/*"], exclude=["**/*.py", "**/* *", "BUILD", "WORKSPACE"]),
+  # This makes this directory a top-level in the python import
+  # search path for anything that depends on this.
+  imports = ["."],
+  deps = [{dependencies}],
+     )""".format(
+       requirements=args.requirements,
+       dependencies=','.join([
+         'packages("%s")' % d
+         for d in whl.dependencies()
+       ])))
+    
+if __name__ == '__main__':
+  main()
diff --git a/python/whl.sh b/python/whl.sh
diff --git a/python/whl_test.py b/python/whl_test.py
@@ -0,0 +1,54 @@
+# Copyright 2017 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+
+from python import whl
+
+
+def TestData(name):
+  return os.path.join(os.environ['TEST_SRCDIR'], name)
+
+
+class WheelTest(unittest.TestCase):
+
+  def test_grpc_whl(self):
+    td = TestData('grpc_whl/file/grpcio-1.6.0-cp27-cp27m-manylinux1_i686.whl')
+    wheel = whl.Wheel(td)
+    self.assertEqual(wheel.name(), 'grpcio')
+    self.assertEqual(wheel.distribution(), 'grpcio')
+    self.assertEqual(wheel.version(), '1.6.0')
+    self.assertEqual(set(wheel.dependencies()),
+                     set(['enum34', 'futures', 'protobuf', 'six']))
+
+  def test_futures_whl(self):
+    td = TestData('futures_whl/file/futures-3.1.1-py2-none-any.whl')
+    wheel = whl.Wheel(td)
+    self.assertEqual(wheel.name(), 'futures')
+    self.assertEqual(wheel.distribution(), 'futures')
+    self.assertEqual(wheel.version(), '3.1.1')
+    self.assertEqual(set(wheel.dependencies()), set())
+
+  def test_mock_whl(self):
+    td = TestData('mock_whl/file/mock-2.0.0-py2.py3-none-any.whl')
+    wheel = whl.Wheel(td)
+    self.assertEqual(wheel.name(), 'mock')
+    self.assertEqual(wheel.distribution(), 'mock')
+    self.assertEqual(wheel.version(), '2.0.0')
+    self.assertEqual(set(wheel.dependencies()),
+                     set(['pbr', 'six', 'funcsigs']))
+
+if __name__ == '__main__':
+  unittest.main()
