Module variable

examples/pure/pure.pyi

@@ -6,6 +6,7 @@ import pathlib
 import typing
 from enum import Enum, auto
 
+MY_CONSTANT: int
 class A:
     x: int
     def __new__(cls,x:int): ...

examples/pure/src/lib.rs

@@ -3,7 +3,7 @@ mod readme {}
 
 use ahash::RandomState;
 use pyo3::{exceptions::PyRuntimeError, prelude::*, types::*};
-use pyo3_stub_gen::{create_exception, define_stub_info_gatherer, derive::*};
+use pyo3_stub_gen::{create_exception, define_stub_info_gatherer, derive::*, module_variable};
 use std::{collections::HashMap, path::PathBuf};
 
 /// Returns the sum of two numbers as a string.
@@ -99,10 +99,13 @@ pub enum Number {
     Integer,
 }
 
+module_variable!("pure", "MY_CONSTANT", usize);
+
 /// Initializes the Python module
 #[pymodule]
 fn pure(m: &Bound<PyModule>) -> PyResult<()> {
     m.add("MyError", m.py().get_type_bound::<MyError>())?;
+    m.add("MY_CONSTANT", 19937)?;
     m.add_class::<A>()?;
     m.add_class::<Number>()?;
     m.add_function(wrap_pyfunction!(sum, m)?)?;

pyo3-stub-gen/src/generate.rs

@@ -10,6 +10,7 @@ mod method;
 mod module;
 mod new;
 mod stub_info;
+mod variable;
 
 pub use arg::*;
 pub use class::*;
@@ -21,6 +22,7 @@ pub use method::*;
 pub use module::*;
 pub use new::*;
 pub use stub_info::*;
+pub use variable::*;
 
 use crate::stub_type::ModuleRef;
 use std::collections::HashSet;

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

@@ -13,6 +13,7 @@ pub struct Module {
     pub enum_: BTreeMap<TypeId, EnumDef>,
     pub function: BTreeMap<&'static str, FunctionDef>,
     pub error: BTreeMap<&'static str, ErrorDef>,
+    pub variables: BTreeMap<&'static str, VariableDef>,
     pub name: String,
     pub default_module_name: String,
     /// Direct submodules of this module.
@@ -51,6 +52,9 @@ impl fmt::Display for Module {
         }
         writeln!(f)?;
 
+        for var in self.variables.values() {
+            writeln!(f, "{}", var)?;
+        }
         for class in self.class.values().sorted_by_key(|class| class.name) {
             write!(f, "{}", class)?;
         }

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

@@ -108,6 +108,12 @@ impl StubInfoBuilder {
             .insert(info.name, ErrorDef::from(info));
     }
 
+    fn add_variable(&mut self, info: &PyVariableInfo) {
+        self.get_module(Some(info.module))
+            .variables
+            .insert(info.name, VariableDef::from(info));
+    }
+
     fn add_methods(&mut self, info: &PyMethodsInfo) {
         let struct_id = (info.struct_id)();
         for module in self.modules.values_mut() {
@@ -143,6 +149,9 @@ impl StubInfoBuilder {
         for info in inventory::iter::<PyErrorInfo> {
             self.add_error(info);
         }
+        for info in inventory::iter::<PyVariableInfo> {
+            self.add_variable(info);
+        }
         for info in inventory::iter::<PyMethodsInfo> {
             self.add_methods(info);
         }

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

@@ -0,0 +1,24 @@
+use std::fmt;
+
+use crate::{type_info::PyVariableInfo, TypeInfo};
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct VariableDef {
+    pub name: &'static str,
+    pub type_: TypeInfo,
+}
+
+impl From<&PyVariableInfo> for VariableDef {
+    fn from(info: &PyVariableInfo) -> Self {
+        Self {
+            name: info.name,
+            type_: (info.r#type)(),
+        }
+    }
+}
+
+impl fmt::Display for VariableDef {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}: {}", self.name, self.type_)
+    }
+}

pyo3-stub-gen/src/lib.rs

@@ -175,3 +175,16 @@ macro_rules! define_stub_info_gatherer {
         }
     };
 }
+
+#[macro_export]
+macro_rules! module_variable {
+    ($module:expr, $name:expr, $ty:ty) => {
+        $crate::inventory::submit! {
+            $crate::type_info::PyVariableInfo{
+                name: $name,
+                module: $module,
+                r#type: <$ty as $crate::PyStubType>::type_output,
+            }
+        }
+    };
+}

pyo3-stub-gen/src/type_info.rs

@@ -137,3 +137,12 @@ pub struct PyErrorInfo {
 }
 
 inventory::collect!(PyErrorInfo);
+
+#[derive(Debug)]
+pub struct PyVariableInfo {
+    pub name: &'static str,
+    pub module: &'static str,
+    pub r#type: fn() -> TypeInfo,
+}
+
+inventory::collect!(PyVariableInfo);