Patch to remove tests from wheels (#20653)

* wip

* remove changes

* remove tests

* remove funcions

* patch

* remove imports for patch

* patch

* print in patch

* fix typo

* repack wheels

* fix

* add wheel library

* comment patch

* after repair

* patch

* uncomment patch call

* skip unchanged

* skip unchanged

* skip unchanged

* debug

* debug

* debug

* after repair

* after repair

* fix

* classify wheels

* fix

* classify wheels

* remove import

* patch

* patch fix

* fix

* remove quotes

* replace quote

* fix quotes in patch

* remove patch

* remove patch

* simplify utils and .toml

* toml

* toml

* toml

* typo

* change toml format to gitignore patterns

* publish wheels for testing

* test

* rename wheels

* uncomment

* remove built extra index

* remove built flag

* upload wheels

* Uncomment condition to publish wheels

* Fix typo

* Fix conflicts

* Add formatting

* Remove unnecessary continues

Author: lucia-sb

.builders/images/runner_dependencies.txt

@@ -3,3 +3,5 @@ urllib3==2.2.0
 auditwheel==6.0.0; sys_platform == 'linux'
 delvewheel==1.5.2; sys_platform == 'win32'
 delocate==0.13.0; sys_platform == 'darwin'
+wheel==0.45.1
+pathspec==0.12.1

.builders/scripts/build_wheels.py

@@ -1,26 +1,40 @@
 from __future__ import annotations
 
 import argparse
+import email
 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 extract_metadata, normalize_project_name
+from utils import iter_wheels
+from wheel.cli.pack import pack
+from wheel.cli.unpack import unpack
 
 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
 
+
 if sys.platform == 'win32':
     PY3_PATH = Path('C:\\py3\\Scripts\\python.exe')
     PY2_PATH = Path('C:\\py2\\Scripts\\python.exe')
@@ -62,6 +76,161 @@ def check_process(*args, **kwargs) -> subprocess.CompletedProcess:
     return process
 
 
+def extract_metadata(wheel: Path) -> email.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(wheel_path: Path) -> bool:
+    '''
+    Unpack the wheel, remove excluded test files, then repack it to rebuild RECORD correctly.
+    '''
+    # First, check whether the wheel contains any files that should be excluded. If not, leave it untouched.
+    with ZipFile(wheel_path, 'r') as zf:
+        excluded_members = [name for name in zf.namelist() if is_excluded_from_wheel(name)]
+
+    if not excluded_members:
+        # Nothing to strip, so skip rewriting the wheel
+        return False
+    with TemporaryDirectory() as td:
+        td_path = Path(td)
+
+        # Unpack the wheel into temp dir
+        unpack(wheel_path, dest=td_path)
+        unpacked_dir = next(td_path.iterdir())
+        # Remove excluded files/folders
+        for root, dirs, files in os.walk(td, topdown=False):
+            for d in list(dirs):
+                full_dir = Path(root) / d
+                rel = full_dir.relative_to(unpacked_dir).as_posix()
+                if is_excluded_from_wheel(rel):
+                    shutil.rmtree(full_dir)
+                    dirs.remove(d)
+            for f in files:
+                rel = Path(root).joinpath(f).relative_to(unpacked_dir).as_posix()
+                if is_excluded_from_wheel(rel):
+                    os.remove(Path(root) / f)
+
+        print(f'Tests removed from {wheel_path.name}')
+
+        dest_dir = wheel_path.parent
+        before = {p.resolve() for p in dest_dir.glob("*.whl")}
+        # Repack to same directory, regenerating RECORD
+        pack(unpacked_dir, dest_dir=dest_dir, build_number=None)
+
+        # The wheel might not be platform-specific, so repacking restores its original name.
+        # We need to move the repacked wheel to wheel_path, which was changed to be platform-specific.
+        after = {p.resolve() for p in wheel_path.parent.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
+
+
+@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
+
+
+def add_dependency(dependencies: dict[str, str], sizes: dict[str, WheelSizes], 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:
@@ -92,6 +261,13 @@ def main():
 
     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)
@@ -131,25 +307,29 @@ def main():
             str(MOUNT_DIR / 'requirements.in'),
             '--wheel-dir',
             str(staged_wheel_dir),
-            # Temporarily removing extra index urls. See below.
-            # '--extra-index-url', CUSTOM_EXTERNAL_INDEX,
+            '--extra-index-url',
+            CUSTOM_EXTERNAL_INDEX,
         ]
-        # Temporarily disable extra index urls. There are broken wheels in the gcloud bucket
-        # while working on removing tests from them. Adding extra indices causes undefined behavior
-        # and can pull a broken image, preventing the building from running.
-        # if args.use_built_index:
-        #     command_args.extend(['--extra-index-url', CUSTOM_BUILT_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-dir',
-                str(staged_wheel_dir),
+                '--source-built-dir',
+                str(staged_built_wheels_dir),
+                '--source-external-dir',
+                str(staged_external_wheels_dir),
                 '--built-dir',
                 str(built_wheels_dir),
                 '--external-dir',
@@ -160,15 +340,22 @@ def main():
     dependencies: dict[str, tuple[str, str]] = {}
     sizes: dict[str, WheelSizes] = {}
 
-    for wheel_dir in wheels_dir.iterdir():
-        for wheel in wheel_dir.iterdir():
-            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)}
+    # 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 = remove_test_files(wheel)
+        if was_modified:
+            # A modified wheel is no longer external → move it to built directory
+            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)
+
+    # Handle wheels already in the built directory
+    for wheel in iter_wheels(built_wheels_dir):
+        remove_test_files(wheel)
+        add_dependency(dependencies, sizes, wheel)
 
     output_path = MOUNT_DIR / 'sizes.json'
     with output_path.open('w', encoding='utf-8') as fp:

.builders/scripts/files_to_remove.toml

@@ -0,0 +1,45 @@
+excluded_paths = [
+  # --- Type annotation ---
+  "krb5/**/*.pyi",
+  "krb5/**/py.typed",
+
+  "Cryptodome/**/*.pyi",
+  "Cryptodome/**/py.typed",
+
+  "ddtrace/**/*.pyi",
+  "ddtrace/**/py.typed",
+
+  "pyVmomi/**/*.pyi",
+  "pyVmomi/**/py.typed",
+
+  "gssapi/**/*.pyi",
+  "gssapi/**/py.typed",
+
+  # --- Tests ---
+ 
+  "idlelib/idle_test/",
+  "bs4/tests/",
+  "Cryptodome/SelfTest/",
+  "gssapi/tests/",
+  "keystoneauth1/tests/",
+  "lazy_loader/tests/",
+  "openstack/tests/",
+  "os_service_types/tests/",
+  "pbr/tests/",
+  "pkg_resources/tests/",
+  "pip/_vendor/colorama/tests/",
+  "psutil/tests/",
+  "requests_unixsocket/tests/",
+  "securesystemslib/_vendor/ed25519/test_data/",
+  "setuptools/_distutils/compilers/C/tests/",
+  "setuptools/_vendor/packaging/tests/",
+  "setuptools/_distutils/tests/",
+  "setuptools/tests/",
+  "simplejson/tests/",
+  "stevedore/tests/",
+  "supervisor/tests/",
+  "/test/",
+  "vertica_python/tests/",
+  "websocket/tests/",
+  "win32com/test/",
+]