fix: Track all submodules instead of one (#106)

Currently , submodule relation is tracked as `map: BTreeMap<String,
String>` in `register_submodules` function:


https://github.com/Jij-Inc/pyo3-stub-gen/blob/bb791fb6d20e6c019aef32437d59a9b6349ff19e/pyo3-stub-gen/src/generate/stub_info.rs#L70-L85

This resulted in wrong behaviour when a module has multiple submodules -
it contains only one `from . import submod` for the last processed
submodule.
This PR fixes this by using `BTreeMap<String, BTreeSet<String>>`
instead. It also adds `examples/mixed_sub_multiple` case to check the
intended behaviour for this case.

.github/workflows/python.yml

@@ -17,6 +17,7 @@ jobs:
           - pure
           - mixed
           - mixed_sub
+          - mixed_sub_multiple
     steps:
       - name: Checkout
         uses: actions/checkout@v4

Cargo.toml

@@ -5,6 +5,7 @@ members = [
     "examples/pure",
     "examples/mixed",
     "examples/mixed_sub",
+    "examples/mixed_sub_multiple",
 ]
 resolver = "2"
 

examples/mixed_sub_multiple/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "mixed_sub_multiple"
+version.workspace = true
+edition.workspace = true
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+env_logger.workspace = true
+pyo3-stub-gen = { path = "../../pyo3-stub-gen" }
+pyo3.workspace = true
+
+[[bin]]
+name = "stub_gen"
+doc = false

examples/mixed_sub_multiple/pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["maturin>=1.1,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "mixed_sub_multiple"
+requires-python = ">=3.9"
+
+[project.optional-dependencies]
+test = ["pytest", "pyright", "ruff"]
+
+[tool.maturin]
+python-source = "python"
+module-name = "mixed_sub_multiple.main_mod"
+features = ["pyo3/extension-module"]

examples/mixed_sub_multiple/python/mixed_sub_multiple/main_mod/__init__.pyi

@@ -0,0 +1,9 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+from . import mod_a
+from . import mod_b
+
+def greet_main() -> None:
+    ...
+

examples/mixed_sub_multiple/python/mixed_sub_multiple/main_mod/mod_a.pyi

@@ -0,0 +1,7 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+
+def greet_a() -> None:
+    ...
+

examples/mixed_sub_multiple/python/mixed_sub_multiple/main_mod/mod_b.pyi

@@ -0,0 +1,7 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: E501, F401
+
+
+def greet_b() -> None:
+    ...
+

examples/mixed_sub_multiple/src/bin/stub_gen.rs

@@ -0,0 +1,8 @@
+use pyo3_stub_gen::Result;
+
+fn main() -> Result<()> {
+    env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
+    let stub = mixed_sub_multiple::stub_info()?;
+    stub.generate()?;
+    Ok(())
+}

examples/mixed_sub_multiple/src/lib.rs

@@ -0,0 +1,57 @@
+use pyo3::prelude::*;
+use pyo3_stub_gen::{define_stub_info_gatherer, derive::*};
+
+#[gen_stub_pyfunction(module = "mixed_sub_multiple.main_mod.mod_a")]
+#[pyfunction(name = "greet_a")]
+pub fn greet_a() {
+    println!("Hello from mod_A!")
+}
+
+#[gen_stub_pyfunction(module = "mixed_sub_multiple.main_mod")]
+#[pyfunction(name = "greet_main")]
+pub fn greet_main() {
+    println!("Hello from main_mod!")
+}
+
+#[gen_stub_pyfunction(module = "mixed_sub_multiple.main_mod.mod_b")]
+#[pyfunction(name = "greet_b")]
+pub fn greet_b() {
+    println!("Hello from mod_B!")
+}
+
+#[pymodule]
+fn main_mod(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(greet_main, m)?)?;
+    mod_a(m)?;
+    mod_b(m)?;
+    Ok(())
+}
+
+#[pymodule]
+fn mod_a(parent: &Bound<PyModule>) -> PyResult<()> {
+    let py = parent.py();
+    let sub = PyModule::new_bound(py, "mod_a")?;
+    sub.add_function(wrap_pyfunction!(greet_a, &sub)?)?;
+    parent.add_submodule(&sub)?;
+    Ok(())
+}
+
+#[pymodule]
+fn mod_b(parent: &Bound<PyModule>) -> PyResult<()> {
+    let py = parent.py();
+    let sub = PyModule::new_bound(py, "mod_b")?;
+    sub.add_function(wrap_pyfunction!(greet_b, &sub)?)?;
+    parent.add_submodule(&sub)?;
+    Ok(())
+}
+
+define_stub_info_gatherer!(stub_info);
+
+/// Test of unit test for testing link problem
+#[cfg(test)]
+mod test {
+    #[test]
+    fn test() {
+        assert_eq!(2 + 2, 4);
+    }
+}

examples/mixed_sub_multiple/tests/test_mixed_sub_multiple.py

@@ -0,0 +1,13 @@
+from mixed_sub_multiple import main_mod
+
+
+def test_main_mod():
+    main_mod.greet_main()
+
+
+def test_sub_mod_a():
+    main_mod.mod_a.greet_a()
+
+
+def test_sub_mod_b():
+    main_mod.mod_b.greet_b()

pyo3-stub-gen/src/generate/stub_info.rs

@@ -1,6 +1,11 @@
 use crate::{generate::*, pyproject::PyProject, type_info::*};
 use anyhow::{Context, Result};
-use std::{collections::BTreeMap, fs, io::Write, path::*};
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    fs,
+    io::Write,
+    path::*,
+};
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct StubInfo {
@@ -68,18 +73,20 @@ impl StubInfoBuilder {
     }
 
     fn register_submodules(&mut self) {
-        let mut map = BTreeMap::new();
+        let mut map: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
         for module in self.modules.keys() {
             let path = module.split('.').collect::<Vec<_>>();
             let n = path.len();
             if n <= 1 {
                 continue;
             }
-            map.insert(path[..n - 1].join("."), path[n - 1].to_string());
+            map.entry(path[..n - 1].join("."))
+                .or_default()
+                .insert(path[n - 1].to_string());
         }
-        for (parent, child) in map {
+        for (parent, children) in map {
             if let Some(module) = self.modules.get_mut(&parent) {
-                module.submodules.insert(child);
+                module.submodules.extend(children);
             }
         }
     }