Merge pull request #7999 from radarhere/accept

hugovk · web-flow · commit 0c8be388781a · 2024-05-11T23:02:51.000+03:00
Added MPEG accept function
diff --git a/Tests/test_file_mpeg.py b/Tests/test_file_mpeg.py
@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from io import BytesIO
+
+import pytest
+
+from PIL import Image, MpegImagePlugin
+
+
+def test_identify() -> None:
+    # Arrange
+    b = BytesIO(b"\x00\x00\x01\xb3\x01\x00\x01")
+
+    # Act
+    with Image.open(b) as im:
+        # Assert
+        assert im.format == "MPEG"
+
+        assert im.mode == "RGB"
+        assert im.size == (16, 1)
+
+
+def test_invalid_file() -> None:
+    # Arrange
+    invalid_file = "Tests/images/flower.jpg"
+
+    # Act / Assert
+    with pytest.raises(SyntaxError):
+        MpegImagePlugin.MpegImageFile(invalid_file)
+
+
+def test_load() -> None:
+    # Arrange
+    b = BytesIO(b"\x00\x00\x01\xb3\x01\x00\x01")
+
+    with Image.open(b) as im:
+        # Act / Assert: cannot load
+        with pytest.raises(OSError):
+            im.load()
diff --git a/src/PIL/MpegImagePlugin.py b/src/PIL/MpegImagePlugin.py
@@ -53,6 +53,10 @@ def read(self, bits: int) -> int:
         return v
 
 
+def _accept(prefix: bytes) -> bool:
+    return prefix[:4] == b"\x00\x00\x01\xb3"
+
+
 ##
 # Image plugin for MPEG streams.  This plugin can identify a stream,
 # but it cannot read it.
@@ -77,7 +81,7 @@ def _open(self) -> None:
 # --------------------------------------------------------------------
 # Registry stuff
 
-Image.register_open(MpegImageFile.format, MpegImageFile)
+Image.register_open(MpegImageFile.format, MpegImageFile, _accept)
 
 Image.register_extensions(MpegImageFile.format, [".mpg", ".mpeg"])
 
