diff --git a/src/avif.rs b/src/avif.rs index 3823284..6e4565f 100644 --- a/src/avif.rs +++ b/src/avif.rs @@ -21,7 +21,7 @@ pub fn get_size(data: &[u8]) -> Option { width as u64, height as u64, MIME_TYPE.to_string(), - animated, + animated.into(), )) } diff --git a/src/bmp.rs b/src/bmp.rs index 5de8f2b..ec5e4af 100644 --- a/src/bmp.rs +++ b/src/bmp.rs @@ -3,7 +3,7 @@ use std::io::{Seek, SeekFrom}; use byteorder::{LittleEndian, ReadBytesExt}; use crate::utils::cursor_parser; -use crate::Size; +use crate::{Animation, Size}; const MIME_TYPE: &str = "image/bmp"; @@ -19,7 +19,7 @@ pub fn get_size(data: &[u8]) -> Option { width as u64, height as u64, MIME_TYPE.to_string(), - false, + Animation::No, ))) } 40 | 64 | 108 | 124 => { @@ -29,7 +29,12 @@ pub fn get_size(data: &[u8]) -> Option { if cursor.read_u8()? == 0xff { height = 4294967296 - height; } - Ok(Some(Size::new(width, height, MIME_TYPE.to_string(), false))) + Ok(Some(Size::new( + width, + height, + MIME_TYPE.to_string(), + Animation::No, + ))) } _ => Ok(None), } diff --git a/src/gif.rs b/src/gif.rs index 4e6855f..aec7a22 100644 --- a/src/gif.rs +++ b/src/gif.rs @@ -4,7 +4,7 @@ use std::io::{Cursor, Seek, SeekFrom}; use byteorder::{LittleEndian, ReadBytesExt}; use crate::utils::cursor_parser; -use crate::Size; +use crate::{Animation, Size}; const MIME_TYPE: &str = "image/gif"; @@ -22,67 +22,79 @@ pub fn get_size(data: &[u8]) -> Option { // skip Global Color Table cursor.seek(SeekFrom::Current(size))?; } - let mut found_image = false; - let mut gce_found = false; - loop { - match cursor.read_u8()? { - // Image Descriptor - 0x2c => { - if found_image { - return Ok(Some(Size::new(width, height, MIME_TYPE.to_string(), true))); - } else if !gce_found { - return Ok(Some(Size::new(width, height, MIME_TYPE.to_string(), false))); - } - found_image = true; - cursor.seek(SeekFrom::Current(8))?; - let flags = cursor.read_u8()?; - if let Some(size) = color_table_size(flags) { - cursor.seek(SeekFrom::Current(size))?; - } - // skip LZW Minimum Code Size - cursor.seek(SeekFrom::Current(1))?; - skip_data_sub_blocks(&mut cursor)?; + let animation = detect_animation(&mut cursor)?; + Ok(animation.map(|animation| Size::new(width, height, MIME_TYPE.to_string(), animation))) + }) +} + +fn detect_animation(cursor: &mut Cursor<&[u8]>) -> io::Result> { + match detect_animation_inner(cursor) { + Ok(animation) => Ok(animation.map(|a| a.into())), + Err(_) => Ok(Some(Animation::Unknown)), + } +} + +fn detect_animation_inner(cursor: &mut Cursor<&[u8]>) -> io::Result> { + let mut found_image = false; + let mut gce_found = false; + loop { + match cursor.read_u8()? { + // Image Descriptor + 0x2c => { + if found_image { + return Ok(Some(true)); + } else if !gce_found { + return Ok(Some(false)); + } + found_image = true; + cursor.seek(SeekFrom::Current(8))?; + let flags = cursor.read_u8()?; + if let Some(size) = color_table_size(flags) { + cursor.seek(SeekFrom::Current(size))?; } - // Extension - 0x21 => match cursor.read_u8()? { - // Graphic Control Extension - 0xf9 => { - gce_found = true; - // skip block size (always 4) and extension data - cursor.seek(SeekFrom::Current(5))?; - skip_data_sub_blocks(&mut cursor)?; - } - // Comment Extension - 0xfe => { - skip_data_sub_blocks(&mut cursor)?; - } - // Plain Text Extension - 0x01 => { - // skip block size (always 12) and extension data - cursor.seek(SeekFrom::Current(13))?; - skip_data_sub_blocks(&mut cursor)?; - } - // Application Extension - 0xff => { - // skip block size (always 11) and extension data - cursor.seek(SeekFrom::Current(12))?; - skip_data_sub_blocks(&mut cursor)?; - } - _ => { - return Ok(None); - } - }, - // Trailer - 0x3B => { - if found_image { - return Ok(Some(Size::new(width, height, MIME_TYPE.to_string(), false))); - } + // skip LZW Minimum Code Size + cursor.seek(SeekFrom::Current(1))?; + skip_data_sub_blocks(cursor)?; + } + // Extension + 0x21 => match cursor.read_u8()? { + // Graphic Control Extension + 0xf9 => { + gce_found = true; + // skip block size (always 4) and extension data + cursor.seek(SeekFrom::Current(5))?; + skip_data_sub_blocks(cursor)?; + } + // Comment Extension + 0xfe => { + skip_data_sub_blocks(cursor)?; + } + // Plain Text Extension + 0x01 => { + // skip block size (always 12) and extension data + cursor.seek(SeekFrom::Current(13))?; + skip_data_sub_blocks(cursor)?; + } + // Application Extension + 0xff => { + // skip block size (always 11) and extension data + cursor.seek(SeekFrom::Current(12))?; + skip_data_sub_blocks(cursor)?; + } + _ => { return Ok(None); } - _ => return Ok(None), + }, + // Trailer + 0x3B => { + if found_image { + return Ok(Some(false)); + } + return Ok(None); } + _ => return Ok(None), } - }) + } } fn color_table_size(flags: u8) -> Option { diff --git a/src/jpg.rs b/src/jpg.rs index 480828d..4075846 100644 --- a/src/jpg.rs +++ b/src/jpg.rs @@ -3,7 +3,7 @@ use std::io::{Seek, SeekFrom}; use byteorder::{BigEndian, ReadBytesExt}; use crate::utils::cursor_parser; -use crate::Size; +use crate::{Animation, Size}; const MIME_TYPE: &str = "image/jpeg"; const START_OF_FRAMES: [u8; 13] = [ @@ -23,7 +23,7 @@ pub fn get_size(data: &[u8]) -> Option { width as u64, height as u64, MIME_TYPE.to_string(), - false, + Animation::No, ))); } else { let length = cursor.read_u16::()?; diff --git a/src/lib.rs b/src/lib.rs index eb1fcfd..7ea7abf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,23 +10,50 @@ use pyo3::types::PyDict; #[cfg(test)] use serde::Deserialize; use std::array::IntoIter; -use std::collections::hash_map::DefaultHasher; -use std::hash::{Hash, Hasher}; +use std::fmt::{Display, Formatter}; + +#[pyclass(eq, eq_int)] +#[derive(Debug, Eq, PartialEq, Hash, Clone)] +pub enum Animation { + Yes, + No, + Unknown, +} -#[pyclass(get_all)] +impl From for Animation { + fn from(value: bool) -> Self { + if value { + Self::Yes + } else { + Self::No + } + } +} + +impl Display for Animation { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Animation::Yes => "yes", + Animation::No => "no", + Animation::Unknown => "unknown", + }) + } +} + +#[pyclass(get_all, eq, hash, frozen)] #[derive(Debug, Eq, PartialEq, Hash)] #[cfg_attr(test, derive(Deserialize))] pub struct Size { pub width: u64, pub height: u64, pub mime_type: String, - pub is_animated: bool, + pub is_animated: Animation, } #[pymethods] impl Size { #[new] - fn new(width: u64, height: u64, mime_type: String, is_animated: bool) -> Self { + fn new(width: u64, height: u64, mime_type: String, is_animated: Animation) -> Self { Self { width, height, @@ -39,10 +66,6 @@ impl Size { format!("{:?}", self) } - fn __eq__(&self, other: &Self) -> bool { - self == other - } - fn __iter__(slf: PyRef<'_, Self>) -> PyResult> { let itr = SizeIter { inner: [slf.width, slf.height].into_iter(), @@ -50,19 +73,13 @@ impl Size { Py::new(slf.py(), itr) } - fn __hash__(&self) -> u64 { - let mut hasher = DefaultHasher::new(); - self.hash(&mut hasher); - hasher.finish() - } - fn as_dict(&self) -> PyResult> { Python::with_gil(|py| { let dict = PyDict::new_bound(py); dict.set_item("width", self.width)?; dict.set_item("height", self.height)?; dict.set_item("mime_type", self.mime_type.clone())?; - dict.set_item("is_animated", self.is_animated)?; + dict.set_item("is_animated", self.is_animated.clone().into_py(py))?; Ok(dict.unbind()) }) } @@ -111,6 +128,7 @@ fn py_get_size(data: &[u8]) -> PyResult> { fn imgsize(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_get_size, m)?)?; m.add_class::()?; + m.add_class::()?; Ok(()) } @@ -118,9 +136,49 @@ fn imgsize(m: &Bound<'_, PyModule>) -> PyResult<()> { mod tests { use super::*; use paste::paste; + use serde::{de, Deserialize, Deserializer}; use std::collections::HashSet; + use std::fmt::Formatter; use std::path::Path; + struct AnimationVisitor; + + impl<'de> de::Visitor<'de> for AnimationVisitor { + type Value = Animation; + + fn expecting(&self, formatter: &mut Formatter) -> std::fmt::Result { + formatter.write_str("'yes', 'no' or 'unknown' or a boolean") + } + + fn visit_bool(self, v: bool) -> Result + where + E: de::Error, + { + Ok(Animation::from(v)) + } + + fn visit_str(self, v: &str) -> Result + where + E: de::Error, + { + match v { + "yes" => Ok(Animation::Yes), + "no" => Ok(Animation::No), + "unknown" => Ok(Animation::Unknown), + _ => Err(E::custom(format!("Unexpected value: {}", v))), + } + } + } + + impl<'de> Deserialize<'de> for Animation { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_any(AnimationVisitor) + } + } + macro_rules! define_tests { (impl $name:ident) => { paste! { diff --git a/src/png.rs b/src/png.rs index 815a489..76c708c 100644 --- a/src/png.rs +++ b/src/png.rs @@ -3,7 +3,7 @@ use std::io::{Read, Seek, SeekFrom}; use byteorder::{BigEndian, ReadBytesExt}; use crate::utils::cursor_parser; -use crate::Size; +use crate::{Animation, Size}; const MIME_TYPE: &str = "image/png"; @@ -12,7 +12,7 @@ pub fn get_size(data: &[u8]) -> Option { cursor.seek(SeekFrom::Start(8))?; let mut chunk_type_buf = [0u8; 4]; let mut size = None; - let mut animated = false; + let mut animated = Animation::No; loop { let chunk_length = cursor.read_u32::()?; cursor.read_exact(&mut chunk_type_buf)?; @@ -25,14 +25,14 @@ pub fn get_size(data: &[u8]) -> Option { let width = cursor.read_u32::()?; let height = cursor.read_u32::()?; size = Some((width, height)); - if animated { + if animated == Animation::Yes { break; } cursor.seek(SeekFrom::Current(chunk_length as i64 - 4))?; } // acTL [0x61, 0x63, 0x54, 0x4c] => { - animated = true; + animated = Animation::Yes; if size.is_some() { break; }