Introspection: split module and name in decorators (#5618)

Introspection API: Introduces the PythonIdentifier struct to store an optional module name + a name inside the module or local

Author: Tpt

newsfragments/5618.changed.md

@@ -0,0 +1 @@
+Introspection: properly represent decorators as module + name (allows to get imports working)

pyo3-introspection/src/introspection.rs

@@ -1,6 +1,6 @@
 use crate::model::{
-    Argument, Arguments, Attribute, Class, Function, Module, TypeHint, TypeHintExpr,
-    VariableLengthArgument,
+    Argument, Arguments, Attribute, Class, Function, Module, PythonIdentifier, TypeHint,
+    TypeHintExpr, VariableLengthArgument,
 };
 use anyhow::{anyhow, bail, ensure, Context, Result};
 use goblin::elf::section_header::SHN_XINDEX;
@@ -134,7 +134,7 @@ fn convert_members<'a>(
                 parent: _,
                 decorators,
                 returns,
-            } => functions.push(convert_function(name, arguments, decorators, returns)),
+            } => functions.push(convert_function(name, arguments, decorators, returns)?),
             Chunk::Attribute {
                 name,
                 id: _,
@@ -149,14 +149,24 @@ fn convert_members<'a>(
     classes.sort_by(|l, r| l.name.cmp(&r.name));
     functions.sort_by(|l, r| match l.name.cmp(&r.name) {
         Ordering::Equal => {
-            // We put the getter before the setter
-            if l.decorators.iter().any(|d| d == "property") {
+            // We put the getter before the setter. For that, we put @property before the other ones
+            if l.decorators
+                .iter()
+                .any(|d| d.name == "property" && d.module.as_deref() == Some("builtins"))
+            {
                 Ordering::Less
-            } else if r.decorators.iter().any(|d| d == "property") {
+            } else if r
+                .decorators
+                .iter()
+                .any(|d| d.name == "property" && d.module.as_deref() == Some("builtins"))
+            {
                 Ordering::Greater
             } else {
                 // We pick an ordering based on decorators
-                l.decorators.cmp(&r.decorators)
+                l.decorators
+                    .iter()
+                    .map(|d| &d.name)
+                    .cmp(r.decorators.iter().map(|d| &d.name))
             }
         }
         o => o,
@@ -194,12 +204,27 @@ fn convert_class(
 fn convert_function(
     name: &str,
     arguments: &ChunkArguments,
-    decorators: &[String],
+    decorators: &[ChunkTypeHint],
     returns: &Option<ChunkTypeHint>,
-) -> Function {
-    Function {
+) -> Result<Function> {
+    Ok(Function {
         name: name.into(),
-        decorators: decorators.to_vec(),
+        decorators: decorators
+            .iter()
+            .map(|d| match convert_type_hint(d) {
+                TypeHint::Plain(id) => Ok(PythonIdentifier {
+                    module: None,
+                    name: id.clone(),
+                }),
+                TypeHint::Ast(expr) => {
+                    if let TypeHintExpr::Identifier(i) = expr {
+                        Ok(i)
+                    } else {
+                        bail!("A decorator must be the identifier of a Python function")
+                    }
+                }
+            })
+            .collect::<Result<_>>()?,
         arguments: Arguments {
             positional_only_arguments: arguments.posonlyargs.iter().map(convert_argument).collect(),
             arguments: arguments.args.iter().map(convert_argument).collect(),
@@ -214,7 +239,7 @@ fn convert_function(
                 .map(convert_variable_length_argument),
         },
         returns: returns.as_ref().map(convert_type_hint),
-    }
+    })
 }
 
 fn convert_argument(arg: &ChunkArgument) -> Argument {
@@ -253,15 +278,24 @@ fn convert_type_hint(arg: &ChunkTypeHint) -> TypeHint {
 
 fn convert_type_hint_expr(expr: &ChunkTypeHintExpr) -> TypeHintExpr {
     match expr {
-        ChunkTypeHintExpr::Local { id } => TypeHintExpr::Local { id: id.clone() },
-        ChunkTypeHintExpr::Builtin { id } => TypeHintExpr::Builtin { id: id.clone() },
-        ChunkTypeHintExpr::Attribute { module, attr } => TypeHintExpr::Attribute {
-            module: module.clone(),
-            attr: attr.clone(),
-        },
-        ChunkTypeHintExpr::Union { elts } => TypeHintExpr::Union {
-            elts: elts.iter().map(convert_type_hint_expr).collect(),
-        },
+        ChunkTypeHintExpr::Local { id } => PythonIdentifier {
+            module: None,
+            name: id.clone(),
+        }
+        .into(),
+        ChunkTypeHintExpr::Builtin { id } => PythonIdentifier {
+            module: Some("builtins".into()),
+            name: id.clone(),
+        }
+        .into(),
+        ChunkTypeHintExpr::Attribute { module, attr } => PythonIdentifier {
+            module: Some(module.clone()),
+            name: attr.clone(),
+        }
+        .into(),
+        ChunkTypeHintExpr::Union { elts } => {
+            TypeHintExpr::Union(elts.iter().map(convert_type_hint_expr).collect())
+        }
         ChunkTypeHintExpr::Subscript { value, slice } => TypeHintExpr::Subscript {
             value: Box::new(convert_type_hint_expr(value)),
             slice: slice.iter().map(convert_type_hint_expr).collect(),
@@ -419,8 +453,8 @@ enum Chunk {
         #[serde(default)]
         parent: Option<String>,
         #[serde(default)]
-        decorators: Vec<String>,
-        #[serde(default, deserialize_with = "deserialize_type_hint")]
+        decorators: Vec<ChunkTypeHint>,
+        #[serde(default)]
         returns: Option<ChunkTypeHint>,
     },
     Attribute {
@@ -431,7 +465,7 @@ enum Chunk {
         name: String,
         #[serde(default)]
         value: Option<String>,
-        #[serde(default, deserialize_with = "deserialize_type_hint")]
+        #[serde(default)]
         annotation: Option<ChunkTypeHint>,
     },
 }
@@ -455,7 +489,7 @@ struct ChunkArgument {
     name: String,
     #[serde(default)]
     default: Option<String>,
-    #[serde(default, deserialize_with = "deserialize_type_hint")]
+    #[serde(default)]
     annotation: Option<ChunkTypeHint>,
 }
 
@@ -467,6 +501,45 @@ enum ChunkTypeHint {
     Plain(String),
 }
 
+impl<'de> Deserialize<'de> for ChunkTypeHint {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct AnnotationVisitor;
+
+        impl<'de> Visitor<'de> for AnnotationVisitor {
+            type Value = ChunkTypeHint;
+
+            fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+                formatter.write_str("annotation")
+            }
+
+            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
+            where
+                E: Error,
+            {
+                self.visit_string(v.into())
+            }
+
+            fn visit_string<E>(self, v: String) -> Result<Self::Value, E>
+            where
+                E: Error,
+            {
+                Ok(ChunkTypeHint::Plain(v))
+            }
+
+            fn visit_map<M: MapAccess<'de>>(self, map: M) -> Result<ChunkTypeHint, M::Error> {
+                Ok(ChunkTypeHint::Ast(Deserialize::deserialize(
+                    MapAccessDeserializer::new(map),
+                )?))
+            }
+        }
+
+        deserializer.deserialize_any(AnnotationVisitor)
+    }
+}
+
 #[derive(Deserialize)]
 #[serde(tag = "type", rename_all = "lowercase")]
 enum ChunkTypeHintExpr {
@@ -488,39 +561,3 @@ enum ChunkTypeHintExpr {
         slice: Vec<ChunkTypeHintExpr>,
     },
 }
-
-fn deserialize_type_hint<'de, D: Deserializer<'de>>(
-    deserializer: D,
-) -> Result<Option<ChunkTypeHint>, D::Error> {
-    struct AnnotationVisitor;
-
-    impl<'de> Visitor<'de> for AnnotationVisitor {
-        type Value = ChunkTypeHint;
-
-        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
-            formatter.write_str("annotation")
-        }
-
-        fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
-        where
-            E: Error,
-        {
-            self.visit_string(v.into())
-        }
-
-        fn visit_string<E>(self, v: String) -> Result<Self::Value, E>
-        where
-            E: Error,
-        {
-            Ok(ChunkTypeHint::Plain(v))
-        }
-
-        fn visit_map<M: MapAccess<'de>>(self, map: M) -> Result<ChunkTypeHint, M::Error> {
-            Ok(ChunkTypeHint::Ast(Deserialize::deserialize(
-                MapAccessDeserializer::new(map),
-            )?))
-        }
-    }
-
-    Ok(Some(deserializer.deserialize_any(AnnotationVisitor)?))
-}

pyo3-introspection/src/model.rs

@@ -19,7 +19,7 @@ pub struct Class {
 pub struct Function {
     pub name: String,
     /// decorator like 'property' or 'staticmethod'
-    pub decorators: Vec<String>,
+    pub decorators: Vec<PythonIdentifier>,
     pub arguments: Arguments,
     /// return type
     pub returns: Option<TypeHint>,
@@ -77,17 +77,27 @@ pub enum TypeHint {
 /// A type hint annotation as an AST fragment
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub enum TypeHintExpr {
-    /// A local identifier which module is unknown
-    Local { id: String },
-    /// A Python builtin like `int`
-    Builtin { id: String },
-    /// The attribute of a python object like `{value}.{attr}`
-    Attribute { module: String, attr: String },
+    /// An identifier
+    Identifier(PythonIdentifier),
     /// A union `{left} | {right}`
-    Union { elts: Vec<TypeHintExpr> },
+    Union(Vec<TypeHintExpr>),
     /// A subscript `{value}[*slice]`
     Subscript {
         value: Box<TypeHintExpr>,
         slice: Vec<TypeHintExpr>,
     },
 }
+
+impl From<PythonIdentifier> for TypeHintExpr {
+    #[inline]
+    fn from(value: PythonIdentifier) -> Self {
+        Self::Identifier(value)
+    }
+}
+
+/// An Python identifier, either local (with `module = None`) or global (with `module = Some(_)`)
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct PythonIdentifier {
+    pub module: Option<String>,
+    pub name: String,
+}