[Draft] 2nd attempt at registering a external python module as shamro…

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,

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"

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) {

src/main_lib.cpp

@@ -18,4 +18,4 @@
 #include <pybind11/pybind11.h>
 
 /// Call bindings init for the shamrock python module
-PYBIND11_MODULE(shamrock, m) { shambindings::init_lib(m); }
+PYBIND11_MODULE(pyshamrock, m) { shambindings::init_lib(m); }

src/main_test.cpp

@@ -40,7 +40,7 @@
 #include <string>
 
 /// 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[]) {
 

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}")

src/pylib/shamrock/utils/__init__.py

@@ -0,0 +1,2 @@
+def hello_world():
+    print("Hello, World!")

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\"; }
   "
 )
 

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

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 <pybind11/embed.h>
 #include <cstdlib>
+#include <filesystem>
 #include <optional>
+#include <sstream>
 #include <string>
 
 /**
@@ -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<std::string> pylib_path_env_var = shamcmdopt::getenv_str("SHAMROCK_PYLIB_PATH");
+
 namespace shambindings {
 
+    std::optional<std::string> 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<std::string>();
+            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<std::string> 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<std::string> 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("-----------------------------------");

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