Repository: takluyver/flit
Pull request #67 (2016-01-20 13:30:04)

Refactor WheelBuilder not to work out its own target file
With this refactoring, WheelBuilder is no longer responsible for working out its own target file. A new wheel_main function gives it a temporary file to write to, and renames that afterwards (because we don't know the final name until we've processed the metadata).

Besides general loose-coupling principles, this should make it easier for `flit install` to work by building a wheel and calling pip with it: that machinery will use WheelBuilder directly to make the wheel in a temporary location.

I'm not sure whether WheelBuilder should get:
1. The path to a file it should open and write
2. A writable file object that it should write zip data to
3. A ZipFile object that it should write into.

Currently, this does 2, because it works neatly with mkstemp. Any clear advantages to the other options?


Patch:
diff --git a/flit/__init__.py b/flit/__init__.py
index be17e54..2304992 100644
--- a/flit/__init__.py
+++ b/flit/__init__.py
@@ -78,11 +78,11 @@ def main(argv=None):
         sys.exit(0)
 
     if args.subcmd == 'wheel':
-        from .wheel import WheelBuilder
+        from .wheel import wheel_main
         try:
-            WheelBuilder(args.ini_file, upload=args.upload,
+            wheel_main(args.ini_file, upload=args.upload,
                      verify_metadata=args.verify_metadata,
-                     repo=args.repository).build()
+                     repo=args.repository)
         except common.ProblemInModule as e:
             sys.exit(e.args[0])
     elif args.subcmd == 'install':
diff --git a/flit/wheel.py b/flit/wheel.py
index 1bd3655..bf173cd 100644
--- a/flit/wheel.py
+++ b/flit/wheel.py
@@ -1,5 +1,6 @@
 import configparser
 import contextlib
+import tempfile
 from datetime import datetime
 import hashlib
 import io
@@ -26,23 +27,15 @@ class EntryPointsConflict(ValueError):
             'flit.ini, not both.')
 
 class WheelBuilder:
-    def __init__(self, ini_path, upload=False, verify_metadata=False, repo='pypi'):
+    def __init__(self, ini_path, target_fp):
         """Build a wheel from a module/package
         """
         self.ini_path = ini_path
         self.directory = ini_path.parent
-        self.dist_dir = self.directory / 'dist'
-        try:
-            self.dist_dir.mkdir()
-        except FileExistsError:
-            pass
 
         self.ini_info = inifile.read_pkg_ini(ini_path)
         self.module = common.Module(self.ini_info['module'], ini_path.parent)
         self.metadata = common.make_metadata(self.module, self.ini_info)
-        self.upload=upload
-        self.verify_metadata=verify_metadata
-        self.repo = repo
 
         self.dist_version = self.metadata.name + '-' + self.metadata.version
         self.records = []
@@ -59,13 +52,13 @@ class WheelBuilder:
             self.source_time_stamp = None
 
         # Open the zip file ready to write
-        self.wheel_zip = zipfile.ZipFile(str(self.wheel_path), 'w',
+        self.wheel_zip = zipfile.ZipFile(target_fp, 'w',
                              compression=zipfile.ZIP_DEFLATED)
 
     @property
-    def wheel_path(self):
+    def wheel_filename(self):
         tag = ('py2.' if self.supports_py2 else '') + 'py3-none-any'
-        return self.dist_dir / '{}-{}-{}.whl'.format(
+        return '{}-{}-{}.whl'.format(
                 re.sub("[^\w\d.]+", "_", self.metadata.name, re.UNICODE),
                 re.sub("[^\w\d.]+", "_", self.metadata.version, re.UNICODE),
                 tag)
@@ -98,6 +91,7 @@ class WheelBuilder:
         self.records.append((rel_path, hashsum.hexdigest(), len(b)))
 
     def copy_module(self):
+        log.info('Copying package file(s) from %s', self.module.path)
         if self.module.is_package:
             # Walk the tree and compress it, sorting everything so the order
             # is stable.
@@ -118,6 +112,7 @@ class WheelBuilder:
                                     .startswith(('3', '>3', '>=3'))
 
     def write_metadata(self):
+        log.info('Writing metadata files')
         dist_info = self.dist_version + '.dist-info'
 
         # Write entry points
@@ -149,6 +144,7 @@ class WheelBuilder:
             self.metadata.write_metadata_file(f)
 
     def write_record(self):
+        log.info('Writing the record of files')
         # Write a record of the files in the wheel
         with self._write_to_zip(self.dist_version + '.dist-info/RECORD') as f:
             for path, hash, size in self.records:
@@ -156,15 +152,6 @@ class WheelBuilder:
             # RECORD itself is recorded with no hash or size
             f.write(self.dist_version + '.dist-info/RECORD,,\n')
 
-    def post_build(self):
-        if self.verify_metadata:
-            from .upload import verify
-            verify(self.metadata, self.repo)
-
-        if self.upload:
-            from .upload import do_upload
-            do_upload(self.wheel_path, self.metadata, self.repo)
-
     def build(self):
         self.copy_module()
         self.write_metadata()
@@ -172,7 +159,33 @@ class WheelBuilder:
 
         self.wheel_zip.close()
 
-        self.post_build()
+def wheel_main(ini_path, upload=False, verify_metadata=False, repo='pypi'):
+    """Build a wheel in the dist/ directory, and optionally upload it.
+    """
+    dist_dir = ini_path.parent / 'dist'
+    try:
+        dist_dir.mkdir()
+    except FileExistsError:
+        pass
+
+    # We don't know the final filename until metadata is loaded, so write to
+    # a temporary_file, and rename it afterwards.
+    (fd, temp_path) = tempfile.mkstemp(suffix='.whl', dir=str(dist_dir))
+    with open(fd, 'w+b') as fp:
+        wb = WheelBuilder(ini_path, fp)
+        wb.build()
+
+    wheel_path = dist_dir / wb.wheel_filename
+    os.replace(temp_path, str(wheel_path))
+    log.info("Wheel built: %s", wheel_path)
+
+    if verify_metadata:
+        from .upload import verify
+        verify(wb.metadata, repo)
+
+    if upload:
+        from .upload import do_upload
+        do_upload(wheel_path, wb.metadata, repo)
 
 
 import stat, time


Test patch:
diff --git a/tests/test_upload.py b/tests/test_upload.py
index 2bbc365..de5bdb5 100644
--- a/tests/test_upload.py
+++ b/tests/test_upload.py
@@ -38,9 +38,8 @@ def test_verify():
 def test_upload():
     responses.add(responses.POST, upload.PYPI, status=200)
 
-    wb = wheel.WheelBuilder(samples_dir / 'module1-pkg.ini', upload='pypi')
     with patch('flit.upload.get_repository', return_value=repo_settings):
-        wb.build()
+        wheel.wheel_main(samples_dir / 'module1-pkg.ini', upload='pypi')
 
     assert len(responses.calls) == 1
 
@@ -54,9 +53,8 @@ def test_upload_registers():
         responses.add_callback(responses.POST, upload.PYPI,
                                callback=upload_callback)
 
-        wb = wheel.WheelBuilder(samples_dir / 'module1-pkg.ini', upload='pypi')
         with patch('flit.upload.get_repository', return_value=repo_settings):
-            wb.build()
+            wheel.wheel_main(samples_dir / 'module1-pkg.ini', upload='pypi')
 
     assert len(responses.calls) == 2
     assert register_mock.call_count == 1
diff --git a/tests/test_wheel.py b/tests/test_wheel.py
index 3a830a0..250c803 100644
--- a/tests/test_wheel.py
+++ b/tests/test_wheel.py
@@ -7,7 +7,7 @@ import zipfile
 import pytest
 from testpath import assert_isfile
 
-from flit.wheel import WheelBuilder, EntryPointsConflict
+from flit.wheel import wheel_main, WheelBuilder, EntryPointsConflict
 
 samples_dir = Path(__file__).parent / 'samples'
 
@@ -25,22 +25,22 @@ def unpack(path):
 
 def test_wheel_module():
     clear_samples_dist()
-    WheelBuilder(samples_dir / 'module1-pkg.ini').build()
+    wheel_main(samples_dir / 'module1-pkg.ini')
     assert_isfile(samples_dir / 'dist/module1-0.1-py2.py3-none-any.whl')
 
 def test_wheel_package():
     clear_samples_dist()
-    WheelBuilder(samples_dir / 'package1-pkg.ini').build()
+    wheel_main(samples_dir / 'package1-pkg.ini')
     assert_isfile(samples_dir / 'dist/package1-0.1-py2.py3-none-any.whl')
 
 def test_dist_name():
     clear_samples_dist()
-    WheelBuilder(samples_dir / 'altdistname.ini').build()
+    wheel_main(samples_dir / 'altdistname.ini')
     assert_isfile(samples_dir / 'dist/packagedist1-0.1-py2.py3-none-any.whl')
 
 def test_entry_points():
     clear_samples_dist()
-    WheelBuilder(samples_dir / 'entrypoints_valid.ini').build()
+    wheel_main(samples_dir / 'entrypoints_valid.ini')
     assert_isfile(samples_dir / 'dist/package1-0.1-py2.py3-none-any.whl')
     with unpack(samples_dir / 'dist/package1-0.1-py2.py3-none-any.whl') as td:
         entry_points = Path(td, 'package1-0.1.dist-info', 'entry_points.txt')
@@ -52,6 +52,16 @@ def test_entry_points():
 
 def test_entry_points_conflict():
     clear_samples_dist()
-    wb = WheelBuilder(samples_dir / 'entrypoints_conflict.ini')
     with pytest.raises(EntryPointsConflict):
-        wb.build()
+        wheel_main(samples_dir / 'entrypoints_conflict.ini')
+
+def test_wheel_builder():
+    # Slightly lower level interface
+    with tempfile.TemporaryDirectory() as td:
+        target = Path(td, 'sample.whl')
+        with target.open('wb') as f:
+            wb = WheelBuilder(samples_dir / 'package1-pkg.ini', f)
+            wb.build()
+
+        assert zipfile.is_zipfile(str(target))
+        assert wb.wheel_filename == 'package1-0.1-py2.py3-none-any.whl'