Merge pull request #78 from takluyver/smarter-extract-wheel

Improve how wheels are extracted

nsist/pypi.py

@@ -2,7 +2,10 @@ from distutils.version import LooseVersion
 import errno
 import hashlib
 import logging
+from pathlib import Path
 import re
+import shutil
+from tempfile import mkdtemp
 import zipfile
 
 import yarg
@@ -129,13 +132,65 @@ class CachedRelease(object):
         self.filename = filename
         self.package_type = 'wheel' if filename.endswith('.whl') else ''
 
+def merge_dir_to(src, dst):
+    """Merge all files from one directory into another.
+
+    Subdirectories will be merged recursively. If filenames are the same, those
+    from src will overwrite those in dst. If a regular file clashes with a
+    directory, an error will occur.
+    """
+    for p in src.iterdir():
+        if p.is_dir():
+            dst_p = dst / p.name
+            if dst_p.is_dir():
+                merge_dir_to(p, dst_p)
+            elif dst_p.is_file():
+                raise RuntimeError('Directory {} clashes with file {}'
+                                   .format(p, dst_p))
+            else:
+                shutil.copytree(str(p), str(dst_p))
+        else:
+            # Copy regular file
+            dst_p = dst / p.name
+            if dst_p.is_dir():
+                raise RuntimeError('File {} clashes with directory {}'
+                                   .format(p, dst_p))
+            shutil.copy2(str(p), str(dst_p))
+
 def extract_wheel(whl_file, target_dir):
+    """Extract importable modules from a wheel to the target directory
+    """
+    # Extract to temporary directory
+    td = Path(mkdtemp())
     with zipfile.ZipFile(str(whl_file), mode='r') as zf:
-        names = zf.namelist()
-        # TODO: Do anything with data and dist-info folders?
-        pkg_files = [n for n in names \
-                     if not n.split('/')[0].endswith(('.data', '.dist-info'))]
-        zf.extractall(target_dir, members=pkg_files)
+        zf.extractall(str(td))
+
+    # Move extra lib files out of the .data subdirectory
+    for p in td.iterdir():
+        if p.suffix == '.data':
+            if (p / 'purelib').is_dir():
+                merge_dir_to(p / 'purelib', td)
+            if (p / 'platlib').is_dir():
+                merge_dir_to(p / 'platlib', td)
+
+    # Copy to target directory
+    target = Path(target_dir)
+    copied_something = False
+    for p in td.iterdir():
+        if p.suffix not in {'.data', '.dist-info'}:
+            if p.is_dir():
+                shutil.copytree(str(p), str(target / p.name))
+            else:
+                shutil.copy2(str(p), str(target))
+            copied_something = True
+
+    if not copied_something:
+        raise RuntimeError("Did not find any files to extract from wheel {}"
+                           .format(whl_file))
+
+    # Clean up temporary directory
+    shutil.rmtree(str(td))
+
 
 def fetch_pypi_wheels(requirements, target_dir, py_version, bitness):
     for req in requirements:

nsist/tests/test_pypi.py

@@ -1,9 +1,10 @@
 from nose.tools import *
 from os.path import join as pjoin
+from pathlib import Path
 from testpath import assert_isfile, assert_isdir
 from testpath.tempdir import TemporaryDirectory
 
-from nsist.pypi import WheelDownloader, extract_wheel, CachedRelease
+from nsist.pypi import WheelDownloader, extract_wheel, CachedRelease, merge_dir_to
 
 def test_download():
     wd = WheelDownloader("astsearch==0.1.2", "3.5.1", 64)
@@ -78,3 +79,33 @@ def test_pick_best_wheel():
         CachedRelease('astsearch-0.1.2-py2.py3-none-win_amd64.whl'),
     ]
     assert_equal(wd.pick_best_wheel(releases), releases[1])
+
+def test_merge_dir_to():
+    with TemporaryDirectory() as td1, TemporaryDirectory() as td2:
+        td1 = Path(td1)
+        td2 = Path(td2)
+        with (td1 / 'ab').open('w') as f:
+            f.write(u"original")
+        with (td2 / 'ab').open('w') as f:
+            f.write(u"alternate")
+
+        (td1 / 'subdir').mkdir()
+        with (td1 / 'subdir' / 'foo').open('w'): pass
+        (td2 / 'subdir').mkdir()
+        with (td2 / 'subdir' / 'bar').open('w'): pass
+
+        merge_dir_to(td2, td1)
+
+        assert_isfile(td1 / 'subdir' / 'foo')
+        assert_isfile(td1 / 'subdir' / 'bar')
+        with (td1 / 'ab').open() as f:
+            assert_equal(f.read(), u"alternate")
+
+        # Test with conflicts
+        (td1 / 'conflict').mkdir()
+        with (td2 / 'conflict').open('w'): pass
+
+        with assert_raises(RuntimeError):
+            merge_dir_to(td2, td1)
+        with assert_raises(RuntimeError):
+            merge_dir_to(td1, td2)