Skip to content

Commit

Permalink
Include JpegImageFile layers in state
Browse files Browse the repository at this point in the history
  • Loading branch information
radarhere committed Oct 17, 2024
1 parent 11c654c commit 37ca2ec
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 1 deletion.
11 changes: 11 additions & 0 deletions Tests/test_pickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion src/PIL/Image.py
Original file line number Diff line number Diff line change
Expand Up @@ -763,7 +763,7 @@ def __getstate__(self) -> list[Any]:

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
Expand Down
8 changes: 8 additions & 0 deletions src/PIL/JpegImagePlugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,14 @@ def __getattr__(self, name: str) -> Any:
return getattr(self, "_" + name)
raise AttributeError(name)

def __getstate__(self) -> list[Any]:
return super().__getstate__() + [self.layer, self.layers]

def __setstate__(self, state: list[Any]) -> None:
super().__setstate__(state)
self.layer = state[-2]
self.layers = state[-1]

def load_read(self, read_bytes: int) -> bytes:
"""
internal: read more image data
Expand Down

0 comments on commit 37ca2ec

Please sign in to comment.