Generate doc comments for struct fields

VirxEC · Jan 2, 2025 · 1d6c39d · 1d6c39d
1 parent 29a99ec
commit 1d6c39d
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 7 deletions.
diff --git a/codegen/main.rs b/codegen/main.rs
@@ -86,10 +86,9 @@ impl PythonBindType {
 
         let mut docs = Vec::new();
 
-        let binding = contents[..struct_def_pos].to_string();
-        for line in binding.lines().rev() {
+        for line in contents[..struct_def_pos].lines().rev() {
             if line.starts_with("///") {
-                docs.push(line.trim().trim_start_matches("///").trim());
+                docs.push(line.trim_start_matches("///").trim());
             } else {
                 break;
             }

diff --git a/codegen/pyi.rs b/codegen/pyi.rs
@@ -125,6 +125,17 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                         if variable_type == rust_type {
                             python_types.push(python_type.to_string());
                             write_fmt!(file, "    {variable_name}: {python_type}");
+
+                            if let Some(docs) = variable_info.doc_str.as_ref() {
+                                write_str!(file, "    \"\"\"");
+
+                                for line in docs {
+                                    write_fmt!(file, "    {line}");
+                                }
+
+                                write_str!(file, "    \"\"\"");
+                            }
+
                             continue 'outer;
                         }
                     }
@@ -227,6 +238,16 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                             write_fmt!(file, "    {variable_name}: {type_name}");
                         }
                     }
+
+                    if let Some(docs) = variable_info.doc_str.as_ref() {
+                        write_str!(file, "    \"\"\"");
+
+                        for line in docs {
+                            write_fmt!(file, "    {line}");
+                        }
+
+                        write_str!(file, "    \"\"\"");
+                    }
                 }
 
                 if !gen.types.is_empty() {

diff --git a/codegen/structs.rs b/codegen/structs.rs
@@ -42,6 +42,7 @@ pub struct CustomType {
     pub frozen_needs_py: bool,
     pub is_special_base: Option<SpecialBase>,
     pub snake_case_name: String,
+    pub doc_str: Option<Vec<String>>,
 }
 
 pub struct StructBindGenerator {
@@ -153,10 +154,38 @@ impl StructBindGenerator {
         let raw_types: Vec<_> = struct_definition
             .split('\n')
             .filter_map(|s| {
-                s.trim_start_matches(' ')
+                let (name, raw_type) = s
+                    .trim_start_matches(' ')
                     .trim_start_matches("pub ")
                     .trim_end_matches(',')
-                    .split_once(": ")
+                    .split_once(": ")?;
+
+                let var_def = format!("pub fn {name}(");
+                let var_def_pos = contents.find(&var_def).unwrap();
+
+                let mut docs = Vec::new();
+
+                for line in contents[..var_def_pos].lines().rev().skip(2) {
+                    let line = line.trim();
+                    if line.starts_with("///") {
+                        docs.push(line.trim_start_matches("///").trim());
+                    } else {
+                        break;
+                    }
+                }
+
+                let struct_doc_str = if docs.is_empty() {
+                    None
+                } else {
+                    Some(
+                        docs.into_iter()
+                            .map(|s| s.to_string())
+                            .rev()
+                            .collect::<Vec<_>>(),
+                    )
+                };
+
+                Some((name, raw_type, struct_doc_str))
             })
             .collect();
 
@@ -165,10 +194,10 @@ impl StructBindGenerator {
         Some(custom_types)
     }
 
-    fn raw_types_to_custom(raw_types: Vec<(&str, &str)>) -> Vec<CustomType> {
+    fn raw_types_to_custom(raw_types: Vec<(&str, &str, Option<Vec<String>>)>) -> Vec<CustomType> {
         raw_types
             .into_iter()
-            .map(|(name, raw_type)| {
+            .map(|(name, raw_type, doc_str)| {
                 let (rust_type, inner_type) = if raw_type.starts_with("Vec<") {
                     if raw_type == "Vec<u8>" {
                         (RustType::Vec(InnerVecType::U8), None)
@@ -266,6 +295,7 @@ impl StructBindGenerator {
                     frozen_needs_py,
                     is_special_base,
                     snake_case_name: String::new(),
+                    doc_str,
                 }
             })
             .collect()