Refactor local runtime info functions and improve structure

Author: laramiel

bazel/repo_rules/get_local_runtime_info.py

@@ -11,19 +11,41 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Returns information about the local Python runtime as JSON."""
 
+import glob
 import json
 import os
 import sys
 import sysconfig
+from typing import Any
 
 _IS_WINDOWS = sys.platform == "win32"
 _IS_DARWIN = sys.platform == "darwin"
 
 
-def _search_directories(get_config, base_executable):
+def _get_abi_flags(get_config) -> str:
+  """Returns the ABI flags for the Python runtime."""
+  # sys.abiflags may not exist, but it still may be set in the config.
+  abi_flags = getattr(sys, "abiflags", None)
+  if abi_flags is None:
+    abi_flags = get_config("ABIFLAGS") or get_config("abiflags") or ""
+  return abi_flags
+
+
+def _get_shlib_suffix(get_config) -> str:
+  """Returns the shared library suffix for the platform."""
+  if _IS_DARWIN:
+    return ".dylib"
+  if _IS_WINDOWS:
+    return ".dll"
+  suffix = get_config("SHLIB_SUFFIX")
+  if not suffix:
+    suffix = ".so"
+  return suffix
+
+
+def _search_directories(get_config, base_executable) -> list[str]:
   """Returns a list of library directories to search for shared libraries."""
   # There's several types of libraries with different names and a plethora
   # of settings, and many different config variables to check:
@@ -74,23 +96,12 @@ def _search_directories(get_config, base_executable):
         lib_dirs.append(os.path.join(os.path.dirname(exec_dir), "lib"))
 
   # Dedup and remove empty values, keeping the order.
-  lib_dirs = [v for v in lib_dirs if v]
-  return {k: None for k in lib_dirs}.keys()
+  return list(dict.fromkeys(d for d in lib_dirs if d))
 
 
-def _get_shlib_suffix(get_config) -> str:
-  """Returns the suffix for shared libraries."""
-  if _IS_DARWIN:
-    return ".dylib"
-  if _IS_WINDOWS:
-    return ".dll"
-  suffix = get_config("SHLIB_SUFFIX")
-  if not suffix:
-    suffix = ".so"
-  return suffix
-
-
-def _search_library_names(get_config, shlib_suffix):
+def _search_library_names(
+    get_config, version, abi_flags, shlib_suffix
+) -> list[str]:
   """Returns a list of library files to search for shared libraries."""
   # Quoting configure.ac in the cpython code base:
   # "INSTSONAME is the name of the shared library that will be use to install
@@ -117,57 +128,32 @@ def _search_library_names(get_config, shlib_suffix):
   # Set the prefix and suffix to construct the library name used for linking.
   # The suffix and version are set here to the default values for the OS,
   # since they are used below to construct "default" library names.
-  if _IS_DARWIN:
-    prefix = "lib"
-  elif _IS_WINDOWS:
+  if _IS_WINDOWS:
     prefix = ""
   else:
     prefix = "lib"
 
-  version = get_config("VERSION")
-
   # Ensure that the pythonXY.dll files are included in the search.
   lib_names.append(f"{prefix}python{version}{shlib_suffix}")
 
   # If there are ABIFLAGS, also add them to the python version lib search.
-  abiflags = get_config("ABIFLAGS") or get_config("abiflags") or ""
-  if abiflags:
-    lib_names.append(f"{prefix}python{version}{abiflags}{shlib_suffix}")
+  if abi_flags:
+    lib_names.append(f"{prefix}python{version}{abi_flags}{shlib_suffix}")
 
-  # Add the abi-version includes to the search list.
-  lib_names.append(f"{prefix}python{sys.version_info.major}{shlib_suffix}")
+  return list(dict.fromkeys(k for k in lib_names if k))
 
-  # Dedup and remove empty values, keeping the order.
-  lib_names = [v for v in lib_names if v]
-  return {k: None for k in lib_names}.keys()
-
-
-def _get_python_library_info(base_executable):
-  """Returns a dictionary with the static and dynamic python libraries."""
-  config_vars = sysconfig.get_config_vars()
-
-  # VERSION is X.Y in Linux/macOS and XY in Windows.  This is used to
-  # construct library paths such as python3.12, so ensure it exists.
-  if not config_vars.get("VERSION"):
-    if sys.platform == "win32":
-      config_vars["VERSION"] = (
-          f"{sys.version_info.major}{sys.version_info.minor}"
-      )
-    else:
-      config_vars["VERSION"] = (
-          f"{sys.version_info.major}.{sys.version_info.minor}"
-      )
-
-  shlib_suffix = _get_shlib_suffix(config_vars.get)
-  search_directories = _search_directories(config_vars.get, base_executable)
-  search_libnames = _search_library_names(config_vars.get, shlib_suffix)
-
-  interface_libraries = {}
-  dynamic_libraries = {}
-  static_libraries = {}
 
+def _do_library_search(
+    *,
+    search_directories: list[str],
+    libnames: list[str],
+) -> tuple[dict[str, None], dict[str, None], dict[str, None]]:
+  """Finds existing libnames in search_directories, returning the found libraries."""
+  static_libraries: dict[str, None] = {}
+  dynamic_libraries: dict[str, None] = {}
+  interface_libraries: dict[str, None] = {}
   for root_dir in search_directories:
-    for libname in search_libnames:
+    for libname in libnames:
       # Check whether the library exists.
       composed_path = os.path.join(root_dir, libname)
       if os.path.exists(composed_path) or os.path.isdir(composed_path):
@@ -178,9 +164,9 @@ def _get_python_library_info(base_executable):
 
       interface_path = None
       if libname.endswith(".dll"):
-        # On windows a .lib file may be an "import library" or a static library.
-        # The file could be inspected to determine which it is; typically python
-        # is used as a shared library.
+        # On windows a .lib file may be an "import library" or a static
+        # library. The file could be inspected to determine which it is;
+        # typically python is used as a shared library.
         #
         # On Windows, extensions should link with the pythonXY.lib interface
         # libraries.
@@ -195,32 +181,88 @@ def _get_python_library_info(base_executable):
       # Check whether an interface library exists.
       if interface_path and os.path.exists(interface_path):
         interface_libraries[interface_path] = None
+  return static_libraries, dynamic_libraries, interface_libraries
 
-  if hasattr(sys, "abiflags"):
-    abiflags = sys.abiflags
-  else:
-    abiflags = ""
+
+def _get_python_library_info(base_executable) -> dict[str, Any]:
+  """Returns a dictionary with the static and dynamic python libraries."""
+  config_vars = sysconfig.get_config_vars()
+
+  # VERSION is X.Y in Linux/macOS and XY in Windows.  This is used to
+  # construct library paths such as python3.12, so ensure it exists.
+  version = config_vars.get("VERSION")
+  if not version:
+    if _IS_WINDOWS:
+      version = f"{sys.version_info.major}{sys.version_info.minor}"
+    else:
+      version = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+  # sys.abiflags may not exist, but it still may be set in the config.
+  abi_flags = _get_abi_flags(config_vars.get)
+  shlib_suffix = _get_shlib_suffix(config_vars.get)
+  search_directories = _search_directories(config_vars.get, base_executable)
+  search_libnames = _search_library_names(
+      config_vars.get, version, abi_flags, shlib_suffix
+  )
+
+  # Search for the python libraries for the current python interpreter.
+  static_libraries, dynamic_libraries, interface_libraries = _do_library_search(
+      search_directories=search_directories,
+      libnames=search_libnames,
+  )
+
+  # Search for major version abi libraries which must be dynamic.
+  abi_search_libnames = [
+      s.replace(version, f"{sys.version_info.major}")
+      for s in search_libnames
+      if version in s
+  ]
+  _, abi_dynamic_libraries, abi_interface_libraries = _do_library_search(
+      search_directories=search_directories,
+      libnames=abi_search_libnames,
+  )
+
+  # Additional DLLs are needed on Windows to link properly.
+  dlls = []
+  if _IS_WINDOWS:
+    dlls.extend(
+        glob.glob(os.path.join(os.path.dirname(base_executable), "*.dll"))
+    )
+    dlls = [
+        x
+        for x in dlls
+        if x not in dynamic_libraries and x not in abi_dynamic_libraries
+    ]
+
+  def _unique_basenames(inputs: dict[str, None]) -> list[str]:
+    """Returns a list of paths, keeping only the first path for each basename."""
+    result = []
+    seen = set()
+    for k in inputs:
+      b = os.path.basename(k)
+      if b not in seen:
+        seen.add(b)
+        result.append(k)
+    return result
 
   # When no libraries are found it's likely that the python interpreter is not
   # configured to use shared or static libraries (minilinux).  If this seems
   # suspicious try running `uv tool run find_libpython --list-all -v`
   return {
-      "dynamic_libraries": list(dynamic_libraries.keys()),
-      "static_libraries": list(static_libraries.keys()),
-      "interface_libraries": list(interface_libraries.keys()),
-      "shlib_suffix": "" if _IS_WINDOWS else shlib_suffix,
-      "abi_flags": abiflags,
+      "dynamic_libraries": _unique_basenames(dynamic_libraries),
+      "static_libraries": _unique_basenames(static_libraries),
+      "interface_libraries": _unique_basenames(interface_libraries),
+      "abi_dynamic_libraries": _unique_basenames(abi_dynamic_libraries),
+      "abi_interface_libraries": _unique_basenames(abi_interface_libraries),
+      "abi_flags": abi_flags,
+      "shlib_suffix": shlib_suffix if _IS_DARWIN else "",
+      "additional_dlls": dlls,
   }
 
 
-def _get_base_executable():
+def _get_base_executable() -> str:
   """Returns the base executable path."""
-  try:
-    if sys._base_executable:  # pylint: disable=protected-access
-      return sys._base_executable  # pylint: disable=protected-access
-  except AttributeError:
-    pass
-  return sys.executable
+  return getattr(sys, "_base_executable", None) or sys.executable
 
 
 data = {