Fixes #2, auto generate python files and fix type ref

Author: Scyrc

pyo3-stub-gen-derive/src/gen_stub/stub_type.rs

@@ -21,7 +21,7 @@ impl ToTokens for StubType {
             #[automatically_derived]
             impl ::pyo3_stub_gen::PyStubType for #ty {
                 fn type_output() -> ::pyo3_stub_gen::TypeInfo {
-                    ::pyo3_stub_gen::TypeInfo::with_module(#name, #module_tt)
+                    ::pyo3_stub_gen::TypeInfo::with_type(#name, #module_tt)
                 }
             }
         })

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

@@ -38,11 +38,26 @@ impl fmt::Display for Module {
         writeln!(f, "# This file is automatically generated by pyo3_stub_gen")?;
         writeln!(f, "# ruff: noqa: E501, F401")?;
         writeln!(f)?;
-        for import in self.import().into_iter().sorted() {
-            let name = import.get().unwrap_or(&self.default_module_name);
-            if name != self.name {
-                writeln!(f, "import {}", name)?;
+        let package_name = self.default_module_name.split('.').next().unwrap();
+        let mut grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
+        for mut module_ref in self.import().into_iter().sorted() {
+            if module_ref.module.is_empty() {
+                module_ref.module = self.default_module_name.clone();
             }
+            if module_ref.module != self.name {
+                // writeln!(f, "from {} import {}", module_ref.module, module_ref.name)?;
+                if module_ref.module.starts_with(package_name) {
+                    grouped
+                        .entry(module_ref.module)
+                        .or_default()
+                        .push(module_ref.name);
+                } else {
+                    writeln!(f, "import {}", module_ref.module)?;
+                }
+            }
+        }
+        for (module, names) in grouped {
+            writeln!(f, "from {} import {}", module, names.join(", "))?;
         }
         for submod in &self.submodules {
             writeln!(f, "from . import {}", submod)?;

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

@@ -6,11 +6,32 @@ use std::{
     io::Write,
     path::*,
 };
-
+fn get_destination_paths(module: &Module, python_root: &Path, default_module: &str) -> (PathBuf, Option<PathBuf>) {
+    let path = module.name.replace(".", "/");
+    if module.submodules.is_empty() {
+        if module.name != default_module {
+            (
+                python_root.join(format!("{path}.pyi")),
+                Some(python_root.join(format!("{path}.py"))),
+            )
+        } else {
+            (
+                python_root.join(format!("{path}.pyi")),
+                None,
+            )
+        }
+    } else {
+        (
+            python_root.join(path).join("__init__.pyi"),
+            None,
+        )
+    }
+}
 #[derive(Debug, Clone, PartialEq)]
 pub struct StubInfo {
     pub modules: BTreeMap<String, Module>,
     pub python_root: PathBuf,
+    pub default_module: String
 }
 
 impl StubInfo {
@@ -30,13 +51,7 @@ impl StubInfo {
 
     pub fn generate(&self) -> Result<()> {
         for (name, module) in self.modules.iter() {
-            let path = name.replace(".", "/");
-            let dest = if module.submodules.is_empty() {
-                self.python_root.join(format!("{path}.pyi"))
-            } else {
-                self.python_root.join(path).join("__init__.pyi")
-            };
-
+            let (dest, dest_py) = get_destination_paths(module, &self.python_root, &self.default_module);
             let dir = dest.parent().context("Cannot get parent directory")?;
             if !dir.exists() {
                 fs::create_dir_all(dir)?;
@@ -48,6 +63,15 @@ impl StubInfo {
                 "Generate stub file of a module `{name}` at {dest}",
                 dest = dest.display()
             );
+
+            if let Some(dest_py) = dest_py {
+                let mut f_py = fs::File::create(&dest_py)?;
+                write!(f_py, "# Fixed pylance reportMissingModuleSource warning when using \n# \"from {} import xxx \" and \"import {}\"", module.name, module.name)?;
+                log::info!(
+                    "Generate python file of a (not main)module `{name}` at {dest}",
+                    dest = dest_py.display()
+                );
+            }
         }
         Ok(())
     }
@@ -189,6 +213,7 @@ impl StubInfoBuilder {
         StubInfo {
             modules: self.modules,
             python_root: self.python_root,
+            default_module: self.default_module_name
         }
     }
 }

pyo3-stub-gen/src/stub_type.rs

@@ -7,35 +7,27 @@ mod numpy;
 
 use maplit::hashset;
 use std::{collections::HashSet, fmt, ops};
-
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
-pub enum ModuleRef {
-    Named(String),
-
-    /// Default module that PyO3 creates.
-    ///
-    /// - For pure Rust project, the default module name is the crate name specified in `Cargo.toml`
-    ///   or `project.name` specified in `pyproject.toml`
-    /// - For mixed Rust/Python project, the default module name is `tool.maturin.module-name` specified in `pyproject.toml`
-    ///
-    /// Because the default module name cannot be known at compile time, it will be resolved at the time of the stub file generation.
-    /// This is a placeholder for the default module name.
-    #[default]
-    Default,
+/**
+ * Indicates the dependent module or type;
+ * dependent module('import module', eg builtins): name="" , module = module name;
+ * dependent type(eg class enum)('from module import type'): name=type name , module = module name(which type defined).
+ */
+ #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
+pub struct  ModuleRef {
+    pub name: String,
+    pub module: String,
 }
-
-impl ModuleRef {
-    pub fn get(&self) -> Option<&str> {
-        match self {
-            Self::Named(name) => Some(name),
-            Self::Default => None,
-        }
+/// From the dependent module
+impl From<&str> for ModuleRef {
+    fn from(m: &str) -> Self {
+        Self{name:"".to_string(), module:m.to_string()}
     }
 }
 
-impl From<&str> for ModuleRef {
-    fn from(s: &str) -> Self {
-        Self::Named(s.to_string())
+impl fmt::Display for ModuleRef {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        writeln!(f, "name: {}, module: {}", self.name, self.module)?;
+        Ok(())
     }
 }
 
@@ -136,14 +128,33 @@ impl TypeInfo {
     /// ```
     /// pyo3_stub_gen::TypeInfo::with_module("pathlib.Path", "pathlib".into());
     /// ```
-    pub fn with_module(name: &str, module: ModuleRef) -> Self {
+    pub fn with_module(name: &str, module: ModuleRef) -> Self
+    {
         let mut import = HashSet::new();
         import.insert(module);
         Self {
             name: name.to_string(),
             import,
         }
     }
+
+    /// A type annotation of a type that must be imported.
+    ///
+    /// ```
+    /// ClassA defined in ModuleA
+    /// pyo3_stub_gen::TypeInfo::with_type("ClassA", "ModuleA");
+    /// ```
+    pub fn with_type(type_name: &str, module: ModuleRef) -> Self {
+        let mut import = HashSet::new();
+        let mut type_ref: ModuleRef = module.clone();
+        type_ref.name = type_name.to_string();
+        import.insert(type_ref);
+
+        Self {
+            name: type_name.to_string(),
+            import,
+        }
+    }
 }
 
 impl ops::BitOr for TypeInfo {