Implement variant-specific patches

Fromager now supports variant-specific patches in additional to
versioned and unversioned patches. Variant-specific patches are in
subfolders, e.g. `atches/test_pkg-1.0.2/cpu/005-cpuver.patch`.

Patch files are sorted alphanumerically by their *base* name, not their
full path. This allows users to sort versioned/unversioned and
plain/variant patches as they like.

Internal changes: `sources.patch_source()` now uses
`PackageBuildInfo` to get patches. `patches_for_requirement()` has been
removed.

Fixes: python-wheel-build#486
Signed-off-by: Christian Heimes <[email protected]>

docs/customization.md

@@ -170,34 +170,46 @@ any further dependencies or building the project. The default directory is
 `overrides/patches`.
 
 Patch files should be placed in a subdirectory matching the source directory
-name and use the suffix `.patch`. The filenames are sorted lexicographically, so
-any text between the prefix and suffix can be used to ensure the patches are
-applied in a specific order.
+name and use the suffix `.patch`. The filenames are sorted lexicographically
+by their base name, so any text between the prefix and extension can be used
+to ensure the patches are applied in a specific order. Patch can be
+version-specific or versionless, apply to all variants, or apply to a single
+variant.
 
 Patches are applied by running `patch -p1 filename` while inside the root of the
 source tree.
 
 ```console
-$ ls -1 overrides/patches/*
-clarifai-10.2.1/fix-sdist.patch
-flash_attn-2.5.7/pyproject-toml.patch
-jupyterlab_pygments-0.3.0/pyproject-remove-jupyterlab.patch
-ninja-1.11.1.1/wrap-system-ninja.patch
-pytorch-v2.2.1/001-remove-cmake-build-requirement.patch
-pytorch-v2.2.1/002-dist-info-no-run-build-deps.patch
-pytorch-v2.2.1/003-fbgemm-no-maybe-uninitialized.patch
-pytorch-v2.2.1/004-fix-release-version.patch
-pytorch-v2.2.2/001-remove-cmake-build-requirement.patch
-pytorch-v2.2.2/002-dist-info-no-run-build-deps.patch
-pytorch-v2.2.2/003-fbgemm-no-maybe-uninitialized.patch
-pytorch-v2.2.2/004-fix-release-version.patch
-xformers-0.0.26.post1/pyproject.toml.patch
+$ tree overrides/patches/
+overrides/patches/
+├── test_pkg
+│   ├── 010-unversioned.patch
+│   └── cpu
+│       └── 004-cpu.patch
+└── test_pkg-1.0.2
+    ├── 001-somepatch.patch
+    ├── 002-otherpatch.patch
+    ├── cpu
+    │   └── 005-cpuver.patch
+    └── rocm
+        └── 005-rocmver.patch
 ```
 
-Note: A legacy patch organization with the patches in the patches directory, not
-in subdirectories, with the filenames prefixed with the source directory name is
-also supported. The newer format, using subdirectories, is preferred because it
-avoids name collisions between variant source trees.
+For package `test-pkg`, version `1.0.2`, and variant `cpu`, Fromager would apply:
+
+1. `001-somepatch.patch`
+2. `002-otherpatch.patch`
+3. `004-cpu.patch`
+4. `005-cpuver.patch`
+5. `010-unversioned.patch`
+
+For version `1.0.3` and variant `rocm`, Fromager would only apply
+`010-unversioned.patch`.
+
+```{versionchanged} 0.33.0
+
+Added support for variant-specific patches.
+```
 
 ## `project_override` section
 

src/fromager/overrides.py

@@ -1,13 +1,11 @@
 import inspect
-import itertools
 import logging
 import pathlib
 import typing
 from importlib import metadata
 
 from packaging.requirements import Requirement
 from packaging.utils import canonicalize_name
-from packaging.version import Version
 from stevedore import extension
 
 # An interface for reretrieving per-package information which influences
@@ -94,29 +92,6 @@ def log_overrides() -> None:
         )
 
 
-def patches_for_requirement(
-    patches_dir: pathlib.Path,
-    req: Requirement,
-    version: Version,
-) -> typing.Iterable[pathlib.Path]:
-    """Iterator producing patches to apply to the source for a given version of a requirement.
-
-    Yields pathlib.Path() references to patches in the order they should be
-    applied, which is controlled through lexical sorting of the filenames.
-
-    """
-    override_name = pkgname_to_override_module(req.name)
-    unversioned_patch_dir = patches_dir / override_name
-    versioned_patch_dir = patches_dir / f"{override_name}-{version}"
-    return itertools.chain(
-        # Apply all of the unversioned patches first, in order based on
-        # filename.
-        sorted(unversioned_patch_dir.glob("*.patch")),
-        # Then apply any for this specific version, in order based on filename.
-        sorted(versioned_patch_dir.glob("*.patch")),
-    )
-
-
 def get_versioned_patch_directories(
     patches_dir: pathlib.Path,
     req: Requirement,

src/fromager/packagesettings.py

@@ -88,7 +88,7 @@ def _validate_envkey(v: typing.Any) -> str:
 ]
 
 # patch mapping
-PatchMap = dict[Version, typing.Iterable[pathlib.Path]]
+PatchMap = dict[Version | None, list[pathlib.Path]]
 
 # URL or filename with templating
 Template = typing.NewType("Template", str)
@@ -493,18 +493,38 @@ def plugin(self) -> types.ModuleType | None:
                 self._plugin_module = None
         return self._plugin_module
 
-    def get_patches(self) -> PatchMap:
+    def get_all_patches(self) -> PatchMap:
         """Get a mapping of version to list of patches"""
         if self._patches is None:
             patches: PatchMap = {}
-            pattern = f"{self.override_module_name}-*"
-            prefix_len = len(pattern) - 1
+            version: Version | None
+            pattern = f"{self.override_module_name}*"
+            prefix_len = len(self.override_module_name) + 1
             for patchdir in self._patches_dir.glob(pattern):
-                version = Version(patchdir.name[prefix_len:])
-                patches[version] = sorted(patchdir.glob("*.patch"))
+                if patchdir.name == self.override_module_name:
+                    version = None
+                else:
+                    version = Version(patchdir.name[prefix_len:])
+                patches[version] = list(patchdir.glob("*.patch"))
+                # variant-specific patches
+                patches[version].extend(patchdir.joinpath(self.variant).glob("*.patch"))
+                patches[version].sort(key=lambda p: p.name)
+
             self._patches = patches
         return self._patches
 
+    def get_patches(self, version: Version) -> list[pathlib.Path]:
+        """Get patches for a version (and unversioned patches)"""
+        patchfiles: list[pathlib.Path] = []
+        patchmap = self.get_all_patches()
+        # unversioned patches
+        patchfiles.extend(patchmap.get(None, []))
+        # version-specific patches
+        patchfiles.extend(patchmap.get(version, []))
+        # sort by basename
+        patchfiles.sort(key=lambda p: p.name)
+        return patchfiles
+
     @property
     def has_config(self) -> bool:
         """Does the package have a config file?"""

src/fromager/sources.py

@@ -294,23 +294,21 @@ def patch_source(
     req: Requirement,
     version: Version,
 ) -> None:
+    pbi = ctx.package_build_info(req)
     patch_count = 0
-    for p in overrides.patches_for_requirement(
-        patches_dir=ctx.settings.patches_dir,
-        req=req,
-        version=version,
-    ):
+
+    for p in pbi.get_patches(version):
         _apply_patch(p, source_root_dir)
         patch_count += 1
 
     logger.debug("%s: applied %d patches", req.name, patch_count)
     # If no patch has been applied, call warn for old patch
-    if not patch_count:
-        _warn_for_old_patch(
-            req=req,
-            version=version,
-            patches_dir=ctx.settings.patches_dir,
-        )
+    patchmap = pbi.get_all_patches()
+    if not patch_count and patchmap:
+        for patchversion in sorted(patchmap):
+            logger.warning(
+                f"{req.name}: patch {patchversion} exists but will not be applied for version {version}"
+            )
 
 
 def _apply_patch(patch: pathlib.Path, source_root_dir: pathlib.Path):
@@ -319,23 +317,6 @@ def _apply_patch(patch: pathlib.Path, source_root_dir: pathlib.Path):
         external_commands.run(["patch", "-p1"], stdin=f, cwd=source_root_dir)
 
 
-def _warn_for_old_patch(
-    req: Requirement,
-    version: Version,
-    patches_dir: pathlib.Path,
-) -> None:
-    # Filter the patch directories using regex
-    patch_directories = overrides.get_versioned_patch_directories(
-        patches_dir=patches_dir, req=req
-    )
-
-    for dirs in patch_directories:
-        for p in dirs.iterdir():
-            logger.warning(
-                f"{req.name}: patch {p} exists but will not be applied for version {version}"
-            )
-
-
 def write_build_meta(
     unpack_dir: pathlib.Path,
     req: Requirement,

tests/test_overrides.py

@@ -1,41 +1,11 @@
-import pathlib
 from unittest import mock
 from unittest.mock import patch
 
 import pytest
-from packaging.requirements import Requirement
-from packaging.version import Version
 
 from fromager import overrides
 
 
-def test_patches_for_requirement(tmp_path: pathlib.Path):
-    patches_dir = tmp_path / "patches"
-    patches_dir.mkdir()
-
-    project_patch_dir = patches_dir / "project-1.2.3"
-    project_patch_dir.mkdir()
-
-    p1 = project_patch_dir / "001.patch"
-    p2 = project_patch_dir / "002.patch"
-    np1 = project_patch_dir / "not-a-patch.txt"
-
-    # Create all of the test files
-    for p in [p1, p2]:
-        p.write_text("this is a patch file")
-    for f in [np1]:
-        f.write_text("this is not a patch file")
-
-    results = list(
-        overrides.patches_for_requirement(
-            patches_dir=patches_dir,
-            req=Requirement("project"),
-            version=Version("1.2.3"),
-        )
-    )
-    assert results == [p1, p2]
-
-
 def test_invoke_override_with_exact_args():
     def foo(arg1, arg2):
         return arg1 is not None and arg2 is not None

tests/test_packagesettings.py

@@ -252,16 +252,37 @@ def test_pbi_test_pkg(testdata_context: context.WorkContext) -> None:
     sdist_root_dir = pathlib.Path("/sdist-root")
     assert pbi.build_dir(sdist_root_dir) == sdist_root_dir / "python"
 
-    patchdir = (
-        testdata_context.settings.patches_dir / f"{TEST_PKG.replace('-', '_')}-1.0.2"
+
+def test_pbi_test_pkg_patches(testdata_context: context.WorkContext) -> None:
+    pbi = testdata_context.settings.package_build_info(TEST_PKG)
+    norm_test_pkg = TEST_PKG.replace("-", "_")
+    unversioned_patchdir = testdata_context.settings.patches_dir / norm_test_pkg
+    versioned_patchdir = (
+        testdata_context.settings.patches_dir / f"{norm_test_pkg}-1.0.2"
     )
-    assert pbi.get_patches() == {
-        Version("1.0.2"): [
-            patchdir / "001-somepatch.patch",
-            patchdir / "002-otherpatch.patch",
-        ],
+    patch001 = versioned_patchdir / "001-somepatch.patch"
+    patch002 = versioned_patchdir / "002-otherpatch.patch"
+    patch004 = unversioned_patchdir / "cpu" / "004-cpu.patch"
+    patch005 = versioned_patchdir / "cpu" / "005-cpuver.patch"
+    patch010 = unversioned_patchdir / "010-unversioned.patch"
+
+    assert pbi.get_all_patches() == {
+        None: [patch004, patch010],
+        Version("1.0.2"): [patch001, patch002, patch005],
     }
-    assert pbi.get_patches() is pbi.get_patches()
+    assert pbi.get_all_patches() is pbi.get_all_patches()
+
+    assert pbi.get_patches(Version("1.0.2")) == [
+        patch001,
+        patch002,
+        patch004,
+        patch005,
+        patch010,
+    ]
+    assert pbi.get_patches(Version("1.0.1")) == [
+        patch004,
+        patch010,
+    ]
 
 
 def test_pbi_other(testdata_context: context.WorkContext) -> None:
@@ -289,12 +310,12 @@ def test_pbi_other(testdata_context: context.WorkContext) -> None:
         testdata_context.settings.patches_dir
         / f"{TEST_OTHER_PKG.replace('-', '_')}-1.0.0"
     )
-    assert pbi.get_patches() == {
+    assert pbi.get_all_patches() == {
         Version("1.0.0"): [
             patchdir / "001-mypatch.patch",
         ],
     }
-    assert pbi.get_patches() is pbi.get_patches()
+    assert pbi.get_all_patches() is pbi.get_all_patches()
 
 
 def test_type_envvars():