Introspection: Introduce TypeHint struct

newsfragments/5438.changed.md

@@ -0,0 +1 @@
+Introspection: introduce `TypeHint` and make use of it to encode type hint annotations.

pyo3-introspection/Cargo.toml

@@ -13,7 +13,6 @@ anyhow = "1"
 goblin = ">=0.9, <0.11"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
-unicode-ident = "1"
 
 [dev-dependencies]
 tempfile = "3.12.0"

pyo3-introspection/src/introspection.rs

@@ -1,5 +1,6 @@
 use crate::model::{
-    Argument, Arguments, Attribute, Class, Function, Module, VariableLengthArgument,
+    Argument, Arguments, Attribute, Class, Function, Module, TypeHint, TypeHintImport,
+    VariableLengthArgument,
 };
 use anyhow::{bail, ensure, Context, Result};
 use goblin::elf::Elf;
@@ -8,11 +9,12 @@ use goblin::mach::symbols::{NO_SECT, N_SECT};
 use goblin::mach::{Mach, MachO, SingleArch};
 use goblin::pe::PE;
 use goblin::Object;
-use serde::Deserialize;
+use serde::de::{Error, MapAccess, Visitor};
+use serde::{de, Deserialize, Deserializer};
 use std::cmp::Ordering;
 use std::collections::HashMap;
-use std::fs;
 use std::path::Path;
+use std::{fmt, fs};
 
 /// Introspect a cdylib built with PyO3 and returns the definition of a Python module.
 ///
@@ -191,7 +193,7 @@ fn convert_function(
     name: &str,
     arguments: &ChunkArguments,
     decorators: &[String],
-    returns: &Option<String>,
+    returns: &Option<ChunkTypeHint>,
 ) -> Function {
     Function {
         name: name.into(),
@@ -209,30 +211,48 @@ fn convert_function(
                 .as_ref()
                 .map(convert_variable_length_argument),
         },
-        returns: returns.clone(),
+        returns: returns.as_ref().map(convert_type_hint),
     }
 }
 
 fn convert_argument(arg: &ChunkArgument) -> Argument {
     Argument {
         name: arg.name.clone(),
         default_value: arg.default.clone(),
-        annotation: arg.annotation.clone(),
+        annotation: arg.annotation.as_ref().map(convert_type_hint),
     }
 }
 
 fn convert_variable_length_argument(arg: &ChunkArgument) -> VariableLengthArgument {
     VariableLengthArgument {
         name: arg.name.clone(),
-        annotation: arg.annotation.clone(),
+        annotation: arg.annotation.as_ref().map(convert_type_hint),
     }
 }
 
-fn convert_attribute(name: &str, value: &Option<String>, annotation: &Option<String>) -> Attribute {
+fn convert_attribute(
+    name: &str,
+    value: &Option<String>,
+    annotation: &Option<ChunkTypeHint>,
+) -> Attribute {
     Attribute {
         name: name.into(),
         value: value.clone(),
-        annotation: annotation.clone(),
+        annotation: annotation.as_ref().map(convert_type_hint),
+    }
+}
+
+fn convert_type_hint(arg: &ChunkTypeHint) -> TypeHint {
+    TypeHint {
+        annotation: arg.annotation.clone(),
+        imports: arg.imports.iter().map(convert_type_hint_import).collect(),
+    }
+}
+
+fn convert_type_hint_import(arg: &ChunkTypeHintImport) -> TypeHintImport {
+    TypeHintImport {
+        module: arg.module.clone(),
+        name: arg.name.clone(),
     }
 }
 
@@ -414,8 +434,8 @@ enum Chunk {
         parent: Option<String>,
         #[serde(default)]
         decorators: Vec<String>,
-        #[serde(default)]
-        returns: Option<String>,
+        #[serde(default, deserialize_with = "deserialize_annotation")]
+        returns: Option<ChunkTypeHint>,
     },
     Attribute {
         #[serde(default)]
@@ -425,8 +445,8 @@ enum Chunk {
         name: String,
         #[serde(default)]
         value: Option<String>,
-        #[serde(default)]
-        annotation: Option<String>,
+        #[serde(default, deserialize_with = "deserialize_annotation")]
+        annotation: Option<ChunkTypeHint>,
     },
 }
 
@@ -449,6 +469,56 @@ struct ChunkArgument {
     name: String,
     #[serde(default)]
     default: Option<String>,
+    #[serde(default, deserialize_with = "deserialize_annotation")]
+    annotation: Option<ChunkTypeHint>,
+}
+
+#[derive(Deserialize)]
+struct ChunkTypeHint {
+    annotation: String,
     #[serde(default)]
-    annotation: Option<String>,
+    imports: Vec<ChunkTypeHintImport>,
+}
+
+#[derive(Deserialize)]
+struct ChunkTypeHintImport {
+    module: String,
+    name: String,
+}
+
+fn deserialize_annotation<'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 {
+                annotation: v,
+                imports: Vec::new(),
+            })
+        }
+
+        fn visit_map<M: MapAccess<'de>>(self, map: M) -> Result<ChunkTypeHint, M::Error> {
+            Deserialize::deserialize(de::value::MapAccessDeserializer::new(map))
+        }
+    }
+
+    Ok(Some(deserializer.deserialize_any(AnnotationVisitor)?))
 }

pyo3-introspection/src/model.rs

@@ -22,7 +22,7 @@ pub struct Function {
     pub decorators: Vec<String>,
     pub arguments: Arguments,
     /// return type
-    pub returns: Option<String>,
+    pub returns: Option<TypeHint>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
@@ -31,7 +31,7 @@ pub struct Attribute {
     /// Value as a Python expression if easily expressible
     pub value: Option<String>,
     /// Type annotation as a Python expression
-    pub annotation: Option<String>,
+    pub annotation: Option<TypeHint>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
@@ -54,13 +54,27 @@ pub struct Argument {
     /// Default value as a Python expression
     pub default_value: Option<String>,
     /// Type annotation as a Python expression
-    pub annotation: Option<String>,
+    pub annotation: Option<TypeHint>,
 }
 
 /// A variable length argument ie. *vararg or **kwarg
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct VariableLengthArgument {
     pub name: String,
     /// Type annotation as a Python expression
-    pub annotation: Option<String>,
+    pub annotation: Option<TypeHint>,
+}
+
+/// A type hint annotation with the required modules to import
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct TypeHint {
+    pub annotation: String,
+    pub imports: Vec<TypeHintImport>,
+}
+
+/// An import required to make the type hint valid like `from {module} import {name}`
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct TypeHintImport {
+    pub module: String,
+    pub name: String,
 }