Merge pull request #6 from dillon-giacoppo/jonathon/support-purelib-platlib

Author: Jonathon Belotti

src/BUILD

@@ -9,5 +9,12 @@ py_test(
     name = "test",
     srcs = glob(["tests/*.py"]),
     main = "tests/__main__.py",
-    deps = [":src"],
+    deps = [
+        ":src",
+        "//third_party/python:pkginfo",
+    ],
+    size="small",
+    imports = [
+        "../third_party/python/",
+    ]
 )

src/extract_wheels.py

@@ -6,6 +6,7 @@
 
 
 from . import namespace_pkgs
+from . import purelib
 from .wheel import Wheel
 
 BUILD_TEMPLATE = """\
@@ -78,6 +79,8 @@ def extract_wheel(whl, directory, extras):
 
     whl.unzip(directory)
 
+    # Note: Order of operations matters here
+    purelib.spread_purelib_into_root(directory)
     _setup_namespace_pkg_compatibility(directory)
 
     with open(os.path.join(directory, "BUILD"), "w") as f:

src/namespace_pkgs.py

@@ -1,7 +1,7 @@
 import os
-import glob
 import sys
 
+from . import wheel
 from typing import Set
 
 
@@ -21,13 +21,7 @@ def pkg_resources_style_namespace_packages(extracted_whl_directory) -> Set[str]:
     """
     namespace_pkg_dirs = set()
 
-    dist_info_dirs = glob.glob(os.path.join(extracted_whl_directory, "*.dist-info"))
-    if not dist_info_dirs:
-        raise ValueError(f"No *.dist-info directory found. {extracted_whl_directory} is not a valid Wheel.")
-    elif len(dist_info_dirs) > 1:
-        raise ValueError(f"Found more than 1 *.dist-info directory. {extracted_whl_directory} is not a valid Wheel.")
-    else:
-        dist_info = dist_info_dirs[0]
+    dist_info = wheel.get_dist_info(extracted_whl_directory)
     namespace_packages_record_file = os.path.join(dist_info, "namespace_packages.txt")
     if os.path.exists(namespace_packages_record_file):
         with open(namespace_packages_record_file) as nspkg:

src/purelib.py

@@ -0,0 +1,33 @@
+import os
+import pathlib
+import shutil
+
+from . import wheel
+
+
+def spread_purelib_into_root(extracted_whl_directory: str) -> None:
+    dist_info = wheel.get_dist_info(extracted_whl_directory)
+    wheel_metadata_file_path = pathlib.Path(dist_info, "WHEEL")
+    wheel_metadata_dict = wheel.parse_WHEEL_file(str(wheel_metadata_file_path))
+
+    if "Root-Is-Purelib" not in wheel_metadata_dict:
+        raise ValueError(f"Invalid WHEEL file '{wheel_metadata_file_path}'. Expected key 'Root-Is-Purelib'.")
+    root_is_purelib = wheel_metadata_dict["Root-Is-Purelib"]
+
+    if root_is_purelib.lower() == "true":
+        # The Python package code is in the root of the Wheel, so no need to 'spread' anything.
+        return
+
+    dot_data_dir = wheel.get_dot_data_directory(extracted_whl_directory)
+
+    # 'Root-Is-Purelib: false' is no guarantee a .date directory exists with
+    # package code in it. eg. the 'markupsafe' package.
+    if dot_data_dir:
+        for child in pathlib.Path(dot_data_dir).iterdir():
+            # TODO(Jonathon): Should all other potential folders get ignored? eg. 'platlib'
+            if str(child).endswith("purelib"):
+                for grandchild in child.iterdir():
+                    shutil.move(
+                        src=str(grandchild),
+                        dst=extracted_whl_directory,
+                    )

src/tests/test_namespace_pkgs.py

@@ -1,4 +1,3 @@
-import os
 import pathlib
 import unittest
 import tempfile

src/wheel.py

@@ -1,6 +1,10 @@
+import os
 import pkginfo
 import zipfile
 import pkg_resources
+import glob
+
+from typing import Dict, Optional
 
 
 class Wheel(object):
@@ -37,3 +41,42 @@ def dependencies(self, extras_requested=None):
     def unzip(self, directory):
         with zipfile.ZipFile(self.path(), "r") as whl:
             whl.extractall(directory)
+
+
+def get_dist_info(extracted_whl_directory) -> str:
+    dist_info_dirs = glob.glob(os.path.join(extracted_whl_directory, "*.dist-info"))
+    if not dist_info_dirs:
+        raise ValueError(f"No *.dist-info directory found. {extracted_whl_directory} is not a valid Wheel.")
+    elif len(dist_info_dirs) > 1:
+        raise ValueError(f"Found more than 1 *.dist-info directory. {extracted_whl_directory} is not a valid Wheel.")
+    else:
+        dist_info = dist_info_dirs[0]
+    return dist_info
+
+
+def get_dot_data_directory(extracted_whl_directory) -> Optional[str]:
+    # See: https://www.python.org/dev/peps/pep-0491/#the-data-directory
+    dot_data_dirs = glob.glob(os.path.join(extracted_whl_directory, "*.data"))
+    if not dot_data_dirs:
+        return None
+    elif len(dot_data_dirs) > 1:
+        raise ValueError(f"Found more than 1 *.data directory. {extracted_whl_directory} is not a valid Wheel.")
+    else:
+        dot_data_dir = dot_data_dirs[0]
+    return dot_data_dir
+
+
+def parse_WHEEL_file(whl_file_path: str) -> Dict[str, str]:
+    contents = {}
+    with open(whl_file_path, "r") as f:
+        for line in f:
+            cleaned = line.strip()
+            if not cleaned:
+                continue
+            try:
+                key, value = cleaned.split(":", maxsplit=1)
+                contents[key] = value.strip()
+            except ValueError:
+                raise RuntimeError(f"Encounted invalid line in WHEEL file: '{cleaned}'")
+    return contents
+