Introduce --custom-patch option

README.md

@@ -70,6 +70,7 @@ The path separator to use in the following options is `';'` on Windows and `':'`
   - On Windows, `--namespace-pkg package1.package2;package3` declares `package1`, `package1\package2`, and `package3` as namespace packages.
 - `--include-symbols`: include `.pdb` symbol files with the vendored DLLs. To be included, a symbol file must be in the same directory as the DLL and have the same filename before the extension, e.g. `example.dll` and `example.pdb`.
 - `--include-imports`: include `.lib` import library files with the vendored DLLs. To be included, an import library file must be in the same directory as the DLL and have the same filename before the extension, e.g. `example.dll` and `example.lib`.
+- `--custom-patch`: Normally, we patch or create `__init__.py` in each top-level package to inject code that adds the vendored DLL location to the DLL search path at runtime. To precisely control where the DLL search path is modified, use this option to instead specify the exact location(s) to place the patch. When this option is enabled, every line in a `.py` file consisting of the unindented comment `# delvewheel: patch` is replaced with the patch.
 
 ## Version Scheme
 

delvewheel/__main__.py

@@ -55,7 +55,9 @@ def main():
     parser_repair.add_argument('--no-mangle-all', action='store_true', help="don't mangle any DLL names")
     parser_repair.add_argument('--strip', action='store_true', help='strip DLLs that contain trailing data when name-mangling')
     parser_repair.add_argument('-L', '--lib-sdir', default='.libs', type=_dir_suffix, help='directory suffix to store vendored DLLs (default .libs)')
-    parser_repair.add_argument('--namespace-pkg', default='', metavar='PKGS', type=_namespace_pkgs, help=f'namespace package(s), {os.pathsep!r}-delimited')
+    group = parser_repair.add_mutually_exclusive_group()
+    group.add_argument('--namespace-pkg', default='', metavar='PKGS', type=_namespace_pkgs, help=f'namespace package(s), {os.pathsep!r}-delimited')
+    group.add_argument('--custom-patch', action='store_true', help='customize the location of the DLL search path patch')
     parser_repair.add_argument('--no-diagnostic', action='store_true', help=argparse.SUPPRESS)  # don't write diagnostic information to DELVEWHEEL metadata file
     parser_repair.add_argument('--include-symbols', action='store_true', help='include .pdb symbol files with vendored DLLs')
     parser_repair.add_argument('--include-imports', action='store_true', help='include .lib import library files with the vendored DLLs')
@@ -82,7 +84,7 @@ def main():
             else:  # args.command == 'repair'
                 no_mangles = set(dll_name.lower() for dll_name in os.pathsep.join(args.no_mangle).split(os.pathsep) if dll_name)
                 namespace_pkgs = set(tuple(namespace_pkg.split('.')) for namespace_pkg in args.namespace_pkg.split(os.pathsep) if namespace_pkg)
-                wr.repair(args.target, no_mangles, args.no_mangle_all, args.strip, args.lib_sdir, not args.no_diagnostic and 'SOURCE_DATE_EPOCH' not in os.environ, namespace_pkgs, args.include_symbols, args.include_imports)
+                wr.repair(args.target, no_mangles, args.no_mangle_all, args.strip, args.lib_sdir, not args.no_diagnostic and 'SOURCE_DATE_EPOCH' not in os.environ, namespace_pkgs, args.include_symbols, args.include_imports, args.custom_patch)
     else:  # args.command == 'needed'
         for dll_name in sorted(_dll_utils.get_direct_needed(args.file, args.v), key=str.lower):
             print(dll_name)

delvewheel/_wheel_repair.py

@@ -493,7 +493,8 @@ def _patch_package(self, package_dir: str, namespace_pkgs: typing.Set[typing.Tup
         load_order_filename is the name of the .load-order file, or None if the
             file is not used
         depth is the number of parent directories to traverse to reach the
-            site-packages directory at runtime starting from package_dir"""
+            site-packages directory at runtime starting from inside
+            package_dir"""
         package_name = os.path.basename(package_dir)
         namespace_root_ext_modules = set()
         if any(x[0] == package_name for x in namespace_pkgs):
@@ -510,6 +511,44 @@ def _patch_package(self, package_dir: str, namespace_pkgs: typing.Set[typing.Tup
             self._patch_py_file(self._get_init(package_dir) or os.path.join(package_dir, '__init__.py'), libs_dir, load_order_filename, depth)
         return namespace_root_ext_modules
 
+    def _patch_custom(self, item_path: str, libs_dir: str, load_order_filename: str, depth: int) -> bool:
+        """Patch a package or .py file so that vendored DLLs can be found at
+        runtime. The patch is placed at every line consisting of the comment
+        '# delvewheel: patch'. Return True iff the patch was applied at least
+        once.
+
+        item_path is the absolute path to the extracted package or .py file
+        libs_dir is the name of the directory where DLLs are stored.
+        load_order_filename is the name of the .load-order file, or None if the
+            file is not used
+        depth is the number of parent directories to traverse to reach the
+            site-packages directory at runtime starting from item_path"""
+        if os.path.isdir(item_path):
+            result = False
+            for item in os.listdir(item_path):
+                if os.path.isdir(new_item_path := os.path.join(item_path, item)) or \
+                        os.path.isfile(new_item_path) and item[-3:].lower() == '.py':
+                    result |= self._patch_custom(new_item_path, libs_dir, load_order_filename, depth + 1)
+            return result
+        with open(item_path) as file:
+            contents = file.read()
+        if not re.search(pattern := '^# *delvewheel *: *repair *$', contents, flags=re.MULTILINE):
+            return False
+        print(f'patching {os.path.relpath(item_path, self._extract_dir)}', end='')
+        with open(item_path, newline='') as file:
+            line = file.readline()
+        for newline in ('\r\n', '\r', '\n'):
+            if line.endswith(newline):
+                break
+        else:
+            newline = '\r\n'
+        patch_py_contents = self._patch_py_contents(False, libs_dir, load_order_filename, depth).rstrip()
+        contents, count = re.subn(pattern, patch_py_contents, contents, flags=re.MULTILINE)
+        with open(item_path, 'w', newline=newline) as file:
+            file.write(contents)
+        print(f' (count {count})' if count > 1 else '')
+        return True
+
     def _get_repair_version(self) -> str:
         """If this wheel has already been repaired, return the delvewheel
         version that performed the repair or '(unknown version)' if the version
@@ -689,7 +728,8 @@ def repair(
             log_diagnostics: bool,
             namespace_pkgs: typing.Set[typing.Tuple[str]],
             include_symbols: bool,
-            include_imports: bool) -> None:
+            include_imports: bool,
+            custom_patch: bool) -> None:
         """Repair the wheel in a manner similar to auditwheel.
 
         target is the target directory for storing the repaired wheel
@@ -704,7 +744,9 @@ def repair(
             corresponding to the namespace packages. Each path is represented
             as a tuple of path components
         include_symbols is True if .pdb symbol files should be included with
-            the vendored DLLs"""
+            the vendored DLLs
+        custom_patch is True to indicate that the DLL patch location is
+            custom"""
         print(f'repairing {self._whl_path}')
 
         # check whether wheel has already been repaired
@@ -857,65 +899,81 @@ def repair(
         else:
             load_order_filename = None
 
-        # Patch each package to load dependent DLLs from correct location at
-        # runtime.
-        #
-        # However, if a module and a folder are next to each other and have the
-        # same name and case,
-        # - If the folder does not contain __init__.py, do not patch
-        #   the folder. Otherwise, the import resolution order of the module
-        #   and the folder may be swapped.
-        # - If the folder contains __init__.py, patch the module (if it is pure
-        #   Python) and the folder.
         dist_info_foldername = f'{self._distribution_name}-{self._version}.dist-info'
-        namespace_root_ext_modules = set()
-        for item in os.listdir(self._extract_dir):
-            if os.path.isdir(package_dir := os.path.join(self._extract_dir, item)) and \
-                    item != dist_info_foldername and \
-                    item != os.path.basename(self._data_dir) and \
-                    item != libs_dir_name and \
-                    (item not in self._root_level_module_names(self._extract_dir) or self._get_init(package_dir)):
-                namespace_root_ext_modules.update(self._patch_package(package_dir, namespace_pkgs, libs_dir_name, load_order_filename, 1))
-        for extra_dir in (self._purelib_dir, self._platlib_dir):
-            if os.path.isdir(extra_dir):
-                for item in os.listdir(extra_dir):
-                    if os.path.isdir(package_dir := os.path.join(extra_dir, item)) and \
-                            (item not in self._root_level_module_names(self._extract_dir) or self._get_init(package_dir)):
-                        namespace_root_ext_modules.update(self._patch_package(package_dir, namespace_pkgs, libs_dir_name, load_order_filename, 1))
-
-        # Copy libraries next to all extension modules that are at the root of
-        # a namespace package. If a namespace package contains extension
-        # modules that are split across at least 2 of the following:
-        # 1. the wheel root,
-        # 2. the platlib directory,
-        # 3. the purelib directory,
-        # then copy the libraries to the first in the above list containing
-        # this namespace package.
-        if namespace_root_ext_modules:
-            dirnames = set(self._get_site_packages_relpath(os.path.dirname(x)) for x in namespace_root_ext_modules)
-            filenames = set(map(os.path.basename, namespace_root_ext_modules))
-            warnings.warn(
-                f'Namespace package{"s" if len(dirnames) > 1 else ""} '
-                f'{os.pathsep.join(dirnames)} contain'
-                f'{"s" if len(dirnames) == 1 else ""} root-level extension '
-                f'module{"s" if len(filenames) > 1 else ""} '
-                f'{os.pathsep.join(filenames)} and need'
-                f'{"s" if len(dirnames) == 1 else ""} '
-                f'{"an " if len(dirnames) == 1 else ""}extra '
-                f'cop{"ies" if len(dirnames) > 1 else "y"} of the '
-                'vendored DLLs. To avoid duplicate DLLs, move extension '
-                'modules into regular (non-namespace) packages.')
-            dirnames = list(set(map(os.path.dirname, namespace_root_ext_modules)))
-            dirnames.sort(key=self._namespace_pkg_sortkey)
-            seen_relative = set()
-            for dirname in dirnames:
-                if (dirname_relative := self._get_site_packages_relpath(dirname)) not in seen_relative:
-                    for filename in os.listdir(libs_dir):
-                        filepath = os.path.join(libs_dir, filename)
-                        if self._verbose >= 1:
-                            print(f'copying {filepath} -> {os.path.join(dirname, filename)}')
-                        shutil.copy2(filepath, dirname)
-                    seen_relative.add(dirname_relative)
+        if custom_patch:
+            # replace all instances of '# delvewheel: patch' with the patch
+            custom_patch_occurred = False
+            for item in os.listdir(self._extract_dir):
+                if os.path.isdir(item_path := os.path.join(self._extract_dir, item)) and item != dist_info_foldername and item != os.path.basename(self._data_dir) and item != libs_dir_name or \
+                        os.path.isfile(item_path) and item_path[-3:].lower() == '.py':
+                    custom_patch_occurred |= self._patch_custom(item_path, libs_dir_name, load_order_filename, 0)
+            for extra_dir in (self._purelib_dir, self._platlib_dir):
+                if os.path.isdir(extra_dir):
+                    for item in os.listdir(extra_dir):
+                        if os.path.isdir(item_path := os.path.join(extra_dir, item)) or \
+                                os.path.isfile(item_path) and item_path[-3:].lower() == '.py':
+                            custom_patch_occurred |= self._patch_custom(item_path, libs_dir_name, load_order_filename, 0)
+            if not custom_patch_occurred:
+                raise RuntimeError("'# delvewheel: patch' comment not found")
+        else:
+            # Patch each package to load dependent DLLs from correct location
+            # at runtime.
+            #
+            # However, if a module and a folder are next to each other and have
+            # the same name and case,
+            # - If the folder does not contain __init__.py, do not patch
+            #   the folder. Otherwise, the import resolution order of the
+            #   module and the folder may be swapped.
+            # - If the folder contains __init__.py, patch the module (if it is
+            #   pure Python) and the folder.
+            namespace_root_ext_modules = set()
+            for item in os.listdir(self._extract_dir):
+                if os.path.isdir(package_dir := os.path.join(self._extract_dir, item)) and \
+                        item != dist_info_foldername and \
+                        item != os.path.basename(self._data_dir) and \
+                        item != libs_dir_name and \
+                        (item not in self._root_level_module_names(self._extract_dir) or self._get_init(package_dir)):
+                    namespace_root_ext_modules.update(self._patch_package(package_dir, namespace_pkgs, libs_dir_name, load_order_filename, 1))
+            for extra_dir in (self._purelib_dir, self._platlib_dir):
+                if os.path.isdir(extra_dir):
+                    for item in os.listdir(extra_dir):
+                        if os.path.isdir(package_dir := os.path.join(extra_dir, item)) and \
+                                (item not in self._root_level_module_names(self._extract_dir) or self._get_init(package_dir)):
+                            namespace_root_ext_modules.update(self._patch_package(package_dir, namespace_pkgs, libs_dir_name, load_order_filename, 1))
+
+            # Copy libraries next to all extension modules that are at the root of
+            # a namespace package. If a namespace package contains extension
+            # modules that are split across at least 2 of the following:
+            # 1. the wheel root,
+            # 2. the platlib directory,
+            # 3. the purelib directory,
+            # then copy the libraries to the first in the above list containing
+            # this namespace package.
+            if namespace_root_ext_modules:
+                dirnames = set(self._get_site_packages_relpath(os.path.dirname(x)) for x in namespace_root_ext_modules)
+                filenames = set(map(os.path.basename, namespace_root_ext_modules))
+                warnings.warn(
+                    f'Namespace package{"s" if len(dirnames) > 1 else ""} '
+                    f'{os.pathsep.join(dirnames)} contain'
+                    f'{"s" if len(dirnames) == 1 else ""} root-level '
+                    f'extension module{"s" if len(filenames) > 1 else ""} '
+                    f'{os.pathsep.join(filenames)} and need'
+                    f'{"s" if len(dirnames) == 1 else ""} '
+                    f'{"an " if len(dirnames) == 1 else ""}extra '
+                    f'cop{"ies" if len(dirnames) > 1 else "y"} of the '
+                    'vendored DLLs. To avoid duplicate DLLs, move extension '
+                    'modules into regular (non-namespace) packages.')
+                dirnames = list(set(map(os.path.dirname, namespace_root_ext_modules)))
+                dirnames.sort(key=self._namespace_pkg_sortkey)
+                seen_relative = set()
+                for dirname in dirnames:
+                    if (dirname_relative := self._get_site_packages_relpath(dirname)) not in seen_relative:
+                        for filename in os.listdir(libs_dir):
+                            filepath = os.path.join(libs_dir, filename)
+                            if self._verbose >= 1:
+                                print(f'copying {filepath} -> {os.path.join(dirname, filename)}')
+                            shutil.copy2(filepath, dirname)
+                        seen_relative.add(dirname_relative)
 
         if load_order_filename is not None:
             # Create .load-order file containing list of DLLs to load during

tests/run_tests.py

@@ -1301,6 +1301,22 @@ def test_free_threaded(self):
         """Free-threaded wheel can be repaired"""
         check_call(['delvewheel', 'repair', '--add-path', 'simpleext/x64', 'simpleext/simpleext-0.0.1-cp313-cp313t-win_amd64.whl'])
 
+    def test_mutually_exclusive(self):
+        """--namespace-pkg and --custom-patch can't both be used"""
+        with self.assertRaises(subprocess.CalledProcessError):
+            check_call(['delvewheel', 'repair', '--add-path', 'simpleext/x64', '--namespace-pkg', 'ns0', '--custom-patch', 'simpleext/simpleext-0.0.1-0namespace-cp312-cp312-win_amd64.whl'])
+
+    def test_custom_none(self):
+        """Exception is raised when --custom-patch is specified and no location
+        was found."""
+        with self.assertRaises(subprocess.CalledProcessError):
+            check_call(['delvewheel', 'repair', '--add-path', 'simpleext/x64', '--custom-patch', 'simpleext/simpleext-0.0.1-0custom-cp312-cp312-win_amd64.whl'])
+
+    def test_custom(self):
+        """--custom-patch with multiple locations"""
+        p = subprocess.run(['delvewheel', 'repair', '--add-path', 'simpleext/x64', '--custom-patch', 'simpleext/simpleext-0.0.1-1custom-cp312-cp312-win_amd64.whl'], capture_output=True, text=True, check=True)
+        self.assertEqual(8, p.stdout.count('patching '))
+        self.assertEqual(1, p.stdout.count(' (count 2)'))
 
 class NeededTestCase(unittest.TestCase):
     """Tests for delvewheel needed"""