Skip to content

Commit

Permalink
Support "unknown" animation status
Browse files Browse the repository at this point in the history
GIF might require reading deeply into the data stream to conclusively
decide whether an image is animated or not. This change allows "unknown"
animation statuses, for the case where we found the image
type/dimensions but have not enough data to conclusively say whether
it's animated or not.
  • Loading branch information
ojii committed Aug 22, 2024
1 parent dcff5f5 commit 2b0c55e
Show file tree
Hide file tree
Showing 6 changed files with 158 additions and 83 deletions.
2 changes: 1 addition & 1 deletion src/avif.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
width as u64,
height as u64,
MIME_TYPE.to_string(),
animated,
animated.into(),
))
}

Expand Down
11 changes: 8 additions & 3 deletions src/bmp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand All @@ -19,7 +19,7 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
width as u64,
height as u64,
MIME_TYPE.to_string(),
false,
Animation::No,
)))
}
40 | 64 | 108 | 124 => {
Expand All @@ -29,7 +29,12 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
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),
}
Expand Down
126 changes: 69 additions & 57 deletions src/gif.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand All @@ -22,67 +22,79 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
// 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<Option<Animation>> {
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<Option<bool>> {
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<i64> {
Expand Down
4 changes: 2 additions & 2 deletions src/jpg.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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] = [
Expand All @@ -23,7 +23,7 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
width as u64,
height as u64,
MIME_TYPE.to_string(),
false,
Animation::No,
)));
} else {
let length = cursor.read_u16::<BigEndian>()?;
Expand Down
90 changes: 74 additions & 16 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<bool> 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,
Expand All @@ -39,30 +66,20 @@ impl Size {
format!("{:?}", self)
}

fn __eq__(&self, other: &Self) -> bool {
self == other
}

fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<SizeIter>> {
let itr = SizeIter {
inner: [slf.width, slf.height].into_iter(),
};
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<Py<PyDict>> {
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())
})
}
Expand Down Expand Up @@ -111,16 +128,57 @@ fn py_get_size(data: &[u8]) -> PyResult<Option<Size>> {
fn imgsize(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(py_get_size, m)?)?;
m.add_class::<Size>()?;
m.add_class::<Animation>()?;
Ok(())
}

#[cfg(test)]
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<E>(self, v: bool) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(Animation::from(v))
}

fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
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<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
deserializer.deserialize_any(AnimationVisitor)
}
}

macro_rules! define_tests {
(impl $name:ident) => {
paste! {
Expand Down
8 changes: 4 additions & 4 deletions src/png.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand All @@ -12,7 +12,7 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
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::<BigEndian>()?;
cursor.read_exact(&mut chunk_type_buf)?;
Expand All @@ -25,14 +25,14 @@ pub fn get_size(data: &[u8]) -> Option<Size> {
let width = cursor.read_u32::<BigEndian>()?;
let height = cursor.read_u32::<BigEndian>()?;
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;
}
Expand Down

0 comments on commit 2b0c55e

Please sign in to comment.