1from __future__ import print_function 2 3import os.path 4import re 5import sys 6 7from wheel.cli import WheelError 8from wheel.wheelfile import WheelFile 9 10DIST_INFO_RE = re.compile(r"^(?P<namever>(?P<name>.+?)-(?P<ver>\d.*?))\.dist-info$") 11BUILD_NUM_RE = re.compile(br'Build: (\d\w*)$') 12 13 14def pack(directory, dest_dir, build_number): 15 """Repack a previously unpacked wheel directory into a new wheel file. 16 17 The .dist-info/WHEEL file must contain one or more tags so that the target 18 wheel file name can be determined. 19 20 :param directory: The unpacked wheel directory 21 :param dest_dir: Destination directory (defaults to the current directory) 22 """ 23 # Find the .dist-info directory 24 dist_info_dirs = [fn for fn in os.listdir(directory) 25 if os.path.isdir(os.path.join(directory, fn)) and DIST_INFO_RE.match(fn)] 26 if len(dist_info_dirs) > 1: 27 raise WheelError('Multiple .dist-info directories found in {}'.format(directory)) 28 elif not dist_info_dirs: 29 raise WheelError('No .dist-info directories found in {}'.format(directory)) 30 31 # Determine the target wheel filename 32 dist_info_dir = dist_info_dirs[0] 33 name_version = DIST_INFO_RE.match(dist_info_dir).group('namever') 34 35 # Read the tags and the existing build number from .dist-info/WHEEL 36 existing_build_number = None 37 wheel_file_path = os.path.join(directory, dist_info_dir, 'WHEEL') 38 with open(wheel_file_path) as f: 39 tags = [] 40 for line in f: 41 if line.startswith('Tag: '): 42 tags.append(line.split(' ')[1].rstrip()) 43 elif line.startswith('Build: '): 44 existing_build_number = line.split(' ')[1].rstrip() 45 46 if not tags: 47 raise WheelError('No tags present in {}/WHEEL; cannot determine target wheel filename' 48 .format(dist_info_dir)) 49 50 # Set the wheel file name and add/replace/remove the Build tag in .dist-info/WHEEL 51 build_number = build_number if build_number is not None else existing_build_number 52 if build_number is not None: 53 if build_number: 54 name_version += '-' + build_number 55 56 if build_number != existing_build_number: 57 replacement = ('Build: %s\r\n' % build_number).encode('ascii') if build_number else b'' 58 with open(wheel_file_path, 'rb+') as f: 59 wheel_file_content = f.read() 60 if not BUILD_NUM_RE.subn(replacement, wheel_file_content)[1]: 61 wheel_file_content += replacement 62 63 f.truncate() 64 f.write(wheel_file_content) 65 66 # Reassemble the tags for the wheel file 67 impls = sorted({tag.split('-')[0] for tag in tags}) 68 abivers = sorted({tag.split('-')[1] for tag in tags}) 69 platforms = sorted({tag.split('-')[2] for tag in tags}) 70 tagline = '-'.join(['.'.join(impls), '.'.join(abivers), '.'.join(platforms)]) 71 72 # Repack the wheel 73 wheel_path = os.path.join(dest_dir, '{}-{}.whl'.format(name_version, tagline)) 74 with WheelFile(wheel_path, 'w') as wf: 75 print("Repacking wheel as {}...".format(wheel_path), end='') 76 sys.stdout.flush() 77 wf.write_files(directory) 78 79 print('OK') 80