From 0beb2228f9adb387a87c6b944cda2e427f5f264a Mon Sep 17 00:00:00 2001 From: Andrew Murray Date: Thu, 17 Oct 2024 12:44:25 +1100 Subject: [PATCH] Include JpegImageFile layers in state --- Tests/test_pickle.py | 11 +++++++++++ src/PIL/Image.py | 2 +- src/PIL/JpegImagePlugin.py | 7 +++++++ 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/Tests/test_pickle.py b/Tests/test_pickle.py index d250ba369..c4f8de013 100644 --- a/Tests/test_pickle.py +++ b/Tests/test_pickle.py @@ -74,6 +74,17 @@ def test_pickle_image( helper_pickle_file(tmp_path, protocol, test_file, test_mode) +def test_pickle_jpeg() -> None: + # Arrange + with Image.open("Tests/images/hopper.jpg") as image: + # Act: roundtrip + unpickled_image = pickle.loads(pickle.dumps(image)) + + # Assert + assert len(unpickled_image.layer) == 3 + assert unpickled_image.layers == 3 + + def test_pickle_la_mode_with_palette(tmp_path: Path) -> None: # Arrange filename = str(tmp_path / "temp.pkl") diff --git a/src/PIL/Image.py b/src/PIL/Image.py index 44270392c..ec5ce6cb1 100644 --- a/src/PIL/Image.py +++ b/src/PIL/Image.py @@ -763,7 +763,7 @@ class Image: def __setstate__(self, state: list[Any]) -> None: Image.__init__(self) - info, mode, size, palette, data = state + info, mode, size, palette, data = state[:5] self.info = info self._mode = mode self._size = size diff --git a/src/PIL/JpegImagePlugin.py b/src/PIL/JpegImagePlugin.py index 6510e072e..7dec75218 100644 --- a/src/PIL/JpegImagePlugin.py +++ b/src/PIL/JpegImagePlugin.py @@ -395,6 +395,13 @@ class JpegImageFile(ImageFile.ImageFile): return getattr(self, "_" + name) raise AttributeError(name) + def __getstate__(self) -> list[Any]: + return super().__getstate__() + [self.layers, self.layer] + + def __setstate__(self, state: list[Any]) -> None: + super().__setstate__(state) + self.layers, self.layer = state[5:] + def load_read(self, read_bytes: int) -> bytes: """ internal: read more image data