pack.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. from __future__ import annotations
  2. import email.policy
  3. import os.path
  4. import re
  5. from email.generator import BytesGenerator
  6. from email.parser import BytesParser
  7. from ..wheelfile import WheelError, WheelFile
  8. DIST_INFO_RE = re.compile(r"^(?P<namever>(?P<name>.+?)-(?P<ver>\d.*?))\.dist-info$")
  9. def pack(directory: str, dest_dir: str, build_number: str | None) -> None:
  10. """Repack a previously unpacked wheel directory into a new wheel file.
  11. The .dist-info/WHEEL file must contain one or more tags so that the target
  12. wheel file name can be determined.
  13. :param directory: The unpacked wheel directory
  14. :param dest_dir: Destination directory (defaults to the current directory)
  15. """
  16. # Find the .dist-info directory
  17. dist_info_dirs = [
  18. fn
  19. for fn in os.listdir(directory)
  20. if os.path.isdir(os.path.join(directory, fn)) and DIST_INFO_RE.match(fn)
  21. ]
  22. if len(dist_info_dirs) > 1:
  23. raise WheelError(f"Multiple .dist-info directories found in {directory}")
  24. elif not dist_info_dirs:
  25. raise WheelError(f"No .dist-info directories found in {directory}")
  26. # Determine the target wheel filename
  27. dist_info_dir = dist_info_dirs[0]
  28. name_version = DIST_INFO_RE.match(dist_info_dir).group("namever")
  29. # Read the tags and the existing build number from .dist-info/WHEEL
  30. wheel_file_path = os.path.join(directory, dist_info_dir, "WHEEL")
  31. with open(wheel_file_path, "rb") as f:
  32. info = BytesParser(policy=email.policy.compat32).parse(f)
  33. tags: list[str] = info.get_all("Tag", [])
  34. existing_build_number = info.get("Build")
  35. if not tags:
  36. raise WheelError(
  37. f"No tags present in {dist_info_dir}/WHEEL; cannot determine target "
  38. f"wheel filename"
  39. )
  40. # Set the wheel file name and add/replace/remove the Build tag in .dist-info/WHEEL
  41. build_number = build_number if build_number is not None else existing_build_number
  42. if build_number is not None:
  43. del info["Build"]
  44. if build_number:
  45. info["Build"] = build_number
  46. name_version += "-" + build_number
  47. if build_number != existing_build_number:
  48. with open(wheel_file_path, "wb") as f:
  49. BytesGenerator(f, maxheaderlen=0).flatten(info)
  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 compute_tagline(tags: list[str]) -> str:
  59. """Compute a tagline from a list of tags.
  60. :param tags: A list of tags
  61. :return: A tagline
  62. """
  63. impls = sorted({tag.split("-")[0] for tag in tags})
  64. abivers = sorted({tag.split("-")[1] for tag in tags})
  65. platforms = sorted({tag.split("-")[2] for tag in tags})
  66. return "-".join([".".join(impls), ".".join(abivers), ".".join(platforms)])