From f690b7f6915ed0cfee0b1b4246741d17aa2e70b1 Mon Sep 17 00:00:00 2001 From: Andrew Murray Date: Mon, 22 Apr 2024 13:39:35 +1000 Subject: [PATCH] Added MPEG accept function --- Tests/test_file_mpeg.py | 39 ++++++++++++++++++++++++++++++++++++++ src/PIL/MpegImagePlugin.py | 6 +++++- 2 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 Tests/test_file_mpeg.py diff --git a/Tests/test_file_mpeg.py b/Tests/test_file_mpeg.py new file mode 100644 index 000000000..468aef8a9 --- /dev/null +++ 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 index 1565612f8..ad4d3e937 100644 --- a/src/PIL/MpegImagePlugin.py +++ b/src/PIL/MpegImagePlugin.py @@ -53,6 +53,10 @@ class BitStream: 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 @@ class MpegImageFile(ImageFile.ImageFile): # -------------------------------------------------------------------- # Registry stuff -Image.register_open(MpegImageFile.format, MpegImageFile) +Image.register_open(MpegImageFile.format, MpegImageFile, _accept) Image.register_extensions(MpegImageFile.format, [".mpg", ".mpeg"])