Allow reading stl from stdin

unlimitedbacon · unlimitedbacon · commit 4b6011553de5 · 2022-05-29T22:11:54.000-07:00
diff --git a/README.md b/README.md
@@ -58,8 +58,8 @@ $ stl-thumb <STL_FILE> [IMG_FILE]
 
 | Option        | Description                                             |
 | ------------- | ------------------------------------------------------- |
-| <STL_FILE>    | The STL file you want a picture of.                     |
-| [IMG_FILE]    | The thumbnail image file that will be created. If this is omitted, the image data will be dumped to stdout. |
+| <STL_FILE>    | The STL file you want a picture of. Use - to read from stdin instead of a file. |
+| <IMG_FILE>    | The thumbnail image file that will be created. Use - to write to stdout instead of a file. |
 | -s, --size \<size\>   | Specify width of the image. It will always be a square. |
 | -f, --format \<format\> | The format of the image file. If not specified it will be determined from the file extension, or default to PNG if there is no extension. Supported formats: PNG, JPEG, GIF, ICO, BMP |
 | -m, --material \<ambient\> \<diffuse\> \<specular\> | Colors for rendering the mesh using the Phong reflection model. Requires 3 colors as rgb hex values: ambient, diffuse, and specular. Defaults to blue. |
diff --git a/src/config.rs b/src/config.rs
@@ -20,7 +20,7 @@ pub enum AAMethod {
 #[derive(Clone)]
 pub struct Config {
     pub stl_filename: String,
-    pub img_filename: Option<String>,
+    pub img_filename: String,
     pub format: ImageOutputFormat,
     pub width: u32,
     pub height: u32,
@@ -35,7 +35,7 @@ impl Default for Config {
     fn default() -> Self {
         Config {
             stl_filename: "".to_string(),
-            img_filename: None,
+            img_filename: "".to_string(),
             format: ImageOutputFormat::Png,
             width: 1024,
             height: 768,
@@ -60,13 +60,14 @@ impl Config {
             .author(env!("CARGO_PKG_AUTHORS"))
             .arg(
                 clap::Arg::new("STL_FILE")
-                    .help("STL file")
+                    .help("STL file. Use - to read from stdin instead of a file.")
                     .required(true)
                     .index(1),
             )
             .arg(
                 clap::Arg::new("IMG_FILE")
-                    .help("Thumbnail image file. If this is omitted, the image data will be dumped to stdout.")
+                    .help("Thumbnail image file. Use - to write to stdout instead of a file.")
+                    .required(true)
                     .index(2),
             )
             .arg(
@@ -125,16 +126,11 @@ impl Config {
         };
 
         c.stl_filename = matches.value_of("STL_FILE").unwrap().to_string();
-        matches
-            .value_of("IMG_FILE")
-            .map(|x| c.img_filename = Some(x.to_string()));
+        c.img_filename = matches.value_of("IMG_FILE").unwrap().to_string();
         match matches.value_of("format") {
             Some(x) => c.format = match_format(x),
-            None => match &c.img_filename {
-                Some(x) => match Path::new(x).extension() {
-                    Some(ext) => c.format = match_format(ext.to_str().unwrap()),
-                    _ => (),
-                },
+            None => match Path::new(&c.img_filename).extension() {
+                Some(ext) => c.format = match_format(ext.to_str().unwrap()),
                 _ => (),
             },
         };
diff --git a/src/lib.rs b/src/lib.rs
@@ -21,7 +21,6 @@ use libc::c_char;
 use mesh::Mesh;
 use std::error::Error;
 use std::ffi::CStr;
-use std::fs::File;
 use std::{io, panic, slice, thread, time};
 
 #[cfg(target_os = "linux")]
@@ -268,9 +267,8 @@ where
 
 pub fn render_to_window(config: Config) -> Result<(), Box<dyn Error>> {
     // Get geometry from STL file
-    // =========================
-    let stl_file = File::open(&config.stl_filename)?;
-    let mesh = Mesh::from_stl(stl_file)?;
+    // ==========================
+    let mesh = Mesh::load(&config.stl_filename)?;
 
     // Create GL context
     // =================
@@ -329,10 +327,7 @@ pub fn render_to_window(config: Config) -> Result<(), Box<dyn Error>> {
 pub fn render_to_image(config: &Config) -> Result<image::DynamicImage, Box<dyn Error>> {
     // Get geometry from STL file
     // =========================
-    // TODO: Add support for URIs instead of plain file names
-    // https://developer.gnome.org/integration-guide/stable/thumbnailer.html.en
-    let stl_file = File::open(&config.stl_filename)?;
-    let mesh = Mesh::from_stl(stl_file)?;
+    let mesh = Mesh::load(&config.stl_filename)?;
 
     let img: image::DynamicImage;
 
@@ -382,9 +377,9 @@ pub fn render_to_file(config: &Config) -> Result<(), Box<dyn Error>> {
 
     // Choose output
     // Write to stdout if user did not specify a file
-    let mut output: Box<dyn io::Write> = match config.img_filename {
-        Some(ref x) => Box::new(std::fs::File::create(&x).unwrap()),
-        None => Box::new(io::stdout()),
+    let mut output: Box<dyn io::Write> = match config.img_filename.as_str() {
+        "-" => Box::new(io::stdout()),
+        _ => Box::new(std::fs::File::create(&config.img_filename).unwrap()),
     };
 
     // write_to() requires a seekable writer for performance reasons.
@@ -503,7 +498,7 @@ mod tests {
         let img_filename = "cube.png".to_string();
         let config = Config {
             stl_filename: "test_data/cube.stl".to_string(),
-            img_filename: Some(img_filename.clone()),
+            img_filename: img_filename.clone(),
             format: image::ImageOutputFormat::Png,
             ..Default::default()
         };
diff --git a/src/main.rs b/src/main.rs
@@ -26,10 +26,7 @@ fn main() {
         .unwrap();
 
     info!("STL File: {}", config.stl_filename);
-    match config.img_filename {
-        Some(ref name) => info!("Thumbnail File: {}\n", &name),
-        None => info!("Output: stdout\n"),
-    };
+    info!("IMG File: {}", config.img_filename);
 
     if config.visible {
         if let Err(e) = stl_thumb::render_to_window(config) {
diff --git a/src/mesh.rs b/src/mesh.rs
@@ -2,8 +2,9 @@ extern crate cgmath;
 extern crate stl_io;
 
 use std::error::Error;
-use std::fmt;
 use std::fs::File;
+use std::io::{Cursor, Read, Seek};
+use std::{fmt, io};
 
 #[derive(Copy, Clone)]
 pub struct Vertex {
@@ -14,7 +15,7 @@ pub struct Vertex {
 implement_vertex!(Vertex, position);
 //implement_vertex!(Vertex, position, texcoords);
 
-#[derive(Copy, Clone)]
+#[derive(Debug, Copy, Clone)]
 pub struct Normal {
     normal: [f32; 3],
 }
@@ -96,9 +97,30 @@ pub struct Mesh {
 }
 
 impl Mesh {
-    pub fn from_stl(mut stl_file: File) -> Result<Mesh, Box<dyn Error>> {
-        //let stl = stl_io::read_stl(&mut stl_file)?;
-        //debug!("{:?}", stl);
+    // Load mesh data from file (if provided) or stdin
+    pub fn load(stl_filename: &String) -> Result<Mesh, Box<dyn Error>> {
+        // TODO: Add support for URIs instead of plain file names
+        // https://developer.gnome.org/integration-guide/stable/thumbnailer.html.en
+        match stl_filename.as_str() {
+            "-" => {
+                // create_stl_reader requires Seek, so we must read the entire stream into memory before proceeding.
+                // So I guess this can just consume all RAM if it gets bad input. Hmmm....
+                let mut input_buffer = Vec::new();
+                io::stdin().read_to_end(&mut input_buffer)?;
+                Mesh::from_stl(Cursor::new(input_buffer))
+            }
+            _ => {
+                // TODO: Try BufReader and see if it's faster
+                let stl_file = File::open(&stl_filename)?;
+                Mesh::from_stl(stl_file)
+            }
+        }
+    }
+
+    pub fn from_stl<R>(mut stl_file: R) -> Result<Mesh, Box<dyn Error>>
+    where
+        R: Read + Seek,
+    {
         let mut stl_iter = stl_io::create_stl_reader(&mut stl_file)?;
 
         // Get starting point for finding bounding box
