diff --git a/env/helpers/_pysetup.py b/env/helpers/_pysetup.py index 16e426930..85838052a 100644 --- a/env/helpers/_pysetup.py +++ b/env/helpers/_pysetup.py @@ -84,7 +84,7 @@ def build_extension(self, ext: ShamEnvExtension) -> None: author_email="tim.shamrock@proton.me", description="SHAMROCK Code for astrophysics", long_description="", - ext_modules=[ShamEnvExtension("shamrock.shamrock")], + ext_modules=[ShamEnvExtension("shamrock.pyshamrock")], data_files=[("bin", ["shamrock"])], cmdclass={"build_ext": ShamEnvBuild}, zip_safe=False, diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index e70f3b8a8..128e2228e 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -96,7 +96,7 @@ target_link_lib_and_install(shamrock_exe) target_link_lib_and_install(shamrock_pylib) set_property(TARGET shamrock_exe PROPERTY OUTPUT_NAME shamrock) -set_property(TARGET shamrock_pylib PROPERTY OUTPUT_NAME shamrock) +set_property(TARGET shamrock_pylib PROPERTY OUTPUT_NAME pyshamrock) ######################################################################################## @@ -131,10 +131,9 @@ if(CMAKE_INSTALL_PYTHONDIR) # Set the install dir for the shamrock package set(SHAMROCK_PYTHON_INSTALL_DIR "${CMAKE_INSTALL_PYTHONDIR}/shamrock") - # Create and install __init__.py file for Python package - file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/__init__.py" "from .shamrock import *\n") - install(FILES "${CMAKE_CURRENT_BINARY_DIR}/__init__.py" - DESTINATION "${SHAMROCK_PYTHON_INSTALL_DIR}") + # Copy content of pylib/shamrock to the install directory + install(DIRECTORY "${CMAKE_SOURCE_DIR}/src/pylib/shamrock" + DESTINATION "${CMAKE_INSTALL_PYTHONDIR}") file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/math/__init__.py" "from ..shamrock.math import *\n") install(FILES "${CMAKE_CURRENT_BINARY_DIR}/math/__init__.py" diff --git a/src/main.cpp b/src/main.cpp index 00724bbbd..b5274803e 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -45,7 +45,7 @@ //%Impl status : Should rewrite /// Call bindings init for the shamrock python module -PYBIND11_EMBEDDED_MODULE(shamrock, m) { shambindings::init_embed(m); } +PYBIND11_EMBEDDED_MODULE(pyshamrock, m) { shambindings::init_embed(m); } int main(int argc, char *argv[]) { @@ -190,13 +190,13 @@ int main(int argc, char *argv[]) { "cannot run ipython mode with > 1 processes"); } - shambindings::start_ipython(true); + shambindings::start_ipython(true, argc, argv); } else if (opts::has_option("--rscript")) { __shamrock_stack_entry(); std::string fname = std::string(opts::get_option("--rscript")); - shambindings::run_py_file(fname, shamcomm::world_rank() == 0); + shambindings::run_py_file(fname, shamcomm::world_rank() == 0, argc, argv); } else { if (shamcomm::world_rank() == 0) { diff --git a/src/main_lib.cpp b/src/main_lib.cpp index 9460cc3a2..d7111d255 100644 --- a/src/main_lib.cpp +++ b/src/main_lib.cpp @@ -18,4 +18,4 @@ #include /// Call bindings init for the shamrock python module -PYBIND11_MODULE(shamrock, m) { shambindings::init_lib(m); } +PYBIND11_MODULE(pyshamrock, m) { shambindings::init_lib(m); } diff --git a/src/main_test.cpp b/src/main_test.cpp index 3880f8558..1da1722c6 100644 --- a/src/main_test.cpp +++ b/src/main_test.cpp @@ -40,7 +40,7 @@ #include /// Call bindings init for the shamrock python module -PYBIND11_EMBEDDED_MODULE(shamrock, m) { shambindings::init_embed(m); } +PYBIND11_EMBEDDED_MODULE(pyshamrock, m) { shambindings::init_embed(m); } int main(int argc, char *argv[]) { diff --git a/src/pylib/shamrock/__init__.py b/src/pylib/shamrock/__init__.py new file mode 100644 index 000000000..7d5419050 --- /dev/null +++ b/src/pylib/shamrock/__init__.py @@ -0,0 +1,13 @@ +try: + # try to import from the global namespace (works if embedded python interpreter is used) + from pyshamrock import * + + IMPORT_LOG = "global" +except ImportError: + # then it is a library mode, we import from the local namespace + from .pyshamrock import * + + IMPORT_LOG = "local" + +print(f"pyshamrock imported from {__file__}") +print(f"import log: {IMPORT_LOG}") diff --git a/src/pylib/shamrock/utils/__init__.py b/src/pylib/shamrock/utils/__init__.py new file mode 100644 index 000000000..f3abad5c0 --- /dev/null +++ b/src/pylib/shamrock/utils/__init__.py @@ -0,0 +1,2 @@ +def hello_world(): + print("Hello, World!") diff --git a/src/shambindings/CMakeLists.txt b/src/shambindings/CMakeLists.txt index 33e9707f2..a9cb22fb1 100644 --- a/src/shambindings/CMakeLists.txt +++ b/src/shambindings/CMakeLists.txt @@ -58,6 +58,7 @@ file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/configure_time_py_sys_path.cpp " const char* configure_time_py_sys_path() { return \"${__PYTHON_SYS_PATH_OUT}\"; } const char* configure_time_py_executable() { return \"${PYTHON_EXECUTABLE}\"; } + const char* configure_time_pylib_path() { return \"${CMAKE_SOURCE_DIR}/src/pylib\"; } " ) diff --git a/src/shambindings/include/shambindings/start_python.hpp b/src/shambindings/include/shambindings/start_python.hpp index 7fb9084e4..7178463fa 100644 --- a/src/shambindings/include/shambindings/start_python.hpp +++ b/src/shambindings/include/shambindings/start_python.hpp @@ -27,6 +27,13 @@ namespace shambindings { */ void setpypath(std::string path); + /** + * @brief set the value of sys.argv + * + * This function will throw if bindings were not initialized in embed mode + */ + void set_sys_argv(int argc, char *argv[]); + /** * @brief set the value of sys.path before init from the supplied binary * @@ -42,7 +49,7 @@ namespace shambindings { * @warning This function shall not be called if more than one processes are running * @param do_print print log at python startup */ - void start_ipython(bool do_print); + void start_ipython(bool do_print, int argc, char *argv[]); /** * @brief run python runscript @@ -52,7 +59,7 @@ namespace shambindings { * @param do_print print log at python startup * @param file_path path to the runscript */ - void run_py_file(std::string file_path, bool do_print); + void run_py_file(std::string file_path, bool do_print, int argc, char *argv[]); /** * @brief Modify Python sys.path to point to one detected during cmake invocation diff --git a/src/shambindings/src/start_python.cpp b/src/shambindings/src/start_python.cpp index 3245b6a07..60d6efcd5 100644 --- a/src/shambindings/src/start_python.cpp +++ b/src/shambindings/src/start_python.cpp @@ -14,14 +14,18 @@ * */ +#include "shambase/exception.hpp" #include "shambase/popen.hpp" #include "shambase/print.hpp" #include "shambindings/pybindaliases.hpp" #include "shambindings/pybindings.hpp" #include "shambindings/start_python.hpp" +#include "shamcmdopt/env.hpp" #include #include +#include #include +#include #include /** @@ -34,6 +38,9 @@ extern const char *configure_time_py_sys_path(); /// @brief path of the python executable that was used to configure sys.path extern const char *configure_time_py_executable(); +/// @brief Path to shamrock utils lib a config time +extern const char *configure_time_pylib_path(); + /** * @brief Script to run ipython * @@ -81,8 +88,74 @@ if not cur_path.startswith(sysprefix): )"; +// env var to set the path to the pylib +std::optional pylib_path_env_var = shamcmdopt::getenv_str("SHAMROCK_PYLIB_PATH"); + namespace shambindings { + std::optional get_binary_path() { + + // first try /proc/self/exe + try { + return std::filesystem::read_symlink("/proc/self/exe"); + } catch (const std::filesystem::filesystem_error &e) { + return std::nullopt; + } + + // then try sys.executable from python because why not XD + try { + py::module_ sys = py::module_::import("sys"); + std::string executable = sys.attr("executable").cast(); + return executable; + } catch (const std::exception &e) { + return std::nullopt; + } + } + + std::string locate_pylib_path(bool do_print) { + + auto get_binary_dir = []() -> std::filesystem::path { + auto bpath = get_binary_path(); + if (bpath.has_value()) { + return std::filesystem::path(bpath.value()).parent_path(); + } + return std::filesystem::path("."); + }; + + // Get the path to the current binary + std::filesystem::path binary_dir = get_binary_dir(); + + std::filesystem::path pyshamrock_path_relative1 = binary_dir / ".." / "pylib"; + std::filesystem::path pyshamrock_path_relative2 = binary_dir / ".." / "src" / "pylib"; + + std::vector possible_paths + = {"pyshamrock", + pyshamrock_path_relative1, + pyshamrock_path_relative2, + std::string(configure_time_pylib_path())}; + + if (pylib_path_env_var.has_value()) { + possible_paths.push_back(pylib_path_env_var.value()); + } + + std::string ret = std::string(configure_time_pylib_path()); + + for (const auto &path : possible_paths) { + if (std::filesystem::is_directory(path.c_str())) { + ret = path; + break; + } else { + shambase::println("pylib path " + path + " does not exist, skipping"); + } + } + + if (do_print) { + shambase::println("using pylib path : " + ret); + } + + return ret; + } + void setpypath(std::string path) { runtime_set_pypath = path; } void setpypath_from_binary(std::string binary_path) { @@ -107,12 +180,34 @@ namespace shambindings { std::string modify_path = std::string("paths = ") + get_pypath() + "\n"; modify_path += R"(import sys;sys.path = paths)"; py::exec(modify_path); + + std::string pylib_path = locate_pylib_path(do_print); + std::string modify_path_lib = std::string("sys.path.insert(0, \"") + pylib_path + "\")\n"; + py::exec(modify_path_lib); + } + + void set_sys_argv(int argc, char *argv[]) { + std::vector sys_argv; + for (int i = 0; i < argc; i++) { + sys_argv.push_back(argv[i]); + } + std::stringstream ss; + ss << "["; + for (const auto &arg : sys_argv) { + ss << "\"" << arg << "\", "; + } + ss << "]"; + + std::string cmd = "import sys; sys.argv = " + ss.str(); + + py::exec(cmd); } - void start_ipython(bool do_print) { + void start_ipython(bool do_print, int argc, char *argv[]) { py::scoped_interpreter guard{}; modify_py_sys_path(do_print); + set_sys_argv(argc, argv); if (do_print) { shambase::println("--------------------------------------------"); @@ -127,9 +222,10 @@ namespace shambindings { } } - void run_py_file(std::string file_path, bool do_print) { + void run_py_file(std::string file_path, bool do_print, int argc, char *argv[]) { py::scoped_interpreter guard{}; modify_py_sys_path(do_print); + set_sys_argv(argc, argv); if (do_print) { shambase::println("-----------------------------------"); diff --git a/src/shamtest/shamtest.cpp b/src/shamtest/shamtest.cpp index 65ea57474..5cb813347 100644 --- a/src/shamtest/shamtest.cpp +++ b/src/shamtest/shamtest.cpp @@ -501,6 +501,7 @@ namespace shamtest { ON_RANK_0(shamcomm::logs::print_faint_row()); shambindings::modify_py_sys_path(shamcomm::world_rank() == 0); + shambindings::set_sys_argv(argc, argv); ON_RANK_0(shamcomm::logs::print_faint_row()); // import shamrock in pybind