infra/tools/dockerbuild/patch_wheel.py

# Copyright 2026 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Patches a wheel file to update its version.

This is used to fix PEP 440 versioning issues for Artifact Registry
without modifying the original wheel used for CIPD.
"""

import base64
import hashlib
import os
import re
import tempfile
import zipfile

from typing import Optional


def create_dummy_wheel(wheel_path: str, escaped_name: str,
                       version: str, requirements: Optional[list[str]] = None) -> None:
  """Creates a dummy metapackage wheel containing only METADATA and WHEEL files.

  Args:
    wheel_path: Path where the new dummy wheel should be written.
    escaped_name: The escaped package name (e.g. pyobjc).
    version: The version string.
    requirements: Optional list of package requirement strings.
  """
  dist_info_dir = f"{escaped_name}-{version}.dist-info"
  metadata_content = ("Metadata-Version: 2.1\n"
                      f"Name: {escaped_name}\n"
                      f"Version: {version}\n"
                      "Summary: Dummy metapackage wheel\n")
  if requirements:
    for req in requirements:
      metadata_content += f"Requires-Dist: {req}\n"
  wheel_content = ("Wheel-Version: 1.0\n"
                   "Root-Is-Purelib: true\n"
                   "Tag: py2-none-any\n"
                   "Tag: py3-none-any\n")

  with zipfile.ZipFile(wheel_path, 'w', zipfile.ZIP_DEFLATED) as z:
    z.writestr(f"{dist_info_dir}/METADATA", metadata_content)
    z.writestr(f"{dist_info_dir}/WHEEL", wheel_content)
    z.writestr(f"{dist_info_dir}/RECORD", "")


def patch_wheel(wheel_path: str, new_version: str, escaped_name: str) -> None:
  """Patches a wheel file with a new version.

  Args:
    wheel_path: Path to the original wheel file.
    new_version: The new version string to apply.
    escaped_name: The escaped package name (e.g. aioquic, NOT aioquic-py3).
  """
  with tempfile.TemporaryDirectory() as tdir:
    # Extract wheel
    with zipfile.ZipFile(wheel_path, 'r') as z:
      z.extractall(tdir)

    # Find .dist-info
    dist_info_dirs = [
        name for name in os.listdir(tdir) if name.endswith('.dist-info')
    ]

    if not dist_info_dirs:
      raise ValueError(f"No .dist-info directory found in {wheel_path}")

    # We need to find the one matching our escaped_name by checking internal METADATA.
    dist_info_dir = None
    expected_name = escaped_name.replace('-', '_').lower()

    for d in dist_info_dirs:
      candidate_dir = os.path.join(tdir, d)
      metadata_path = os.path.join(candidate_dir, 'METADATA')
      if not os.path.exists(metadata_path):
        continue

      with open(metadata_path, 'r', encoding='utf-8') as f:
        content = f.read()

      name_match = re.search(r'^Name:\s*(.*)$', content, re.MULTILINE)
      if name_match:
        actual_name = name_match.group(1).strip().replace('-', '_').lower()
        if actual_name == expected_name:
          dist_info_dir = candidate_dir
          metadata_content = content
          break

    if not dist_info_dir:
      raise ValueError(
          f"Could not find a .dist-info directory for {escaped_name} in {wheel_path}. "
          f"Found candidates: {dist_info_dirs}")

    # Read METADATA (already read in the loop)
    metadata_path = os.path.join(dist_info_dir, 'METADATA')

    # Update Version
    # We use a strict regex to avoid matching Version if it appears in description.
    new_metadata_content, count = re.subn(
        r'^(Version:\s*).*$',
        rf'\g<1>{new_version}',
        metadata_content,
        count=1,
        flags=re.MULTILINE)

    if count == 0:
      raise ValueError(f"Could not find Version: field in {metadata_path}")

    # If nothing changed, we are done (idempotency)
    if new_metadata_content == metadata_content:
      return

    with open(metadata_path, 'w', encoding='utf-8') as f:
      f.write(new_metadata_content)

    # Update RECORD
    record_path = os.path.join(dist_info_dir, 'RECORD')
    if not os.path.exists(record_path):
      raise ValueError(f"No RECORD file found in {dist_info_dir}")

    with open(record_path, 'r', encoding='utf-8') as f:
      record_lines = f.readlines()

    new_record_lines = []

    # We need to calculate the hash of the new METADATA
    hasher = hashlib.sha256()
    hasher.update(new_metadata_content.encode('utf-8'))
    metadata_hash = base64.urlsafe_b64encode(
        hasher.digest()).decode('utf-8').rstrip('=')
    metadata_size = len(new_metadata_content.encode('utf-8'))

    dist_info_name = os.path.basename(dist_info_dir)
    new_dist_info_name = f"{escaped_name}-{new_version}.dist-info"

    for line in record_lines:
      parts = line.strip().split(',')
      if len(parts) < 2:
        new_record_lines.append(line)
        continue

      file_path = parts[0]

      # Update paths if we are renaming dist-info
      if file_path.startswith(dist_info_name + '/'):
        file_path = file_path.replace(dist_info_name, new_dist_info_name, 1)
        parts[0] = file_path

      # Update METADATA hash
      if file_path == f"{new_dist_info_name}/METADATA":
        parts[1] = f"sha256={metadata_hash}"
        parts[2] = str(metadata_size)

      # RECORD itself has empty hash and size
      if file_path == f"{new_dist_info_name}/RECORD":
        parts[1] = ""
        parts[2] = ""

      new_record_lines.append(','.join(parts) + '\n')

    with open(record_path, 'w', encoding='utf-8') as f:
      f.writelines(new_record_lines)

    # Rename .dist-info
    new_dist_info_dir = os.path.join(tdir, new_dist_info_name)
    if dist_info_dir != new_dist_info_dir:
      os.rename(dist_info_dir, new_dist_info_dir)

    # Zip back
    temp_output_path = wheel_path + '.tmp'
    with zipfile.ZipFile(temp_output_path, 'w', zipfile.ZIP_DEFLATED) as z:
      for root, _, files in os.walk(tdir):
        for file in files:
          file_path = os.path.join(root, file)
          arcname = os.path.relpath(file_path, tdir)
          z.write(file_path, arcname)

    os.replace(temp_output_path, wheel_path)