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

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/161307

>From 55eeaa53bd07d5a9cfc8c489f02b663abbda6cd5 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Mon, 29 Sep 2025 20:09:57 -0700
Subject: [PATCH] [MLIR][Python] Fix stubgen/PYTHONPATH collision/bug

---
 mlir/cmake/modules/AddMLIRPython.cmake | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 208cbdd1dd535..c01d0e06bcd7a 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(NATIVE_PATH _install_prefix NORMALIZE "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/..")
+  string(REGEX REPLACE "/$" "" _install_prefix "${_install_prefix}")
+  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()
+  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}"