png_achunk/
encode.rs

1use std::io::Write;
2use std::path::Path;
3
4use crate::chunk::{self, Chunk};
5
6#[derive(Debug, thiserror::Error)]
7pub enum EncodingError {
8    #[error("I/O error: {0}")]
9    Io(#[from] std::io::Error),
10    #[error("Invalid chunk name {0:?}")]
11    Chunk(#[from] chunk::Error),
12}
13
14/// PNG encoder, including custom chunks
15pub struct Encoder<W: Write> {
16    writer: W,
17    custom_chunks: Vec<Chunk>,
18}
19
20impl<W: Write> Encoder<W> {
21    /// Create a neww encoder, adding the custom chunks before the image data.
22    pub fn new(writer: W) -> Self {
23        Self {
24            writer,
25            custom_chunks: vec![],
26        }
27    }
28    /// Add a custom chunk, which will be placed before the image data.
29    pub fn with_custom_chunk(mut self, chunk: Chunk) -> Self {
30        self.custom_chunks.push(chunk);
31        self
32    }
33}
34impl Encoder<std::io::BufWriter<std::fs::File>> {
35    /// Create a new encoder writing to a file
36    pub fn new_to_file(output: impl AsRef<Path>) -> Result<Self, EncodingError> {
37        let output = std::fs::OpenOptions::new()
38            .write(true)
39            .create(true)
40            .open(output)?;
41        let output = std::io::BufWriter::new(output);
42        Ok(Self::new(output))
43    }
44}
45
46fn encoding_error(err: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> image::ImageError {
47    image::ImageError::Encoding(image::error::EncodingError::new(
48        image::ImageFormat::Png.into(),
49        err,
50    ))
51}
52
53// This is similar to the implementation for image::codecs::png::PngEncoder
54impl<W: Write> image::ImageEncoder for Encoder<W> {
55    fn write_image(
56        self,
57        buf: &[u8],
58        width: u32,
59        height: u32,
60        color_type: image::ColorType,
61    ) -> image::ImageResult<()> {
62        let (ct, bits) = match color_type {
63            image::ColorType::L8 => (png::ColorType::Grayscale, png::BitDepth::Eight),
64            image::ColorType::L16 => (png::ColorType::Grayscale, png::BitDepth::Sixteen),
65            image::ColorType::La8 => (png::ColorType::GrayscaleAlpha, png::BitDepth::Eight),
66            image::ColorType::La16 => (png::ColorType::GrayscaleAlpha, png::BitDepth::Sixteen),
67            image::ColorType::Rgb8 => (png::ColorType::Rgb, png::BitDepth::Eight),
68            image::ColorType::Rgb16 => (png::ColorType::Rgb, png::BitDepth::Sixteen),
69            image::ColorType::Rgba8 => (png::ColorType::Rgba, png::BitDepth::Eight),
70            image::ColorType::Rgba16 => (png::ColorType::Rgba, png::BitDepth::Sixteen),
71            _ => {
72                return Err(image::ImageError::Unsupported(
73                    image::error::UnsupportedError::from_format_and_kind(
74                        image::ImageFormat::Png.into(),
75                        image::error::UnsupportedErrorKind::Color(color_type.into()),
76                    ),
77                ))
78            }
79        };
80        let mut encoder = png::Encoder::new(self.writer, width, height);
81        encoder.set_color(ct);
82        encoder.set_depth(bits);
83        encoder.set_compression(png::Compression::Default);
84        encoder.set_filter(png::FilterType::Sub);
85        encoder.set_adaptive_filter(png::AdaptiveFilterType::Adaptive);
86        let mut writer = encoder.write_header().map_err(encoding_error)?;
87        for chunk in self.custom_chunks {
88            writer
89                .write_chunk(chunk.chunk_type.into(), &chunk.data)
90                .map_err(encoding_error)?
91        }
92        writer.write_image_data(buf).map_err(encoding_error)?;
93        writer.finish().map_err(encoding_error)?;
94        Ok(())
95    }
96}