[Mlir-commits] [mlir] e911eba - [MLIR][Python] Fix stubgen/PYTHONPATH collision/bug (#161307)

[llvmlistbot at llvm.org]

Author: Maksim Levental
Date: 2025-09-30T05:37:54Z
New Revision: e911eba3ef6990da2c4620c573b99a4f20f53da4

URL: https://github.com/llvm/llvm-project/commit/e911eba3ef6990da2c4620c573b99a4f20f53da4
DIFF: https://github.com/llvm/llvm-project/commit/e911eba3ef6990da2c4620c573b99a4f20f53da4.diff

LOG: [MLIR][Python] Fix stubgen/PYTHONPATH collision/bug (#161307)

If `PYTHONPATH` is set and points to the build location of the python
bindings package then when stubgen runs, `_mlir` will get imported twice
and bad things will happen (e.g., `Assertion !instance && “PyGlobals
already constructed”’`). This happens because `mlir` is a namespace
package and the importer/loader can't distinguish between
`mlir._mlir_libs._mlir` and `_mlir_libs._mlir` imported from `CWD`. Or
something like that. The fix is to filter out any entries in
`PYTHONPATH` that end in `MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/..` (e.g.,
`python_packages/mlir_core/`).

Added: 
    

Modified: 
    mlir/cmake/modules/AddMLIRPython.cmake

Removed: 
    


################################################################################
diff  --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 208cbdd1dd535..fa6aec8a603a9 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -173,9 +173,32 @@ function(mlir_generate_type_stubs)
   if(ARG_VERBOSE)
     message(STATUS "Generating type-stubs outputs ${_generated_type_stubs}")
   endif()
+
+  # If PYTHONPATH is set and points to the build location of the python package then when stubgen runs, _mlir will get
+  # imported twice and bad things will happen (e.g., Assertion `!instance && “PyGlobals already constructed”’).
+  # This happens because mlir is a namespace package and the importer/loader can't distinguish between
+  # mlir._mlir_libs._mlir and _mlir_libs._mlir imported from CWD.
+  # So try to filter out any entries in PYTHONPATH that end in "MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/.."
+  # (e.g., python_packages/mlir_core/).
+  set(_pythonpath "$ENV{PYTHONPATH}")
+  cmake_path(CONVERT "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." TO_NATIVE_PATH_LIST _install_prefix NORMALIZE)
+  if(WIN32)
+    set(_path_sep ";")
+    set(_trailing_sep "\\")
+  else()
+    set(_path_sep ":")
+    set(_trailing_sep "/")
+    # `;` is the CMake list delimiter so Windows paths are automatically lists
+    # and Unix paths can be made into lists by replacing `:` with `;`
+    string(REPLACE "${_path_sep}" ";" _pythonpath "${_pythonpath}")
+  endif()
+  string(REGEX REPLACE "${_trailing_sep}$" "" _install_prefix "${_install_prefix}")
+  list(FILTER _pythonpath EXCLUDE REGEX "(${_install_prefix}|${_install_prefix}${_trailing_sep})$")
+  # Note, ${_pythonpath} is a list but "${_pythonpath}" is not a list - it's a string with ";" chars in it.
+  string(JOIN "${_path_sep}" _pythonpath ${_pythonpath})
   add_custom_command(
     OUTPUT ${_generated_type_stubs}
-    COMMAND ${_nb_stubgen_cmd}
+    COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH="${_pythonpath}" ${_nb_stubgen_cmd}
     WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
     DEPENDS "${ARG_DEPENDS_TARGETS}"
     DEPFILE "${_depfile}"