feat(CMakeLists.txt): custom Python executable path

CMakeLists.txt

@@ -39,59 +39,108 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -pthread -fPIC -fopenmp")
 set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3")
 set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} -O3")
 
+function(system)
+    set(options STRIP)
+    set(oneValueArgs OUTPUT_VARIABLE ERROR_VARIABLE WORKING_DIRECTORY)
+    set(multiValueArgs COMMAND)
+    cmake_parse_arguments(SYSTEM
+                          "${options}"
+                          "${oneValueArgs}"
+                          "${multiValueArgs}"
+                          "${ARGN}")
+
+    if (NOT DEFINED SYSTEM_WORKING_DIRECTORY)
+        set(SYSTEM_WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}")
+    endif()
+
+    execute_process(
+        COMMAND ${SYSTEM_COMMAND}
+        OUTPUT_VARIABLE STDOUT
+        ERROR_VARIABLE STDERR
+        WORKING_DIRECTORY "${SYSTEM_WORKING_DIRECTORY}"
+    )
+    if("${SYSTEM_STRIP}")
+        string(STRIP "${STDOUT}" STDOUT)
+        string(STRIP "${STDERR}" STDERR)
+    endif()
+
+    set("${SYSTEM_OUTPUT_VARIABLE}" "${STDOUT}" PARENT_SCOPE)
+    if (DEFINED SYSTEM_ERROR_VARIABLE)
+        set("${SYSTEM_ERROR_VARIABLE}" "${STDERR}" PARENT_SCOPE)
+    endif()
+endfunction()
+
+if (NOT DEFINED PYTHON_EXECUTABLE)
+    set(PYTHON_EXECUTABLE python3)
+endif()
+
+system(
+    STRIP OUTPUT_VARIABLE PYTHON_EXECUTABLE
+    COMMAND which "${PYTHON_EXECUTABLE}"
+)
+
+system(
+    STRIP OUTPUT_VARIABLE PYTHON_VERSION
+    COMMAND "${PYTHON_EXECUTABLE}" --version
+)
+
+message("-- Use Python version: ${PYTHON_VERSION}")
+message("-- Use Python executable: \"${PYTHON_EXECUTABLE}\"")
+
 if(NOT DEFINED PYTHON_INCLUDE_DIR)
     message("-- Auto detecting Python include directory...")
-    execute_process (
-        COMMAND python3 -c "import sys; import os; path = sys.path[2]; splited_path = path.split('/'); splited_path[-2] = 'include'; print(os.path.join('/', *splited_path))"
-        OUTPUT_VARIABLE PYTHON_INCLUDE_DIR)
-    string(STRIP ${PYTHON_INCLUDE_DIR} PYTHON_INCLUDE_DIR)
+    system(
+        STRIP OUTPUT_VARIABLE PYTHON_INCLUDE_DIR
+        COMMAND "${PYTHON_EXECUTABLE}" -c "print(__import__('sysconfig').get_path('include'))"
+    )
 endif()
 
 if("${PYTHON_INCLUDE_DIR}" STREQUAL "")
     message(FATAL_ERROR "-- Python include directory not found")
 else()
-    message("-- Detected Python include directory: ${PYTHON_INCLUDE_DIR}")
+    message("-- Detected Python include directory: \"${PYTHON_INCLUDE_DIR}\"")
     include_directories(${PYTHON_INCLUDE_DIR})
 endif()
 
 if(NOT DEFINED TORCH_INCLUDE_PATH)
     message("-- Auto detecting PyTorch include directory...")
-    execute_process (
-        COMMAND python3 -c "from torch.utils import cpp_extension; print(cpp_extension.include_paths()[0], end='')"
-        OUTPUT_VARIABLE TORCH_INCLUDE_PATH)
-    string(STRIP ${TORCH_INCLUDE_PATH} TORCH_INCLUDE_PATH)
+    system(
+        STRIP OUTPUT_VARIABLE TORCH_INCLUDE_PATH
+        COMMAND "${PYTHON_EXECUTABLE}" -c "print(__import__('torch.utils.cpp_extension', fromlist=[None]).include_paths()[0])"
+    )
 endif()
 
 if("${TORCH_INCLUDE_PATH}" STREQUAL "")
     message(FATAL_ERROR "-- Torch include directory not found")
 else()
-    message("-- Detected Torch include directory: ${TORCH_INCLUDE_PATH}")
+    message("-- Detected Torch include directory: \"${TORCH_INCLUDE_PATH}\"")
     include_directories(${TORCH_INCLUDE_PATH})
 endif()
 
-
 if(NOT DEFINED TORCH_LIBRARY_PATH)
     message("-- Auto detecting PyTorch library directory...")
-    execute_process (
-        COMMAND python3 -c "from torch.utils import cpp_extension; print(cpp_extension.library_paths()[0], end='')"
-        OUTPUT_VARIABLE TORCH_LIBRARY_PATH)
-    string(STRIP ${TORCH_LIBRARY_PATH} TORCH_LIBRARY_PATH)
+    system(
+        STRIP OUTPUT_VARIABLE TORCH_LIBRARY_PATH
+        COMMAND "${PYTHON_EXECUTABLE}" -c "print(__import__('torch.utils.cpp_extension', fromlist=[None]).library_paths()[0])"
+    )
 endif()
 
 if("${TORCH_LIBRARY_PATH}" STREQUAL "")
     message(FATAL_ERROR "-- Torch library directory not found")
 else()
-    message("-- Detected Torch library directory: ${TORCH_LIBRARY_PATH}")
+    message("-- Detected Torch library directory: \"${TORCH_LIBRARY_PATH}\"")
 endif()
 
+unset(TORCH_LIBRARIES)
+foreach(VAR_PATH ${TORCH_LIBRARY_PATH})
+    file(GLOB TORCH_LIBRARY "${VAR_PATH}/*.so")
+    list(APPEND TORCH_LIBRARIES "${TORCH_LIBRARY}")
+endforeach()
+message("-- Detected Torch libraries: \"${TORCH_LIBRARIES}\"")
+
 add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)
 
 add_subdirectory("third_party/pybind11")
 include_directories(include)
 
-foreach(TMP_PATH ${TORCH_LIBRARY_PATH})
-    file(GLOB TORCH_LIBRARY ${TMP_PATH}/*.so)
-    set(TORCH_LIBRARIES "${TORCH_LIBRARIES};${TORCH_LIBRARY};")
-endforeach()
-
 add_subdirectory(src)