Skip to content

Commit

Permalink
Module variable (#63)
Browse files Browse the repository at this point in the history
Annotation in Rust side:

```rust
module_variable!("pure", "MY_CONSTANT", usize);
```

will creates module variable in `pure.pyi`

```python
MY_CONSTANT: int
```
  • Loading branch information
termoshtt authored Aug 22, 2024
1 parent cea14d1 commit 4721c02
Show file tree
Hide file tree
Showing 8 changed files with 66 additions and 1 deletion.
1 change: 1 addition & 0 deletions examples/pure/pure.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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): ...
Expand Down
5 changes: 4 additions & 1 deletion examples/pure/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)?)?;
Expand Down
2 changes: 2 additions & 0 deletions pyo3-stub-gen/src/generate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ mod method;
mod module;
mod new;
mod stub_info;
mod variable;

pub use arg::*;
pub use class::*;
Expand All @@ -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;
Expand Down
4 changes: 4 additions & 0 deletions pyo3-stub-gen/src/generate/module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)?;
}
Expand Down
9 changes: 9 additions & 0 deletions pyo3-stub-gen/src/generate/stub_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -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);
}
Expand Down
24 changes: 24 additions & 0 deletions pyo3-stub-gen/src/generate/variable.rs
Original file line number Diff line number Diff line change
@@ -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_)
}
}
13 changes: 13 additions & 0 deletions pyo3-stub-gen/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}
};
}
9 changes: 9 additions & 0 deletions pyo3-stub-gen/src/type_info.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

0 comments on commit 4721c02

Please sign in to comment.