Add doc strings for struct definitions

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

Hide file tree

Showing 3 changed files with 39 additions and 0 deletions.
diff --git a/codegen/main.rs b/codegen/main.rs
@@ -81,11 +81,37 @@ impl PythonBindType {
         #[cfg(windows)]
         let contents = contents.replace("\r\n", "\n");
 
+        let struct_def = format!("pub struct {struct_name}");
+        let struct_def_pos = contents.find(&struct_def).unwrap();
+
+        let mut docs = Vec::new();
+
+        let binding = contents[..struct_def_pos].to_string();
+        for line in binding.lines().rev() {
+            if line.starts_with("///") {
+                docs.push(line.trim().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<_>>(),
+            )
+        };
+
         if let Some(types) = StructBindGenerator::get_types(&contents, &struct_t_name) {
             return Some(Self::Struct(StructBindGenerator::new(
                 filename.to_string(),
                 struct_name,
                 struct_t_name,
+                struct_doc_str,
                 contents,
                 types,
             )?));

diff --git a/codegen/pyi.rs b/codegen/pyi.rs
@@ -105,6 +105,16 @@ pub fn generator(type_data: &[PythonBindType]) -> io::Result<()> {
                 write_str!(file, "    def __hash__(self) -> str: ...");
             }
             PythonBindType::Struct(gen) => {
+                if let Some(docs) = gen.struct_doc_str.as_ref() {
+                    write_str!(file, "    \"\"\"");
+
+                    for line in docs {
+                        write_fmt!(file, "    {line}");
+                    }
+
+                    write_str!(file, "    \"\"\"");
+                }
+
                 let mut python_types = Vec::new();
 
                 'outer: for variable_info in &gen.types {

diff --git a/codegen/structs.rs b/codegen/structs.rs
@@ -48,6 +48,7 @@ pub struct StructBindGenerator {
     pub filename: String,
     pub struct_name: String,
     struct_t_name: String,
+    pub struct_doc_str: Option<Vec<String>>,
     pub types: Vec<CustomType>,
     file_contents: Vec<Cow<'static, str>>,
     has_complex_pack: bool,
@@ -74,6 +75,7 @@ impl StructBindGenerator {
         filename: String,
         struct_name: String,
         struct_t_name: String,
+        struct_doc_str: Option<Vec<String>>,
         contents: String,
         types: Vec<CustomType>,
     ) -> Option<Self> {
@@ -118,6 +120,7 @@ impl StructBindGenerator {
             filename,
             struct_name,
             struct_t_name,
+            struct_doc_str,
             types,
             file_contents,
             has_complex_pack,