Skip to content

Commit 74ad737

Browse files
maksleventalmahesh-attarde
authored andcommitted
[MLIR][Python] Fix stubgen/PYTHONPATH collision/bug (llvm#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/`).
1 parent 1bb7c5e commit 74ad737

File tree

1 file changed

+24
-1
lines changed

1 file changed

+24
-1
lines changed

mlir/cmake/modules/AddMLIRPython.cmake

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,9 +173,32 @@ function(mlir_generate_type_stubs)
173173
if(ARG_VERBOSE)
174174
message(STATUS "Generating type-stubs outputs ${_generated_type_stubs}")
175175
endif()
176+
177+
# If PYTHONPATH is set and points to the build location of the python package then when stubgen runs, _mlir will get
178+
# imported twice and bad things will happen (e.g., Assertion `!instance && “PyGlobals already constructed”’).
179+
# This happens because mlir is a namespace package and the importer/loader can't distinguish between
180+
# mlir._mlir_libs._mlir and _mlir_libs._mlir imported from CWD.
181+
# So try to filter out any entries in PYTHONPATH that end in "MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/.."
182+
# (e.g., python_packages/mlir_core/).
183+
set(_pythonpath "$ENV{PYTHONPATH}")
184+
cmake_path(CONVERT "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." TO_NATIVE_PATH_LIST _install_prefix NORMALIZE)
185+
if(WIN32)
186+
set(_path_sep ";")
187+
set(_trailing_sep "\\")
188+
else()
189+
set(_path_sep ":")
190+
set(_trailing_sep "/")
191+
# `;` is the CMake list delimiter so Windows paths are automatically lists
192+
# and Unix paths can be made into lists by replacing `:` with `;`
193+
string(REPLACE "${_path_sep}" ";" _pythonpath "${_pythonpath}")
194+
endif()
195+
string(REGEX REPLACE "${_trailing_sep}$" "" _install_prefix "${_install_prefix}")
196+
list(FILTER _pythonpath EXCLUDE REGEX "(${_install_prefix}|${_install_prefix}${_trailing_sep})$")
197+
# Note, ${_pythonpath} is a list but "${_pythonpath}" is not a list - it's a string with ";" chars in it.
198+
string(JOIN "${_path_sep}" _pythonpath ${_pythonpath})
176199
add_custom_command(
177200
OUTPUT ${_generated_type_stubs}
178-
COMMAND ${_nb_stubgen_cmd}
201+
COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH="${_pythonpath}" ${_nb_stubgen_cmd}
179202
WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
180203
DEPENDS "${ARG_DEPENDS_TARGETS}"
181204
DEPFILE "${_depfile}"

0 commit comments

Comments
 (0)