Support "unknown" animation status

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.
ojii · Aug 22, 2024 · 8f5942b · 8f5942b
1 parent e995c3d
commit 8f5942b
Show file tree

Hide file tree

Showing 10 changed files with 188 additions and 97 deletions.
diff --git a/README.md b/README.md
@@ -16,18 +16,18 @@ Supported formats:
 ## Usage
 
 ```python
-from imgsize import get_size
+from imgsize import get_size, Size, Animation
 
 some_image_data: bytes = ...
 
-size = get_size(some_image_data)
+size: Size | None = get_size(some_image_data)
 if size is None:
     print("Could not handle data")
 else:
-    size.width
-    size.height
-    size.mime_type
-    size.is_animated
+    size.width: int
+    size.height: int
+    size.mime_type: str
+    size.is_animated: Animation
 ```
 
 You should not pass the entire image data, the first kilobyte or so should suffice for most formats, other than GIF

diff --git a/imgsize.pyi b/imgsize.pyi
@@ -1,5 +1,9 @@
 from typing import Any, Iterable, TypedDict
 
+class Animation:
+    class Yes: ...
+    class No: ...
+    class Unknown: ...
 
 class SizeDict(TypedDict):
     width: int

diff --git a/python-tests/test_apis.py b/python-tests/test_apis.py
@@ -1,10 +1,10 @@
 def test_eq(imgsize):
-    assert imgsize.Size(1, 2, 'a', True) == imgsize.Size(1, 2, 'a', True)
-    assert imgsize.Size(1, 2, 'a', True) != imgsize.Size(2, 2, 'a', True)
-    assert imgsize.Size(1, 2, 'a', True) != imgsize.Size(1, 1, 'a', True)
-    assert imgsize.Size(1, 2, 'a', True) != imgsize.Size(1, 2, 'b', True)
-    assert imgsize.Size(1, 2, 'a', True) != imgsize.Size(1, 2, 'a', False)
+    assert imgsize.Size(1, 2, 'a', imgsize.Animation.Yes) == imgsize.Size(1, 2, 'a', imgsize.Animation.Yes)
+    assert imgsize.Size(1, 2, 'a', imgsize.Animation.Yes) != imgsize.Size(2, 2, 'a', imgsize.Animation.Yes)
+    assert imgsize.Size(1, 2, 'a', imgsize.Animation.Yes) != imgsize.Size(1, 1, 'a', imgsize.Animation.Yes)
+    assert imgsize.Size(1, 2, 'a', imgsize.Animation.Yes) != imgsize.Size(1, 2, 'b', imgsize.Animation.Yes)
+    assert imgsize.Size(1, 2, 'a', imgsize.Animation.Yes) != imgsize.Size(1, 2, 'a', imgsize.Animation.No)
 
 
 def test_iter(imgsize):
-    assert list(imgsize.Size(1, 2, 'a', True)) == [1, 2]
+    assert list(imgsize.Size(1, 2, 'a', imgsize.Animation.Yes)) == [1, 2]
diff --git a/python-tests/test_sample_files.py b/python-tests/test_sample_files.py
@@ -13,9 +13,21 @@ def find_examples():
             data = fobj.read()
         with output_path.open('r') as fobj:
             output = json.load(fobj)
-        yield pytest.param(data, output, id=input_path.stem)
+        yield pytest.param(data, fix_output(output), id=input_path.stem)
 
+def fix_output(data):
+    return lambda imgsize: data | {"is_animated": animation(imgsize, data['is_animated'])}
+def animation(imgsize, v):
+    match v:
+        case True | "yes":
+            return imgsize.Animation.Yes
+        case False | "no":
+            return imgsize.Animation.No
+        case "unknown":
+            return imgsize.Animation.Unknown
+        case _ as value:
+            assert False, "unexpected value: {value!r}".format(value=value)
 
 @pytest.mark.parametrize('input,output', find_examples())
 def test_sample_files(input, output, imgsize):
-    assert imgsize.get_size(input).as_dict() == output
+    assert imgsize.get_size(input).as_dict() == output(imgsize)
diff --git a/src/avif.rs b/src/avif.rs
@@ -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(),
     ))
 }
 

diff --git 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<Size> {
                     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<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),
         }

diff --git 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<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> {

diff --git 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<Size> {
                     width as u64,
                     height as u64,
                     MIME_TYPE.to_string(),
-                    false,
+                    Animation::No,
                 )));
             } else {
                 let length = cursor.read_u16::<BigEndian>()?;