Skip to main content

aether_renderer_capture/
sink.rs

1//! Where captured frames go: a folder, a ZIP, or stdout.
2//!
3//! Folder and ZIP use the layout `aether-renderer-core` reads as input:
4//! `frame_00000.png`, `frame_00001.png`, … at the top level, in frame order.
5
6use std::fs::File;
7use std::io::{BufWriter, StdoutLock, Write};
8use std::path::{Path, PathBuf};
9
10use anyhow::{Context, Result};
11use zip::write::SimpleFileOptions;
12use zip::{CompressionMethod, DateTime, ZipWriter};
13
14pub enum FrameSink {
15    Dir(PathBuf),
16    Zip {
17        writer: Box<ZipWriter<BufWriter<File>>>,
18        partial: PathBuf,
19        dest: PathBuf,
20    },
21    Stdout(StdoutLock<'static>),
22}
23
24impl FrameSink {
25    /// `--stdout` wins; otherwise a path ending in `.zip` becomes a ZIP and
26    /// anything else a folder.
27    pub fn open(out: &Path, stdout: bool) -> Result<Self> {
28        if stdout {
29            return Ok(Self::Stdout(std::io::stdout().lock()));
30        }
31        if is_zip(out) {
32            if let Some(parent) = out.parent().filter(|p| !p.as_os_str().is_empty()) {
33                std::fs::create_dir_all(parent)?;
34            }
35            // Written next to the target and renamed on success, so a failed
36            // capture never leaves a half-written ZIP that looks complete.
37            let mut partial = out.as_os_str().to_owned();
38            partial.push(".partial");
39            let partial = PathBuf::from(partial);
40            let file = File::create(&partial)
41                .with_context(|| format!("creating {}", partial.display()))?;
42            return Ok(Self::Zip {
43                writer: Box::new(ZipWriter::new(BufWriter::new(file))),
44                partial,
45                dest: out.to_path_buf(),
46            });
47        }
48        std::fs::create_dir_all(out).with_context(|| format!("creating {}", out.display()))?;
49        Ok(Self::Dir(out.to_path_buf()))
50    }
51
52    pub fn write(&mut self, index: u64, png: &[u8]) -> Result<()> {
53        match self {
54            Self::Dir(dir) => std::fs::write(dir.join(frame_name(index)), png)?,
55            Self::Zip { writer, .. } => {
56                // PNG is already compressed, so deflate would only cost time.
57                // A fixed timestamp keeps the ZIP itself byte-identical per run.
58                let options = SimpleFileOptions::default()
59                    .compression_method(CompressionMethod::Stored)
60                    .last_modified_time(DateTime::default());
61                writer.start_file(frame_name(index), options)?;
62                writer.write_all(png)?;
63            }
64            Self::Stdout(out) => out.write_all(png)?,
65        }
66        Ok(())
67    }
68
69    /// Flush and, for a ZIP, write the central directory and move it into place.
70    pub fn finish(self) -> Result<()> {
71        match self {
72            Self::Dir(_) => {}
73            Self::Zip {
74                writer,
75                partial,
76                dest,
77            } => {
78                writer.finish()?.flush()?;
79                std::fs::rename(&partial, &dest)
80                    .with_context(|| format!("moving {} into place", dest.display()))?;
81            }
82            Self::Stdout(mut out) => out.flush()?,
83        }
84        Ok(())
85    }
86
87    pub fn is_stdout(&self) -> bool {
88        matches!(self, Self::Stdout(_))
89    }
90}
91
92fn is_zip(path: &Path) -> bool {
93    path.extension()
94        .is_some_and(|e| e.eq_ignore_ascii_case("zip"))
95}
96
97fn frame_name(index: u64) -> String {
98    format!("frame_{index:05}.png")
99}