Generate modules to make easy importing

src/bin/generator.rs

@@ -11,9 +11,13 @@ struct Args {
     #[arg(short, long)]
     input: PathBuf,
 
-    /// Path to the output Rust file (.rs)
+    /// Path to the output directory
     #[arg(short, long)]
     output: PathBuf,
+
+    /// Files to generate. If omitted, all files are generated.
+    #[arg(short, long, value_delimiter = ',')]
+    files: Option<Vec<String>>,
 }
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -21,8 +25,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let data = fs::read(&args.input)?;
     let set = FileDescriptorSet::new(&data).expect("Failed to parse FileDescriptorSet");
 
-    let output = generate_rust_code(&set);
+    let files = generate_rust_code(&set, args.files.as_deref(), true);
 
-    fs::write(&args.output, output)?;
+    for (filename, content) in files {
+        let path = args.output.join(filename);
+        if let Some(parent) = path.parent() {
+            fs::create_dir_all(parent)?;
+        }
+        fs::write(path, content)?;
+    }
     Ok(())
 }

src/bin/protoc-gen-roto.rs

@@ -67,36 +67,35 @@ fn handle_request(request: &CodeGeneratorRequest) -> std::result::Result<Vec<u8>
 
     let set = FileDescriptorSet::new(&set_buf)?;
 
+    let files_to_generate: Vec<String> = request.file_to_generate()
+        .filter_map(|res| {
+            let (bytes, _) = res.ok()?;
+            std::str::from_utf8(bytes).ok().map(|s| s.to_string())
+        })
+        .collect();
+
     // Generate the Rust code
     info!("Generating Rust code from descriptor set...");
-    let generated_code = generate_rust_code(&set);
-
-    // Determine the output filename
-    let mut output_filename = "roto_generated.rs".to_string();
-    if let Some(first_file) = request.file_to_generate().next() {
-        if let Ok((name_bytes, _)) = first_file {
-            if let Ok(name) = std::str::from_utf8(name_bytes) {
-                output_filename = format!("{}.rs", name.replace(".proto", ""));
-            }
-        }
-    }
+    let generated_files = generate_rust_code(&set, Some(&files_to_generate), false);
 
     // Construct the response
     let mut response_buf = vec![0u8; 1024 * 1024 * 2]; // Allocate 2MB for response
     let mut resp_builder = CodeGeneratorResponse::builder(&mut response_buf);
 
-    let mut file_buf = vec![0u8; 1024 * 1024 * 2];
-    let final_file = ResponseFile::builder(&mut file_buf)
-        .name(&output_filename)?
-        .content(&generated_code)?
-        .finish()
-        .map_err(|e| {
-            error!("Failed to build ResponseFile: {:?}", e);
-            e
-        })?;
+    for (filename, content) in generated_files {
+        let mut file_buf = vec![0u8; 1024 * 1024 * 2];
+        let final_file = ResponseFile::builder(&mut file_buf)
+            .name(&filename)?
+            .content(&content)?
+            .finish()
+            .map_err(|e| {
+                error!("Failed to build ResponseFile {}: {:?}", filename, e);
+                e
+            })?;
+        resp_builder = resp_builder.add_file(final_file)?;
+    }
 
     let final_response_slice = resp_builder
-        .add_file(final_file)?
         .finish()
         .map_err(|e| {
             error!("Failed to finish CodeGeneratorResponse: {:?}", e);