build_wheels.py

from __future__ import annotations

import argparse
import email.message
import fnmatch
import json
import os
import re
import shutil
import subprocess
import sys
import time
import tomllib
from functools import cache
from hashlib import sha256
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TypedDict
from zipfile import ZipFile

import pathspec
import urllib3
from dotenv import dotenv_values
from utils import iter_wheels


def unpack_wheel(wheel_file: Path, dest_dir: Path | None = None):
    """
    Unpack a wheel file into a directory using the supported CLI entrypoint:
      python -m wheel unpack <wheel> [-d <dest>]
    """
    cmd = [sys.executable, '-m', 'wheel', 'unpack']
    if dest_dir:
        cmd += ['-d', str(dest_dir)]
    cmd.append(str(wheel_file))
    subprocess.check_call(cmd)


def pack_wheel(wheel_dir: Path, dest_dir: Path | None = None):
    """
    Pack a wheel directory into a wheel using:
      python -m wheel pack <wheel-dir> [-d <dest>]
    """
    cmd = [sys.executable, '-m', 'wheel', 'pack']
    if dest_dir:
        cmd += ['-d', str(dest_dir)]
    cmd.append(str(wheel_dir))
    subprocess.check_call(cmd)

INDEX_BASE_URL = 'https://agent-int-packages.datadoghq.com'
CUSTOM_EXTERNAL_INDEX = f'{INDEX_BASE_URL}/external'
CUSTOM_BUILT_INDEX = f'{INDEX_BASE_URL}/built'
UNNORMALIZED_PROJECT_NAME_CHARS = re.compile(r'[-_.]+')


class WheelSizes(TypedDict):
    compressed: int
    uncompressed: int


class VersionedWheelSizes(WheelSizes):
    version: str


if sys.platform == 'win32':
    PY3_PATH = Path('C:\\py3\\Scripts\\python.exe')
    PY2_PATH = Path('C:\\py2\\Scripts\\python.exe')
    MOUNT_DIR = Path('C:\\mnt')
    ENV_FILE = Path('C:\\.env')

    def join_command_args(args: list[str]) -> str:
        return subprocess.list2cmdline(args)

    def path_to_uri(path: str | Path) -> str:
        return f'file:///{os.path.abspath(path).replace(" ", "%20").replace(os.sep, "/")}'

else:
    import shlex

    PY3_PATH = Path(os.environ.get('DD_PY3_BUILDENV_PATH', '/py3/bin/python'))
    PY2_PATH = Path(os.environ.get('DD_PY2_BUILDENV_PATH', '/py2/bin/python'))
    MOUNT_DIR = Path(os.environ.get('DD_MOUNT_DIR', '/home'))
    ENV_FILE = Path(os.environ.get('DD_ENV_FILE', '/.env'))

    def join_command_args(args: list[str]) -> str:
        return shlex.join(args)

    def path_to_uri(path: str | Path) -> str:
        return f'file://{os.path.abspath(path).replace(" ", "%20")}'


def abort(message, *, code=1):
    print(message, file=sys.stderr)
    sys.exit(code)


def check_process(*args, **kwargs) -> subprocess.CompletedProcess:
    print(f'Running: {args[0] if isinstance(args[0], str) else join_command_args(args[0])}', file=sys.stderr)
    process = subprocess.run(*args, **kwargs)
    if process.returncode:
        sys.exit(process.returncode)

    return process


def extract_metadata(wheel: Path) -> email.message.Message:
    with ZipFile(str(wheel)) as zip_archive:
        for path in zip_archive.namelist():
            root = path.split('/', 1)[0]
            if root.endswith('.dist-info'):
                dist_info_dir = root
                break
        else:
            message = f'Could not find the `.dist-info` directory in wheel: {wheel.name}'
            raise RuntimeError(message)

        try:
            with zip_archive.open(f'{dist_info_dir}/METADATA') as zip_file:
                metadata_file_contents = zip_file.read().decode('utf-8')
        except KeyError:
            message = f'Could not find a `METADATA` file in the `{dist_info_dir}` directory'
            raise RuntimeError(message) from None

    return email.message_from_string(metadata_file_contents)


def normalize_project_name(name: str) -> str:
    # https://peps.python.org/pep-0503/#normalized-names
    return UNNORMALIZED_PROJECT_NAME_CHARS.sub('-', name).lower()


@cache
def get_wheel_hashes(project) -> dict[str, str]:
    retry_wait = 2
    while True:
        try:
            response = urllib3.request(
                'GET',
                f'https://pypi.org/simple/{project}',
                headers={"Accept": "application/vnd.pypi.simple.v1+json"},
            )
        except urllib3.exceptions.HTTPError as e:
            err_msg = f'Failed to fetch hashes for `{project}`: {e}'
        else:
            if response.status == 200:
                break

            err_msg = f'Failed to fetch hashes for `{project}`, status code: {response.status}'

        print(err_msg)
        print(f'Retrying in {retry_wait} seconds')
        time.sleep(retry_wait)
        retry_wait *= 2
        continue

    data = response.json()
    return {
        file['filename']: file['hashes']['sha256']
        for file in data['files']
        if file['filename'].endswith('.whl') and 'sha256' in file['hashes']
    }


def wheel_was_built(wheel: Path) -> bool:
    project_metadata = extract_metadata(wheel)
    project_name = normalize_project_name(project_metadata['Name'])
    wheel_hashes = get_wheel_hashes(project_name)
    if wheel.name not in wheel_hashes:
        return True

    file_hash = sha256(wheel.read_bytes()).hexdigest()
    return file_hash != wheel_hashes[wheel.name]


def _remove_test_files_from_dir(unpacked_dir: Path) -> bool:
    """Remove excluded test files and directories from an unpacked wheel directory."""
    removed_any = False
    for root, dirs, files in unpacked_dir.walk(top_down=False):
        for d in list(dirs):
            full_dir = root / d
            rel = full_dir.relative_to(unpacked_dir).as_posix()
            if is_excluded_from_wheel(rel):
                shutil.rmtree(full_dir)
                dirs.remove(d)
                removed_any = True
        for f in files:
            rel = (root / f).relative_to(unpacked_dir).as_posix()
            if is_excluded_from_wheel(rel):
                (root / f).unlink()
                removed_any = True
    return removed_any


@cache
def _load_excluded_spec() -> pathspec.PathSpec:
    """
    Load excluded paths from files_to_remove.toml and compile them
    with .gitignore-style semantics.
    """
    config_path = Path(__file__).parent / "files_to_remove.toml"
    with open(config_path, "rb") as f:
        config = tomllib.load(f)

    patterns = config.get("excluded_paths", [])
    return pathspec.PathSpec.from_lines("gitignore", patterns)


def is_excluded_from_wheel(path: str | Path) -> bool:
    """
    Return True if `path` (file or directory) should be excluded per files_to_remove.toml.
    Matches:
      - type annotation files: **/*.pyi, **/py.typed
      - test directories listed with a trailing '/'
    """
    spec = _load_excluded_spec()
    rel = Path(path).as_posix()

    if spec.match_file(rel) or spec.match_file(rel + "/"):
        return True

    return False


@cache
def _load_line_removal_rules() -> list[dict]:
    '''
    Load the line removal rules from the toml file and compile the regex patterns.
    '''

    config_p = Path(__file__).parent / "lines_to_remove.toml"
    with open(config_p, "rb") as f:
        config = tomllib.load(f)

    rules = config.get("rules", [])
    for rule in rules:
        rule['_compiled_patterns'] = [re.compile(p) for p in rule.get('line_patterns', [])]
    return rules


def _get_matching_rules(wheel_name: str) -> list[dict]:
    '''
    Match the rules in the toml to the wheel name.
    '''
    rules = _load_line_removal_rules()
    matching = []
    for rule in rules:
        pattern = rule.get('package_pattern', '*')
        if fnmatch.fnmatch(wheel_name, pattern):
            matching.append(rule)
    return matching


def _strip_lines_from_dir(unpacked_dir: Path, rules: list[dict]) -> bool:
    """Strip matching lines and their indented blocks from files."""
    modified = False
    for rule in rules:
        file_pattern = rule.get('file_pattern', '*.py')
        compiled_patterns = rule['_compiled_patterns']

        for py_file in unpacked_dir.rglob(file_pattern):
            if not py_file.is_file():
                continue

            try:
                original_content = py_file.read_text(encoding='utf-8')
            except UnicodeDecodeError:
                continue

            lines = original_content.splitlines(keepends=True)
            kept_lines: list[str] = []
            skip_indent: int | None = None

            for line in lines:
                stripped = line.rstrip('\n\r')

                if skip_indent is not None:
                    expanded = stripped.expandtabs()
                    line_indent = len(expanded) - len(expanded.lstrip())
                    if stripped == '' or line_indent > skip_indent:
                        continue
                    skip_indent = None

                if any(pattern.search(line) for pattern in compiled_patterns):
                    expanded = stripped.expandtabs()
                    skip_indent = len(expanded) - len(expanded.lstrip())
                    continue

                kept_lines.append(line)

            if len(kept_lines) != len(lines):
                py_file.write_text(''.join(kept_lines), encoding='utf-8')
                modified = True
    return modified


def clean_wheel(wheel_path: Path) -> bool:
    """
    Unpack the wheel at most once, remove excluded test files and strip
    matching lines, then repack only if changes were made.
    """
    with ZipFile(wheel_path, 'r') as zf:
        has_excluded_files = any(is_excluded_from_wheel(name) for name in zf.namelist())

    line_removal_rules = _get_matching_rules(wheel_path.name)

    if not has_excluded_files and not line_removal_rules:
        return False

    with TemporaryDirectory() as td:
        td_path = Path(td)
        unpack_wheel(wheel_path, dest_dir=td_path)
        unpacked_dir = next(td_path.iterdir())

        modified = False
        if has_excluded_files:
            if _remove_test_files_from_dir(unpacked_dir):
                print(f'Tests removed from {wheel_path.name}')
                modified = True
        if line_removal_rules:
            if _strip_lines_from_dir(unpacked_dir, line_removal_rules):
                print(f'Stripped lines from {wheel_path.name}')
                modified = True

        if not modified:
            return False

        dest_dir = wheel_path.parent
        before = {p.resolve() for p in dest_dir.glob("*.whl")}
        pack_wheel(unpacked_dir, dest_dir=dest_dir)
        # Repacking may restore the wheel's original (non-platform-specific) name.
        # Move the repacked wheel back to the platform-specific path.
        after = {p.resolve() for p in dest_dir.glob("*.whl")}
        new_files = sorted(after - before, key=lambda p: p.stat().st_mtime, reverse=True)

        if new_files:
            shutil.move(str(new_files[0]), str(wheel_path))

    return True

def add_dependency(dependencies: dict[str, str], sizes: dict[str, VersionedWheelSizes], wheel: Path) -> None:
    project_metadata = extract_metadata(wheel)
    project_name = normalize_project_name(project_metadata['Name'])
    project_version = project_metadata['Version']
    dependencies[project_name] = project_version
    sizes[project_name] = {'version': project_version, **calculate_wheel_sizes(wheel)}


def calculate_wheel_sizes(wheel_path: Path) -> WheelSizes:
    compressed_size = wheel_path.stat(follow_symlinks=True).st_size
    with ZipFile(wheel_path) as zf:
        uncompressed_size = sum(zinfo.file_size for zinfo in zf.infolist())
    return {'compressed': compressed_size, 'uncompressed': uncompressed_size}


def main():
    parser = argparse.ArgumentParser(prog='wheel-builder', allow_abbrev=False)
    parser.add_argument('--python', required=True)
    parser.add_argument('--use-built-index', action='store_true', default=False)
    args = parser.parse_args()

    python_version = args.python
    if python_version == '3':
        python_path = PY3_PATH
    elif python_version == '2':
        python_path = PY2_PATH
    else:
        abort(f'Invalid python version: {python_version}')

    wheels_dir = MOUNT_DIR / 'wheels'
    built_wheels_dir = wheels_dir / 'built'
    external_wheels_dir = wheels_dir / 'external'

    # Install build dependencies
    check_process([str(python_path), '-m', 'pip', 'install', '-r', str(MOUNT_DIR / 'build_dependencies.txt')])

    with TemporaryDirectory() as d:
        staged_wheel_dir = Path(d).resolve()
        staged_built_wheels_dir = staged_wheel_dir / 'built'
        staged_external_wheels_dir = staged_wheel_dir / 'external'

        # Create the directories
        staged_built_wheels_dir.mkdir(parents=True, exist_ok=True)
        staged_external_wheels_dir.mkdir(parents=True, exist_ok=True)

        env_vars = dict(os.environ)
        env_vars['PATH'] = f'{python_path.parent}{os.pathsep}{env_vars["PATH"]}'
        env_vars['PIP_WHEEL_DIR'] = str(staged_wheel_dir)
        env_vars['DD_BUILD_PYTHON_VERSION'] = python_version
        env_vars['DD_MOUNT_DIR'] = str(MOUNT_DIR)
        env_vars['DD_ENV_FILE'] = str(ENV_FILE)

        # Off is on, see: https://github.com/pypa/pip/issues/5735
        # We don't want pip build isolation because we manage have one env for everything.
        # We manage our build dependencies in .builders/deps/build_dependencies.txt.
        env_vars['PIP_NO_BUILD_ISOLATION'] = '0'

        # Spaces are used to separate multiple values which means paths themselves cannot contain spaces, see:
        # https://github.com/pypa/pip/issues/10114#issuecomment-1880125475
        env_vars['PIP_FIND_LINKS'] = path_to_uri(staged_wheel_dir)

        # Perform builder-specific logic if required
        if build_command := os.environ.get('DD_BUILD_COMMAND'):
            check_process(build_command, env=env_vars, shell=True)

        # Load environment variables
        if ENV_FILE.is_file():
            for key, value in dotenv_values(str(ENV_FILE)).items():
                if value is None:
                    env_vars.pop(key, None)
                else:
                    env_vars[key] = value

        if constraints_file := env_vars.get('PIP_CONSTRAINT'):
            env_vars['PIP_CONSTRAINT'] = path_to_uri(constraints_file)

        # Fetch or build wheels
        command_args = [
            str(python_path),
            '-m',
            'pip',
            'wheel',
            '-r',
            str(MOUNT_DIR / 'requirements.in'),
            '--wheel-dir',
            str(staged_wheel_dir),
            '--extra-index-url',
            CUSTOM_EXTERNAL_INDEX,
        ]

        check_process(command_args, env=env_vars)

        # Classify wheels
        for wheel in iter_wheels(staged_wheel_dir):
            if wheel_was_built(wheel):
                shutil.move(wheel, staged_built_wheels_dir)
            else:
                shutil.move(wheel, staged_external_wheels_dir)

        # Repair wheels
        check_process(
            [
                sys.executable,
                '-u',
                str(MOUNT_DIR / 'scripts' / 'repair_wheels.py'),
                '--source-built-dir',
                str(staged_built_wheels_dir),
                '--source-external-dir',
                str(staged_external_wheels_dir),
                '--built-dir',
                str(built_wheels_dir),
                '--external-dir',
                str(external_wheels_dir),
            ]
        )

    dependencies: dict[str, str] = {}
    sizes: dict[str, VersionedWheelSizes] = {}

    # Handle wheels already in the built directory
    for wheel in iter_wheels(built_wheels_dir):
        clean_wheel(wheel)
        add_dependency(dependencies, sizes, wheel)

    # Handle wheels currently in the external directory and move them to the built directory if they were modified
    for wheel in iter_wheels(external_wheels_dir):
        was_modified = clean_wheel(wheel)
        if was_modified:
            new_path = built_wheels_dir / wheel.name
            wheel.rename(new_path)
            wheel = new_path
            print(f'Moved {wheel.name} to built directory')

        add_dependency(dependencies, sizes, wheel)

    output_path = MOUNT_DIR / 'sizes.json'
    with output_path.open('w', encoding='utf-8') as fp:
        json.dump(sizes, fp, indent=2, sort_keys=True)

    final_requirements = MOUNT_DIR / 'frozen.txt'
    with final_requirements.open('w', encoding='utf-8') as f:
        for project_name, project_version in sorted(dependencies.items()):
            f.write(f'{project_name}=={project_version}\n')


if __name__ == '__main__':
    main()