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.rs

@@ -23,13 +23,13 @@ pub use module::*;
 pub use stub_info::*;
 pub use variable::*;
 
-use crate::stub_type::ModuleRef;
+use crate::stub_type::CommonRef;
 use std::collections::HashSet;
 
 fn indent() -> &'static str {
     "    "
 }
 
 pub trait Import {
-    fn import(&self) -> HashSet<ModuleRef>;
+    fn import(&self) -> HashSet<CommonRef>;
 }

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

@@ -1,4 +1,4 @@
-use crate::{generate::Import, stub_type::ModuleRef, type_info::*, TypeInfo};
+use crate::{generate::Import, stub_type::CommonRef, type_info::*, TypeInfo};
 use std::{collections::HashSet, fmt};
 
 #[derive(Debug, Clone, PartialEq)]
@@ -9,7 +9,7 @@ pub struct Arg {
 }
 
 impl Import for Arg {
-    fn import(&self) -> HashSet<ModuleRef> {
+    fn import(&self) -> HashSet<CommonRef> {
         self.r#type.import.clone()
     }
 }

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

@@ -12,7 +12,7 @@ pub struct ClassDef {
 }
 
 impl Import for ClassDef {
-    fn import(&self) -> HashSet<ModuleRef> {
+    fn import(&self) -> HashSet<CommonRef> {
         let mut import = HashSet::new();
         for base in &self.bases {
             import.extend(base.import.clone());

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

@@ -11,7 +11,7 @@ pub struct FunctionDef {
 }
 
 impl Import for FunctionDef {
-    fn import(&self) -> HashSet<ModuleRef> {
+    fn import(&self) -> HashSet<CommonRef> {
         let mut import = self.r#return.import.clone();
         for arg in &self.args {
             import.extend(arg.import().into_iter());

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

@@ -13,7 +13,7 @@ pub struct MemberDef {
 }
 
 impl Import for MemberDef {
-    fn import(&self) -> HashSet<ModuleRef> {
+    fn import(&self) -> HashSet<CommonRef> {
         self.r#type.import.clone()
     }
 }

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

@@ -14,7 +14,7 @@ pub struct MethodDef {
 }
 
 impl Import for MethodDef {
-    fn import(&self) -> HashSet<ModuleRef> {
+    fn import(&self) -> HashSet<CommonRef> {
         let mut import = self.r#return.import.clone();
         for arg in &self.args {
             import.extend(arg.import().into_iter());

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

@@ -21,7 +21,7 @@ pub struct Module {
 }
 
 impl Import for Module {
-    fn import(&self) -> HashSet<ModuleRef> {
+    fn import(&self) -> HashSet<CommonRef> {
         let mut imports = HashSet::new();
         for class in self.class.values() {
             imports.extend(class.import());
@@ -38,12 +38,36 @@ 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 type_ref_grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
+        for common_ref in self.import().into_iter().sorted() {
+            match common_ref {
+                CommonRef::Module(module_ref) => {
+                    let name = module_ref.get().unwrap_or(&self.default_module_name);
+                    if name != self.name {
+                        writeln!(f, "import {}", name)?;
+                    }
+                }
+                CommonRef::Type(mut type_ref) => {
+                    if type_ref.module.is_empty() {
+                        type_ref.module = self.default_module_name.clone();
+                    }
+                    if type_ref.module != self.name {
+                        if type_ref.module.starts_with(package_name) {
+                            type_ref_grouped
+                                .entry(type_ref.module)
+                                .or_default()
+                                .push(type_ref.name);
+                        } else {
+                            writeln!(f, "import {}", type_ref.module)?;
+                        }
+                    }
+                }
             }
         }
+        for (module_name, type_names) in type_ref_grouped {
+            writeln!(f, "from {} import {}", module_name, type_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,6 +7,37 @@ mod numpy;
 
 use maplit::hashset;
 use std::{collections::HashSet, fmt, ops};
+use std::cmp::Ordering;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum CommonRef {
+    Module(ModuleRef),
+    Type(TypeRef),
+}
+
+impl PartialOrd for CommonRef {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+
+impl Ord for CommonRef {
+    fn cmp(&self, other: &Self) -> Ordering {
+        match (self, other) {
+            (CommonRef::Module(a), CommonRef::Module(b)) => a.get().cmp(&b.get()),
+            (CommonRef::Type(a), CommonRef::Type(b)) => a.cmp(b),
+            (CommonRef::Module(_), CommonRef::Type(_)) => Ordering::Greater,
+            (CommonRef::Type(_), CommonRef::Module(_)) => Ordering::Less,
+        }
+    }
+}
+
+impl From<&str> for CommonRef {
+    fn from(s: &str) -> Self {
+        CommonRef::Module(s.into())
+    }
+}
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
 pub enum ModuleRef {
@@ -39,6 +70,22 @@ impl From<&str> for ModuleRef {
     }
 }
 
+
+/// Indicates the 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  TypeRef {
+    pub module: String,
+    pub name: String,
+}
+
+impl  TypeRef  {
+    pub fn new(module: String, name: String) -> Self {
+        Self{name, module}
+    }
+}
+
 /// Type information for creating Python stub files annotated by [PyStubType] trait.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TypeInfo {
@@ -49,7 +96,7 @@ pub struct TypeInfo {
     ///
     /// For example, when `name` is `typing.Sequence[int]`, `import` should contain `typing`.
     /// This makes it possible to use user-defined types in the stub file.
-    pub import: HashSet<ModuleRef>,
+    pub import: HashSet<CommonRef>,
 }
 
 impl fmt::Display for TypeInfo {
@@ -138,12 +185,35 @@ impl TypeInfo {
     /// ```
     pub fn with_module(name: &str, module: ModuleRef) -> Self {
         let mut import = HashSet::new();
-        import.insert(module);
+        import.insert(CommonRef::Module(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 module_name = String::new();
+        match module.get() {
+            Some(value) => module_name = value.to_string(),
+            None => module_name = "".to_string(),
+        }
+        
+        let type_ref = TypeRef::new(module_name, type_name.to_string());
+        import.insert(CommonRef::Type(type_ref));
+
+        Self {
+            name: type_name.to_string(),
+            import,
+        }
+    }
 }
 
 impl ops::BitOr for TypeInfo {
@@ -244,7 +314,7 @@ mod test {
     #[test_case(HashMap::<u32, Vec<u32>>::type_output(), "builtins.dict[builtins.int, builtins.list[builtins.int]]", hashset! { "builtins".into() } ; "HashMap_u32_Vec_u32_output")]
     #[test_case(HashSet::<u32>::type_input(), "builtins.set[builtins.int]", hashset! { "builtins".into() } ; "HashSet_u32_input")]
     #[test_case(indexmap::IndexSet::<u32>::type_input(), "builtins.set[builtins.int]", hashset! { "builtins".into() } ; "IndexSet_u32_input")]
-    fn test(tinfo: TypeInfo, name: &str, import: HashSet<ModuleRef>) {
+    fn test(tinfo: TypeInfo, name: &str, import: HashSet<CommonRef>) {
         assert_eq!(tinfo.name, name);
         if import.is_empty() {
             assert!(tinfo.import.is_empty());