Introspection: modules associated constants

newsfragments/5150.added.md

@@ -0,0 +1 @@
+- Add support for module associated consts introspection.

pyo3-introspection/src/introspection.rs

@@ -1,4 +1,4 @@
-use crate::model::{Argument, Arguments, Class, Function, Module, VariableLengthArgument};
+use crate::model::{Argument, Arguments, Class, Const, Function, Module, VariableLengthArgument};
 use anyhow::{bail, ensure, Context, Result};
 use goblin::elf::Elf;
 use goblin::mach::load_command::CommandVariant;
@@ -44,11 +44,12 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
         if let Chunk::Module {
             name,
             members,
+            consts,
             id: _,
         } = chunk
         {
             if name == main_module_name {
-                return convert_module(name, members, &chunks_by_id, &chunks_by_parent);
+                return convert_module(name, members, consts, &chunks_by_id, &chunks_by_parent);
             }
         }
     }
@@ -58,6 +59,7 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
 fn convert_module(
     name: &str,
     members: &[String],
+    consts: &[ConstChunk],
     chunks_by_id: &HashMap<&str, &Chunk>,
     chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
 ) -> Result<Module> {
@@ -69,11 +71,19 @@ fn convert_module(
         chunks_by_id,
         chunks_by_parent,
     )?;
+
     Ok(Module {
         name: name.into(),
         modules,
         classes,
         functions,
+        consts: consts
+            .iter()
+            .map(|c| Const {
+                name: c.name.clone(),
+                value: c.value.clone(),
+            })
+            .collect(),
     })
 }
 
@@ -91,11 +101,13 @@ fn convert_members(
             Chunk::Module {
                 name,
                 members,
+                consts,
                 id: _,
             } => {
                 modules.push(convert_module(
                     name,
                     members,
+                    consts,
                     chunks_by_id,
                     chunks_by_parent,
                 )?);
@@ -354,6 +366,7 @@ enum Chunk {
         id: String,
         name: String,
         members: Vec<String>,
+        consts: Vec<ConstChunk>,
     },
     Class {
         id: String,
@@ -371,6 +384,12 @@ enum Chunk {
     },
 }
 
+#[derive(Deserialize)]
+struct ConstChunk {
+    name: String,
+    value: String,
+}
+
 #[derive(Deserialize)]
 struct ChunkArguments {
     #[serde(default)]

pyo3-introspection/src/model.rs

@@ -4,6 +4,7 @@ pub struct Module {
     pub modules: Vec<Module>,
     pub classes: Vec<Class>,
     pub functions: Vec<Function>,
+    pub consts: Vec<Const>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
@@ -20,6 +21,12 @@ pub struct Function {
     pub arguments: Arguments,
 }
 
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct Const {
+    pub name: String,
+    pub value: String,
+}
+
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Arguments {
     /// Arguments before /

pyo3-introspection/src/stubs.rs

@@ -1,5 +1,5 @@
-use crate::model::{Argument, Class, Function, Module, VariableLengthArgument};
-use std::collections::HashMap;
+use crate::model::{Argument, Class, Const, Function, Module, VariableLengthArgument};
+use std::collections::{BTreeSet, HashMap};
 use std::path::{Path, PathBuf};
 
 /// 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(
 
 /// Generates the module stubs to a String, not including submodules
 fn module_stubs(module: &Module) -> String {
+    let mut modules_to_import = BTreeSet::new();
     let mut elements = Vec::new();
     for class in &module.classes {
         elements.push(class_stubs(class));
     }
     for function in &module.functions {
         elements.push(function_stubs(function));
     }
+    for konst in &module.consts {
+        elements.push(const_stubs(konst, &mut modules_to_import));
+    }
 
-    // We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
     let mut output = String::new();
+
+    for module_to_import in &modules_to_import {
+        output.push_str(&format!("import {module_to_import}\n"));
+    }
+
+    if !modules_to_import.is_empty() {
+        output.push('\n')
+    }
+
+    // We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
     for element in elements {
         let is_multiline = element.contains('\n');
         if is_multiline && !output.is_empty() && !output.ends_with("\n\n") {
@@ -53,6 +66,7 @@ fn module_stubs(module: &Module) -> String {
             output.push('\n');
         }
     }
+
     // We remove a line jump at the end if they are two
     if output.ends_with("\n\n") {
         output.pop();
@@ -111,6 +125,12 @@ fn function_stubs(function: &Function) -> String {
     buffer
 }
 
+fn const_stubs(konst: &Const, modules_to_import: &mut BTreeSet<String>) -> String {
+    modules_to_import.insert("typing".to_string());
+    let Const { name, value } = konst;
+    format!("{name}: typing.Final = {value}")
+}
+
 fn argument_stub(argument: &Argument) -> String {
     let mut output = argument.name.clone();
     if let Some(default_value) = &argument.default_value {

pyo3-macros-backend/src/introspection.rs

@@ -19,6 +19,7 @@ use std::collections::HashMap;
 use std::hash::{Hash, Hasher};
 use std::mem::take;
 use std::sync::atomic::{AtomicUsize, Ordering};
+use syn::ext::IdentExt;
 use syn::{Attribute, Ident, Type, TypePath};
 
 static GLOBAL_COUNTER_FOR_UNIQUE_NAMES: AtomicUsize = AtomicUsize::new(0);
@@ -28,6 +29,9 @@ pub fn module_introspection_code<'a>(
     name: &str,
     members: impl IntoIterator<Item = &'a Ident>,
     members_cfg_attrs: impl IntoIterator<Item = &'a Vec<Attribute>>,
+    consts: impl IntoIterator<Item = &'a Ident>,
+    consts_values: impl IntoIterator<Item = &'a String>,
+    consts_cfg_attrs: impl IntoIterator<Item = &'a Vec<Attribute>>,
 ) -> TokenStream {
     IntrospectionNode::Map(
         [
@@ -52,6 +56,23 @@ pub fn module_introspection_code<'a>(
                         .collect(),
                 ),
             ),
+            (
+                "consts",
+                IntrospectionNode::List(
+                    consts
+                        .into_iter()
+                        .zip(consts_values)
+                        .zip(consts_cfg_attrs)
+                        .filter_map(|((ident, value), attributes)| {
+                            if attributes.is_empty() {
+                                Some(const_introspection_code(ident, value))
+                            } else {
+                                None // TODO: properly interpret cfg attributes
+                            }
+                        })
+                        .collect(),
+                ),
+            ),
         ]
         .into(),
     )
@@ -116,6 +137,20 @@ pub fn function_introspection_code(
     IntrospectionNode::Map(desc).emit(pyo3_crate_path)
 }
 
+fn const_introspection_code<'a>(ident: &'a Ident, value: &'a String) -> IntrospectionNode<'a> {
+    IntrospectionNode::Map(
+        [
+            ("type", IntrospectionNode::String("const".into())),
+            (
+                "name",
+                IntrospectionNode::String(ident.unraw().to_string().into()),
+            ),
+            ("value", IntrospectionNode::String(value.into())),
+        ]
+        .into(),
+    )
+}
+
 fn arguments_introspection_data<'a>(
     signature: &'a FunctionSignature<'a>,
     first_argument: Option<&'a str>,

pyo3-macros-backend/src/method.rs

@@ -7,7 +7,7 @@ use quote::{format_ident, quote, quote_spanned, ToTokens};
 use syn::{ext::IdentExt, spanned::Spanned, Ident, Result};
 
 use crate::pyversions::is_abi3_before;
-use crate::utils::{Ctx, LitCStr};
+use crate::utils::{expr_to_python, Ctx, LitCStr};
 use crate::{
     attributes::{FromPyWithAttribute, TextSignatureAttribute, TextSignatureAttributeValue},
     params::{impl_arg_params, Holders},
@@ -34,31 +34,7 @@ impl RegularArg<'_> {
             ..
         } = self
         {
-            match arg_default {
-                // literal values
-                syn::Expr::Lit(syn::ExprLit { lit, .. }) => match lit {
-                    syn::Lit::Str(s) => s.token().to_string(),
-                    syn::Lit::Char(c) => c.token().to_string(),
-                    syn::Lit::Int(i) => i.base10_digits().to_string(),
-                    syn::Lit::Float(f) => f.base10_digits().to_string(),
-                    syn::Lit::Bool(b) => {
-                        if b.value() {
-                            "True".to_string()
-                        } else {
-                            "False".to_string()
-                        }
-                    }
-                    _ => "...".to_string(),
-                },
-                // None
-                syn::Expr::Path(syn::ExprPath { qself, path, .. })
-                    if qself.is_none() && path.is_ident("None") =>
-                {
-                    "None".to_string()
-                }
-                // others, unsupported yet so defaults to `...`
-                _ => "...".to_string(),
-            }
+            expr_to_python(arg_default)
         } else if let RegularArg {
             option_wrapped_type: Some(..),
             ..

pyo3-macros-backend/src/module.rs

@@ -2,6 +2,7 @@
 
 #[cfg(feature = "experimental-inspect")]
 use crate::introspection::{introspection_id_const, module_introspection_code};
+use crate::utils::expr_to_python;
 use crate::{
     attributes::{
         self, kw, take_attributes, take_pyo3_options, CrateAttribute, GILUsedAttribute,
@@ -150,6 +151,7 @@ pub fn pymodule_module_impl(
 
     let mut pymodule_init = None;
     let mut module_consts = Vec::new();
+    let mut module_consts_values = Vec::new();
     let mut module_consts_cfg_attrs = Vec::new();
 
     for item in &mut *items {
@@ -293,8 +295,8 @@ pub fn pymodule_module_impl(
                 if !find_and_remove_attribute(&mut item.attrs, "pymodule_export") {
                     continue;
                 }
-
                 module_consts.push(item.ident.clone());
+                module_consts_values.push(expr_to_python(&item.expr));
                 module_consts_cfg_attrs.push(get_cfg_attributes(&item.attrs));
             }
             Item::Static(item) => {
@@ -349,6 +351,9 @@ pub fn pymodule_module_impl(
         &name.to_string(),
         &module_items,
         &module_items_cfg_attrs,
+        &module_consts,
+        &module_consts_values,
+        &module_consts_cfg_attrs,
     );
     #[cfg(not(feature = "experimental-inspect"))]
     let introspection = quote! {};
@@ -432,7 +437,8 @@ pub fn pymodule_function_impl(
     );
 
     #[cfg(feature = "experimental-inspect")]
-    let introspection = module_introspection_code(pyo3_path, &name.to_string(), &[], &[]);
+    let introspection =
+        module_introspection_code(pyo3_path, &name.to_string(), &[], &[], &[], &[], &[]);
     #[cfg(not(feature = "experimental-inspect"))]
     let introspection = quote! {};
     #[cfg(feature = "experimental-inspect")]

pyo3-macros-backend/src/utils.rs

@@ -393,3 +393,31 @@ impl TypeExt for syn::Type {
         self
     }
 }
+
+pub fn expr_to_python(expr: &syn::Expr) -> String {
+    match expr {
+        // literal values
+        syn::Expr::Lit(syn::ExprLit { lit, .. }) => match lit {
+            syn::Lit::Str(s) => s.token().to_string(),
+            syn::Lit::Char(c) => c.token().to_string(),
+            syn::Lit::Int(i) => i.base10_digits().to_string(),
+            syn::Lit::Float(f) => f.base10_digits().to_string(),
+            syn::Lit::Bool(b) => {
+                if b.value() {
+                    "True".to_string()
+                } else {
+                    "False".to_string()
+                }
+            }
+            _ => "...".to_string(),
+        },
+        // None
+        syn::Expr::Path(syn::ExprPath { qself, path, .. })
+            if qself.is_none() && path.is_ident("None") =>
+        {
+            "None".to_string()
+        }
+        // others, unsupported yet so defaults to `...`
+        _ => "...".to_string(),
+    }
+}

pytests/src/consts.rs

@@ -0,0 +1,10 @@
+use pyo3::pymodule;
+
+#[pymodule]
+pub mod consts {
+    #[pymodule_export]
+    pub const PI: f64 = std::f64::consts::PI; // Exports PI constant as part of the module
+
+    #[pymodule_export]
+    pub const SIMPLE: &str = "SIMPLE";
+}

pytests/src/lib.rs

@@ -5,6 +5,7 @@ use pyo3::wrap_pymodule;
 pub mod awaitable;
 pub mod buf_and_str;
 pub mod comparisons;
+mod consts;
 pub mod datetime;
 pub mod dict_iter;
 pub mod enums;
@@ -22,7 +23,7 @@ mod pyo3_pytests {
     use super::*;
 
     #[pymodule_export]
-    use {pyclasses::pyclasses, pyfunctions::pyfunctions};
+    use {consts::consts, pyclasses::pyclasses, pyfunctions::pyfunctions};
 
     // Inserting to sys.modules allows importing submodules nicely from Python
     // e.g. import pyo3_pytests.buf_and_str as bas

pytests/stubs/consts.pyi

@@ -0,0 +1,4 @@
+import typing
+
+PI: typing.Final = ...
+SIMPLE: typing.Final = "SIMPLE"