From 80a02d187e24842aa3fb97f0e68029f0bf629f7c Mon Sep 17 00:00:00 2001
From: Paolo Barbolini <paolo@paolo565.org>
Date: Tue, 2 Feb 2021 11:32:56 +0100
Subject: [PATCH] Allow setting the AvifEncoder ColorSpace

---
 src/codecs/avif/encoder.rs | 34 +++++++++++++++++++++++++++++++---
 src/codecs/avif/mod.rs     |  3 +--
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/src/codecs/avif/encoder.rs b/src/codecs/avif/encoder.rs
index 746b9917ee..75d0a5fb18 100644
--- a/src/codecs/avif/encoder.rs
+++ b/src/codecs/avif/encoder.rs
@@ -15,7 +15,7 @@ use crate::error::{EncodingError, ParameterError, ParameterErrorKind, Unsupporte
 
 use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
 use num_traits::Zero;
-use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
+use ravif::{Img, Config, RGBA8, encode_rgba};
 use rgb::AsPixels;
 
 /// AVIF Encoder.
@@ -27,6 +27,28 @@ pub struct AvifEncoder<W> {
     config: Config
 }
 
+/// An enumeration over supported AVIF color spaces
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum ColorSpace {
+    /// sRGB colorspace
+    Srgb,
+    /// BT.709 colorspace
+    Bt709,
+
+    #[doc(hidden)]
+    __NonExhaustive(crate::utils::NonExhaustiveMarker),
+}
+
+impl ColorSpace {
+    fn to_ravif(self) -> ravif::ColorSpace {
+        match self {
+            Self::Srgb => ravif::ColorSpace::RGB,
+            Self::Bt709 => ravif::ColorSpace::YCbCr,
+            Self::__NonExhaustive(marker) => match marker._private {},
+        }
+    }
+}
+
 impl<W: Write> AvifEncoder<W> {
     /// Create a new encoder that writes its output to `w`.
     pub fn new(w: W) -> Self {
@@ -49,11 +71,17 @@ impl<W: Write> AvifEncoder<W> {
                 alpha_quality: quality,
                 speed,
                 premultiplied_alpha: false,
-                color_space: ColorSpace::RGB,
+                color_space: ravif::ColorSpace::RGB,
             } 
         }
     }
 
+    /// Encode with the specified `color_space`.
+    pub fn with_colorspace(mut self, color_space: ColorSpace) -> Self {
+        self.config.color_space = color_space.to_ravif();
+        self
+    }
+
     /// Encode image data with the indicated color type.
     ///
     /// The encoder currently requires all data to be RGBA8, it will be converted internally if
@@ -198,4 +226,4 @@ impl<W: Write> AvifEncoder<W> {
                 )))
         }
     }
-}
\ No newline at end of file
+}
diff --git a/src/codecs/avif/mod.rs b/src/codecs/avif/mod.rs
index 4b8a11d616..3e09ece51b 100644
--- a/src/codecs/avif/mod.rs
+++ b/src/codecs/avif/mod.rs
@@ -3,9 +3,8 @@
 /// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
 ///
 /// [AVIF]: https://aomediacodec.github.io/av1-avif/
-
 pub use self::decoder::AvifDecoder;
-pub use self::encoder::AvifEncoder;
+pub use self::encoder::{AvifEncoder, ColorSpace};
 
 mod decoder;
 mod encoder;