pack.py 4.2 KB

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