Added frombytes() to allow for unlimited parsing

This commit is contained in:
Andrew Murray 2025-03-20 23:12:35 +11:00
parent 0888dc02ac
commit 510bc05577
2 changed files with 38 additions and 9 deletions

View File

@ -16,11 +16,11 @@ def test_sanity() -> None:
GimpPaletteFile(fp) GimpPaletteFile(fp)
with open("Tests/images/bad_palette_file.gpl", "rb") as fp: with open("Tests/images/bad_palette_file.gpl", "rb") as fp:
with pytest.raises(SyntaxError): with pytest.raises(SyntaxError, match="bad palette file"):
GimpPaletteFile(fp) GimpPaletteFile(fp)
with open("Tests/images/bad_palette_entry.gpl", "rb") as fp: with open("Tests/images/bad_palette_entry.gpl", "rb") as fp:
with pytest.raises(ValueError): with pytest.raises(ValueError, match="bad palette entry"):
GimpPaletteFile(fp) GimpPaletteFile(fp)
@ -40,12 +40,26 @@ def test_get_palette(filename: str, size: int) -> None:
assert len(palette) / 3 == size assert len(palette) / 3 == size
def test_palette_limit() -> None: def test_frombytes() -> None:
with open("Tests/images/full_gimp_palette.gpl", "rb") as fp: with open("Tests/images/full_gimp_palette.gpl", "rb") as fp:
data = fp.read() full_data = fp.read()
# Test that __init__ only reads 256 entries # Test that __init__ only reads 256 entries
data = data.replace(b"#\n", b"") + b" 0 0 0 Index 256" data = full_data.replace(b"#\n", b"") + b" 0 0 0 Index 256"
b = BytesIO(data) b = BytesIO(data)
palette = GimpPaletteFile(b) palette = GimpPaletteFile(b)
assert len(palette.palette) / 3 == 256 assert len(palette.palette) / 3 == 256
# Test that frombytes() can read beyond that
palette = GimpPaletteFile.frombytes(data)
assert len(palette.palette) / 3 == 257
# Test that __init__ raises an error if a comment is too long
data = full_data[:-1] + b"a" * 100
b = BytesIO(data)
with pytest.raises(SyntaxError, match="bad palette file"):
palette = GimpPaletteFile(b)
# Test that frombytes() can read the data regardless
palette = GimpPaletteFile.frombytes(data)
assert len(palette.palette) / 3 == 256

View File

@ -16,6 +16,7 @@
from __future__ import annotations from __future__ import annotations
import re import re
from io import BytesIO
from typing import IO from typing import IO
@ -24,13 +25,18 @@ class GimpPaletteFile:
rawmode = "RGB" rawmode = "RGB"
def __init__(self, fp: IO[bytes]) -> None: def _read(self, fp: IO[bytes], limit: bool = True) -> None:
if not fp.readline().startswith(b"GIMP Palette"): if not fp.readline().startswith(b"GIMP Palette"):
msg = "not a GIMP palette file" msg = "not a GIMP palette file"
raise SyntaxError(msg) raise SyntaxError(msg)
palette: list[int] = [] palette: list[int] = []
for _ in range(256 + 3): i = 0
while True:
if limit and i == 256 + 3:
break
i += 1
s = fp.readline() s = fp.readline()
if not s: if not s:
break break
@ -38,7 +44,7 @@ class GimpPaletteFile:
# skip fields and comment lines # skip fields and comment lines
if re.match(rb"\w+:|#", s): if re.match(rb"\w+:|#", s):
continue continue
if len(s) > 100: if limit and len(s) > 100:
msg = "bad palette file" msg = "bad palette file"
raise SyntaxError(msg) raise SyntaxError(msg)
@ -48,10 +54,19 @@ class GimpPaletteFile:
raise ValueError(msg) raise ValueError(msg)
palette += (int(v[i]) for i in range(3)) palette += (int(v[i]) for i in range(3))
if len(palette) == 768: if limit and len(palette) == 768:
break break
self.palette = bytes(palette) self.palette = bytes(palette)
def __init__(self, fp: IO[bytes]) -> None:
self._read(fp)
@classmethod
def frombytes(cls, data: bytes) -> GimpPaletteFile:
self = cls.__new__(cls)
self._read(BytesIO(data), False)
return self
def getpalette(self) -> tuple[bytes, str]: def getpalette(self) -> tuple[bytes, str]:
return self.palette, self.rawmode return self.palette, self.rawmode