fix(windows): use extended paths in Python bootstraps (#4071)

Implicit long-path support is not universal across the Win32 API. The
documented set of APIs covered by the long-path opt-in does not include
DLL loading functions, e.g., LoadLibraryExW:


https://learn.microsoft.com/en-us/windows/win32/fileio/maximum-file-path-limitation#functions-without-max_path-restrictions

Always use extended-length paths in Windows bootstrap code and correctly
convert UNC paths to the \\?\UNC\ form.

Experienced this breakage in JAX after the project being switched over
to Bzlmod, which made some paths too long:
https://github.com/jax-ml/jax/actions/runs/31674295160/job/94365413380

```
    File "c:\botcode\w\bazel-out\x64_windows-opt\bin\jax\experimental\jax2tf\tests\multiprocess\jax2tf_multiprocess_test_cpu.exe.runfiles\rules_python++pip+jax_pypi_312_ml_dtypes_cp312_cp312_win_amd64_c1a95399\site-packages\ml_dtypes\_finfo.py", line 17, in <module>
      from ml_dtypes._ml_dtypes_ext import bfloat16
  ImportError: DLL load failed while importing _ml_dtypes_ext: The filename or extension is too long.
```

Already used as a patch in https://github.com/jax-ml/jax/pull/39961

---------

Co-authored-by: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
diff --git a/news/4071.fixed.md b/news/4071.fixed.md
new file mode 100644
index 0000000..0b24827
--- /dev/null
+++ b/news/4071.fixed.md
@@ -0,0 +1 @@
+(rules) Fixed loading Python extension modules from paths longer than `MAX_PATH` on Windows.
diff --git a/python/private/site_init_template.py b/python/private/site_init_template.py
index 12be98e..c5a4848 100644
--- a/python/private/site_init_template.py
+++ b/python/private/site_init_template.py
@@ -93,35 +93,62 @@
     if not _is_windows() or sys.version_info[0] < 3:
         return path
 
-    # Starting in Windows 10, version 1607(OS build 14393), MAX_PATH limitations have been
-    # removed from common Win32 file and directory functions.
-    # Related doc: https://docs.microsoft.com/en-us/windows/win32/fileio/maximum-file-path-limitation?tabs=cmd#enable-long-paths-in-windows-10-version-1607-and-later
-    import platform
-
-    win32_version = None
-    # Windows 2022 with Python 3.12.8 gives flakey errors, so try a couple times.
-    for _ in range(3):
-        try:
-            win32_version = platform.win32_ver()[1]
-            break
-        except (ValueError, KeyError):
-            pass
-    if win32_version and win32_version >= "10.0.14393":
-        return path
-
     # import sysconfig only now to maintain python 2.6 compatibility
     import sysconfig
 
     if sysconfig.get_platform() == "mingw":
         return path
 
-    # Lets start the unicode fun
-    unicode_prefix = "\\\\?\\"
-    if path.startswith(unicode_prefix):
+    # Implicit long-path support is not universal across the Win32 API. For
+    # example, DLL loading still requires an explicit extended-length prefix.
+    extended_path_prefix = "\\\\?\\"
+    if path.startswith(extended_path_prefix):
         return path
 
     # os.path.abspath returns a normalized absolute path
-    return unicode_prefix + os.path.abspath(path)
+    path = os.path.abspath(path)
+    if path.startswith("\\\\"):
+        return extended_path_prefix + "UNC\\" + path[2:]
+    return extended_path_prefix + path
+
+
+def _install_windows_extension_finder():
+    """Use extended-length paths when loading long Windows extension paths."""
+    if not _is_windows() or sys.version_info[0] < 3:
+        return
+
+    # import these only now to maintain Python 2.6 compatibility
+    import importlib.machinery
+    import sysconfig
+
+    if sysconfig.get_platform() == "mingw":
+        return
+
+    class _WindowsExtensionPathFinder(importlib.machinery.PathFinder):
+        @classmethod
+        def find_spec(cls, fullname, path=None, target=None):
+            spec = super().find_spec(fullname, path, target)
+            if (
+                spec is None
+                or not isinstance(spec.loader, importlib.machinery.ExtensionFileLoader)
+                or len(os.path.abspath(spec.origin)) < 260
+            ):
+                return spec
+
+            # The registry opt-in for long paths only applies to documented
+            # file and directory APIs. It does not include DLL loading APIs,
+            # e.g. LoadLibraryExW. Prefix the actual extension filename instead
+            # of sys.path entries so other APIs continue to receive normal paths.
+            # https://learn.microsoft.com/en-us/windows/win32/fileio/maximum-file-path-limitation#functions-without-max_path-restrictions
+            extended_path = _get_windows_path_with_unc_prefix(spec.origin)
+            spec.origin = extended_path
+            spec.loader.path = extended_path
+            return spec
+
+    for index, finder in enumerate(sys.meta_path):
+        if finder is importlib.machinery.PathFinder:
+            sys.meta_path[index] = _WindowsExtensionPathFinder
+            return
 
 
 def _search_path(name):
@@ -143,7 +170,6 @@
     def _maybe_add_path(path, reason):
         if path in seen:
             return
-        path = _get_windows_path_with_unc_prefix(path)
         if _is_windows():
             path = path.replace("/", os.sep)
 
@@ -241,4 +267,5 @@
 _fixup_sys_base_executable()
 
 COVERAGE_SETUP = _setup_sys_path()
+_install_windows_extension_finder()
 _print_verbose("DONE")
diff --git a/tests/bootstrap_impls/windows_long_path/BUILD.bazel b/tests/bootstrap_impls/windows_long_path/BUILD.bazel
new file mode 100644
index 0000000..2195eb6
--- /dev/null
+++ b/tests/bootstrap_impls/windows_long_path/BUILD.bazel
@@ -0,0 +1,33 @@
+load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
+load("//python:py_test.bzl", "py_test")
+
+# buildifier: disable=bzl-visibility
+load("//python/cc:py_extension.bzl", "py_extension")
+
+_LONG_IMPORT_PATH = "/".join([
+    "long_path_segment_000000000000000000000000000001",
+    "long_path_segment_000000000000000000000000000002",
+    "long_path_segment_000000000000000000000000000003",
+    "long_path_segment_000000000000000000000000000004",
+])
+
+py_extension(
+    name = "ext_long_path_source",
+    srcs = ["ext_long_path.c"],
+    target_compatible_with = ["@platforms//os:windows"],
+)
+
+copy_file(
+    name = "ext_long_path",
+    src = ":ext_long_path_source",
+    out = _LONG_IMPORT_PATH + "/ext_long_path.pyd",
+    target_compatible_with = ["@platforms//os:windows"],
+)
+
+py_test(
+    name = "py_extension_long_path_test",
+    srcs = ["py_extension_long_path_test.py"],
+    data = [":ext_long_path"],
+    imports = [_LONG_IMPORT_PATH],
+    target_compatible_with = ["@platforms//os:windows"],
+)
diff --git a/tests/bootstrap_impls/windows_long_path/ext_long_path.c b/tests/bootstrap_impls/windows_long_path/ext_long_path.c
new file mode 100644
index 0000000..89ea6e5
--- /dev/null
+++ b/tests/bootstrap_impls/windows_long_path/ext_long_path.c
@@ -0,0 +1,22 @@
+#include <Python.h>
+
+static PyObject* get_magic_number(PyObject* self, PyObject* args) {
+    return PyLong_FromLong(42);
+}
+
+static PyMethodDef ModuleMethods[] = {
+    {"get_magic_number", get_magic_number, METH_NOARGS, "Returns 42."},
+    {NULL, NULL, 0, NULL}
+};
+
+static struct PyModuleDef ext_long_path_module = {
+    PyModuleDef_HEAD_INIT,
+    "ext_long_path",
+    NULL,
+    -1,
+    ModuleMethods
+};
+
+PyMODINIT_FUNC PyInit_ext_long_path(void) {
+    return PyModule_Create(&ext_long_path_module);
+}
diff --git a/tests/bootstrap_impls/windows_long_path/py_extension_long_path_test.py b/tests/bootstrap_impls/windows_long_path/py_extension_long_path_test.py
new file mode 100644
index 0000000..f8610ba
--- /dev/null
+++ b/tests/bootstrap_impls/windows_long_path/py_extension_long_path_test.py
@@ -0,0 +1,22 @@
+import sys
+import unittest
+
+import ext_long_path  # pyrefly: ignore[missing-import]
+
+
+class PyExtensionLongPathTest(unittest.TestCase):
+    def test_extension_is_loaded_from_extended_length_path(self):
+        self.assertEqual(ext_long_path.get_magic_number(), 42)
+        self.assertGreaterEqual(len(ext_long_path.__file__), 260)
+        self.assertTrue(ext_long_path.__file__.startswith("\\\\?\\"))
+
+    def test_other_python_paths_are_not_extended_length_paths(self):
+        self.assertFalse(sys.prefix.startswith("\\\\?\\"))
+        self.assertFalse(
+            any(path.startswith("\\\\?\\") for path in sys.path),
+            sys.path,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()