Skip to content

Commit f00808d

Browse files
yogevm15yogevm15
authored
Introspection: modules associated constants (#5150)
* Adds introspection for modules associated constants * Add support for `Final` constants. * Document changes * Review fixes * Fix typo
1 parent 3b40ba5 commit f00808d

File tree

11 files changed

+141
-34
lines changed

11 files changed

+141
-34
lines changed

‎newsfragments/5150.added.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
- Add support for module associated consts introspection.

‎pyo3-introspection/src/introspection.rs‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::model::{Argument, Arguments, Class, Function, Module, VariableLengthArgument};
1+
use crate::model::{Argument, Arguments, Class, Const, Function, Module, VariableLengthArgument};
22
use anyhow::{bail, ensure, Context, Result};
33
use goblin::elf::Elf;
44
use goblin::mach::load_command::CommandVariant;
@@ -44,11 +44,12 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
4444
if let Chunk::Module {
4545
name,
4646
members,
47+
consts,
4748
id: _,
4849
} = chunk
4950
{
5051
if name == main_module_name {
51-
return convert_module(name, members, &chunks_by_id, &chunks_by_parent);
52+
return convert_module(name, members, consts, &chunks_by_id, &chunks_by_parent);
5253
}
5354
}
5455
}
@@ -58,6 +59,7 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
5859
fn convert_module(
5960
name: &str,
6061
members: &[String],
62+
consts: &[ConstChunk],
6163
chunks_by_id: &HashMap<&str, &Chunk>,
6264
chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
6365
) -> Result<Module> {
@@ -69,11 +71,19 @@ fn convert_module(
6971
chunks_by_id,
7072
chunks_by_parent,
7173
)?;
74+
7275
Ok(Module {
7376
name: name.into(),
7477
modules,
7578
classes,
7679
functions,
80+
consts: consts
81+
.iter()
82+
.map(|c| Const {
83+
name: c.name.clone(),
84+
value: c.value.clone(),
85+
})
86+
.collect(),
7787
})
7888
}
7989

@@ -91,11 +101,13 @@ fn convert_members(
91101
Chunk::Module {
92102
name,
93103
members,
104+
consts,
94105
id: _,
95106
} => {
96107
modules.push(convert_module(
97108
name,
98109
members,
110+
consts,
99111
chunks_by_id,
100112
chunks_by_parent,
101113
)?);
@@ -354,6 +366,7 @@ enum Chunk {
354366
id: String,
355367
name: String,
356368
members: Vec<String>,
369+
consts: Vec<ConstChunk>,
357370
},
358371
Class {
359372
id: String,
@@ -371,6 +384,12 @@ enum Chunk {
371384
},
372385
}
373386

387+
#[derive(Deserialize)]
388+
struct ConstChunk {
389+
name: String,
390+
value: String,
391+
}
392+
374393
#[derive(Deserialize)]
375394
struct ChunkArguments {
376395
#[serde(default)]

‎pyo3-introspection/src/model.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ pub struct Module {
44
pub modules: Vec<Module>,
55
pub classes: Vec<Class>,
66
pub functions: Vec<Function>,
7+
pub consts: Vec<Const>,
78
}
89

910
#[derive(Debug, Eq, PartialEq, Clone, Hash)]
@@ -20,6 +21,12 @@ pub struct Function {
2021
pub arguments: Arguments,
2122
}
2223

24+
#[derive(Debug, Eq, PartialEq, Clone, Hash)]
25+
pub struct Const {
26+
pub name: String,
27+
pub value: String,
28+
}
29+
2330
#[derive(Debug, Eq, PartialEq, Clone, Hash)]
2431
pub struct Arguments {
2532
/// Arguments before /

‎pyo3-introspection/src/stubs.rs‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
use crate::model::{Argument, Class, Function, Module, VariableLengthArgument};
2-
use std::collections::HashMap;
1+
use crate::model::{Argument, Class, Const, Function, Module, VariableLengthArgument};
2+
use std::collections::{BTreeSet, HashMap};
33
use std::path::{Path, PathBuf};
44

55
/// Generates the [type stubs](https://typing.readthedocs.io/en/latest/source/stubs.html) of a given module.
@@ -32,16 +32,29 @@ fn add_module_stub_files(
3232

3333
/// Generates the module stubs to a String, not including submodules
3434
fn module_stubs(module: &Module) -> String {
35+
let mut modules_to_import = BTreeSet::new();
3536
let mut elements = Vec::new();
3637
for class in &module.classes {
3738
elements.push(class_stubs(class));
3839
}
3940
for function in &module.functions {
4041
elements.push(function_stubs(function));
4142
}
43+
for konst in &module.consts {
44+
elements.push(const_stubs(konst, &mut modules_to_import));
45+
}
4246

43-
// We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
4447
let mut output = String::new();
48+
49+
for module_to_import in &modules_to_import {
50+
output.push_str(&format!("import {module_to_import}\n"));
51+
}
52+
53+
if !modules_to_import.is_empty() {
54+
output.push('\n')
55+
}
56+
57+
// We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
4558
for element in elements {
4659
let is_multiline = element.contains('\n');
4760
if is_multiline && !output.is_empty() && !output.ends_with("\n\n") {
@@ -53,6 +66,7 @@ fn module_stubs(module: &Module) -> String {
5366
output.push('\n');
5467
}
5568
}
69+
5670
// We remove a line jump at the end if they are two
5771
if output.ends_with("\n\n") {
5872
output.pop();
@@ -111,6 +125,12 @@ fn function_stubs(function: &Function) -> String {
111125
buffer
112126
}
113127

128+
fn const_stubs(konst: &Const, modules_to_import: &mut BTreeSet<String>) -> String {
129+
modules_to_import.insert("typing".to_string());
130+
let Const { name, value } = konst;
131+
format!("{name}: typing.Final = {value}")
132+
}
133+
114134
fn argument_stub(argument: &Argument) -> String {
115135
let mut output = argument.name.clone();
116136
if let Some(default_value) = &argument.default_value {

‎pyo3-macros-backend/src/introspection.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use std::collections::HashMap;
1919
use std::hash::{Hash, Hasher};
2020
use std::mem::take;
2121
use std::sync::atomic::{AtomicUsize, Ordering};
22+
use syn::ext::IdentExt;
2223
use syn::{Attribute, Ident, Type, TypePath};
2324

2425
static GLOBAL_COUNTER_FOR_UNIQUE_NAMES: AtomicUsize = AtomicUsize::new(0);
@@ -28,6 +29,9 @@ pub fn module_introspection_code<'a>(
2829
name: &str,
2930
members: impl IntoIterator<Item = &'a Ident>,
3031
members_cfg_attrs: impl IntoIterator<Item = &'a Vec<Attribute>>,
32+
consts: impl IntoIterator<Item = &'a Ident>,
33+
consts_values: impl IntoIterator<Item = &'a String>,
34+
consts_cfg_attrs: impl IntoIterator<Item = &'a Vec<Attribute>>,
3135
) -> TokenStream {
3236
IntrospectionNode::Map(
3337
[
@@ -52,6 +56,23 @@ pub fn module_introspection_code<'a>(
5256
.collect(),
5357
),
5458
),
59+
(
60+
"consts",
61+
IntrospectionNode::List(
62+
consts
63+
.into_iter()
64+
.zip(consts_values)
65+
.zip(consts_cfg_attrs)
66+
.filter_map(|((ident, value), attributes)| {
67+
if attributes.is_empty() {
68+
Some(const_introspection_code(ident, value))
69+
} else {
70+
None // TODO: properly interpret cfg attributes
71+
}
72+
})
73+
.collect(),
74+
),
75+
),
5576
]
5677
.into(),
5778
)
@@ -116,6 +137,20 @@ pub fn function_introspection_code(
116137
IntrospectionNode::Map(desc).emit(pyo3_crate_path)
117138
}
118139

140+
fn const_introspection_code<'a>(ident: &'a Ident, value: &'a String) -> IntrospectionNode<'a> {
141+
IntrospectionNode::Map(
142+
[
143+
("type", IntrospectionNode::String("const".into())),
144+
(
145+
"name",
146+
IntrospectionNode::String(ident.unraw().to_string().into()),
147+
),
148+
("value", IntrospectionNode::String(value.into())),
149+
]
150+
.into(),
151+
)
152+
}
153+
119154
fn arguments_introspection_data<'a>(
120155
signature: &'a FunctionSignature<'a>,
121156
first_argument: Option<&'a str>,

‎pyo3-macros-backend/src/method.rs‎

Lines changed: 2 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use quote::{format_ident, quote, quote_spanned, ToTokens};
77
use syn::{ext::IdentExt, spanned::Spanned, Ident, Result};
88

99
use crate::pyversions::is_abi3_before;
10-
use crate::utils::{Ctx, LitCStr};
10+
use crate::utils::{expr_to_python, Ctx, LitCStr};
1111
use crate::{
1212
attributes::{FromPyWithAttribute, TextSignatureAttribute, TextSignatureAttributeValue},
1313
params::{impl_arg_params, Holders},
@@ -34,31 +34,7 @@ impl RegularArg<'_> {
3434
..
3535
} = self
3636
{
37-
match arg_default {
38-
// literal values
39-
syn::Expr::Lit(syn::ExprLit { lit, .. }) => match lit {
40-
syn::Lit::Str(s) => s.token().to_string(),
41-
syn::Lit::Char(c) => c.token().to_string(),
42-
syn::Lit::Int(i) => i.base10_digits().to_string(),
43-
syn::Lit::Float(f) => f.base10_digits().to_string(),
44-
syn::Lit::Bool(b) => {
45-
if b.value() {
46-
"True".to_string()
47-
} else {
48-
"False".to_string()
49-
}
50-
}
51-
_ => "...".to_string(),
52-
},
53-
// None
54-
syn::Expr::Path(syn::ExprPath { qself, path, .. })
55-
if qself.is_none() && path.is_ident("None") =>
56-
{
57-
"None".to_string()
58-
}
59-
// others, unsupported yet so defaults to `...`
60-
_ => "...".to_string(),
61-
}
37+
expr_to_python(arg_default)
6238
} else if let RegularArg {
6339
option_wrapped_type: Some(..),
6440
..

‎pyo3-macros-backend/src/module.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
33
#[cfg(feature = "experimental-inspect")]
44
use crate::introspection::{introspection_id_const, module_introspection_code};
5+
use crate::utils::expr_to_python;
56
use crate::{
67
attributes::{
78
self, kw, take_attributes, take_pyo3_options, CrateAttribute, GILUsedAttribute,
@@ -150,6 +151,7 @@ pub fn pymodule_module_impl(
150151

151152
let mut pymodule_init = None;
152153
let mut module_consts = Vec::new();
154+
let mut module_consts_values = Vec::new();
153155
let mut module_consts_cfg_attrs = Vec::new();
154156

155157
for item in &mut *items {
@@ -293,8 +295,8 @@ pub fn pymodule_module_impl(
293295
if !find_and_remove_attribute(&mut item.attrs, "pymodule_export") {
294296
continue;
295297
}
296-
297298
module_consts.push(item.ident.clone());
299+
module_consts_values.push(expr_to_python(&item.expr));
298300
module_consts_cfg_attrs.push(get_cfg_attributes(&item.attrs));
299301
}
300302
Item::Static(item) => {
@@ -349,6 +351,9 @@ pub fn pymodule_module_impl(
349351
&name.to_string(),
350352
&module_items,
351353
&module_items_cfg_attrs,
354+
&module_consts,
355+
&module_consts_values,
356+
&module_consts_cfg_attrs,
352357
);
353358
#[cfg(not(feature = "experimental-inspect"))]
354359
let introspection = quote! {};
@@ -432,7 +437,8 @@ pub fn pymodule_function_impl(
432437
);
433438

434439
#[cfg(feature = "experimental-inspect")]
435-
let introspection = module_introspection_code(pyo3_path, &name.to_string(), &[], &[]);
440+
let introspection =
441+
module_introspection_code(pyo3_path, &name.to_string(), &[], &[], &[], &[], &[]);
436442
#[cfg(not(feature = "experimental-inspect"))]
437443
let introspection = quote! {};
438444
#[cfg(feature = "experimental-inspect")]

‎pyo3-macros-backend/src/utils.rs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,3 +393,31 @@ impl TypeExt for syn::Type {
393393
self
394394
}
395395
}
396+
397+
pub fn expr_to_python(expr: &syn::Expr) -> String {
398+
match expr {
399+
// literal values
400+
syn::Expr::Lit(syn::ExprLit { lit, .. }) => match lit {
401+
syn::Lit::Str(s) => s.token().to_string(),
402+
syn::Lit::Char(c) => c.token().to_string(),
403+
syn::Lit::Int(i) => i.base10_digits().to_string(),
404+
syn::Lit::Float(f) => f.base10_digits().to_string(),
405+
syn::Lit::Bool(b) => {
406+
if b.value() {
407+
"True".to_string()
408+
} else {
409+
"False".to_string()
410+
}
411+
}
412+
_ => "...".to_string(),
413+
},
414+
// None
415+
syn::Expr::Path(syn::ExprPath { qself, path, .. })
416+
if qself.is_none() && path.is_ident("None") =>
417+
{
418+
"None".to_string()
419+
}
420+
// others, unsupported yet so defaults to `...`
421+
_ => "...".to_string(),
422+
}
423+
}

‎pytests/src/consts.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
use pyo3::pymodule;
2+
3+
#[pymodule]
4+
pub mod consts {
5+
#[pymodule_export]
6+
pub const PI: f64 = std::f64::consts::PI; // Exports PI constant as part of the module
7+
8+
#[pymodule_export]
9+
pub const SIMPLE: &str = "SIMPLE";
10+
}

‎pytests/src/lib.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use pyo3::wrap_pymodule;
55
pub mod awaitable;
66
pub mod buf_and_str;
77
pub mod comparisons;
8+
mod consts;
89
pub mod datetime;
910
pub mod dict_iter;
1011
pub mod enums;
@@ -22,7 +23,7 @@ mod pyo3_pytests {
2223
use super::*;
2324

2425
#[pymodule_export]
25-
use {pyclasses::pyclasses, pyfunctions::pyfunctions};
26+
use {consts::consts, pyclasses::pyclasses, pyfunctions::pyfunctions};
2627

2728
// Inserting to sys.modules allows importing submodules nicely from Python
2829
// e.g. import pyo3_pytests.buf_and_str as bas

0 commit comments

Comments
 (0)