Update stub generation

HaarigerHarald · Oct 9, 2023 · a874efc · a874efc
1 parent 3dac008
commit a874efc
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 28 deletions.
diff --git a/doc/generate_stubs.py b/doc/generate_stubs.py
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,6 +4,6 @@ requires = [
     "wheel",
     "cmake>=3.16",
     "ninja; sys_platform != 'win32'",
-    "pybind11-stubgen~=0.16.2",
+    "pybind11-stubgen~=2.3",
 ]
 build-backend = "setuptools.build_meta"
diff --git a/setup.py b/setup.py
@@ -75,11 +75,15 @@ def build_extension(self, ext):
             ["cmake", "--build", "."] + build_args, cwd=self.build_temp
         )
 
+        ext_dir = os.path.dirname(self.get_ext_fullpath(ext.name))
+        python_path = ext_dir
+        if "PYTHONPATH" in os.environ:
+            python_path += os.pathsep + os.environ["PYTHONPATH"]
+
         # Automatically generate stubs
         subprocess.check_call(
-            [sys.executable, "doc/generate_stubs.py", ext.name,
-                os.path.realpath(self.get_ext_fullpath(ext.name))],
-            cwd=os.path.dirname(os.path.abspath(__file__))
+            ["pybind11-stubgen", "-o", ext_dir, "--ignore-all-errors", "geant4_pybind"],
+            env=dict(os.environ, PYTHONPATH=python_path)
         )
 
 

diff --git a/source/geant4_pybind.cc b/source/geant4_pybind.cc
@@ -50,6 +50,15 @@ class G4NullCoutDestination : public G4coutDestination {
    G4int ReceiveG4cerr(const G4String &cerrString) override { return 0; }
 };
 
+static bool endsWith(const std::string &str, const std::string &end)
+{
+   if (str.length() >= end.length()) {
+      return str.compare(str.length() - end.length(), end.length(), end) == 0;
+   } else {
+      return false;
+   }
+}
+
 void export_modG4digit_hits(py::module &);
 void export_modG4event(py::module &);
 void export_modG4geometry(py::module &);
@@ -118,8 +127,9 @@ PYBIND11_MODULE(geant4_pybind, m)
       G4SurfaceProperty::CleanSurfacePropertyTable();
    }));
 
-   py::dict globals = py::module_::import("__main__").attr("__dict__");
-   if (!globals.contains("AUTO_STUB_GENERATION")) {
+   py::dict modules = py::module_::import("sys").attr("modules");
+   if (!modules["__main__"].attr("__dir__")().contains("__file__") ||
+       !endsWith(modules["__main__"].attr("__file__").cast<std::string>(), "pybind11-stubgen")) {
 
       py::dict frameGlobals = py::globals();
       py::exec(