Merge pull request #3441 from trailofbits/perf/lazy-pil-import

perf: lazy import PIL in nes.py and jpeg.py

Author: ESultanik

polyfile/jpeg.py

@@ -4,11 +4,16 @@
 from .fileutils import FileStream, Tempfile
 from .polyfile import Match, register_parser, Submatch
 
-from PIL import Image
+
+def _get_pil_image():
+    """Lazy import PIL.Image only when needed (for JPEG2000 parsing)."""
+    from PIL import Image
+    return Image
 
 
 @register_parser("image/jp2")
 def parse_jpeg2000(file_stream: FileStream, parent: Match):
+    Image = _get_pil_image()
     with Tempfile(file_stream.read(parent.length)) as input_bytes:
         img = Image.open(input_bytes)
         with BytesIO() as img_data:

polyfile/nes.py

@@ -1,10 +1,18 @@
 import base64
 from io import BytesIO
-
-from PIL import Image, ImageDraw
+from typing import TYPE_CHECKING
 
 from .polyfile import register_parser, InvalidMatch, Submatch
 
+if TYPE_CHECKING:
+    from PIL import Image as PILImage
+
+
+def _get_pil():
+    """Lazy import PIL only when needed (for NES ROM CHR bank rendering)."""
+    from PIL import Image, ImageDraw
+    return Image, ImageDraw
+
 
 def parse_ines_header(header, parent=None):
     magic = header[:4]
@@ -130,7 +138,8 @@ def chr_values(chr_bytes: bytes):
                       ((((chr_bytes[offset + y + 8] >> shift) & 0b1)) << 1) | ((chr_bytes[offset + y] >> shift) & 0b1)
 
 
-def render_chr(chr_bytes: bytes) -> Image:
+def render_chr(chr_bytes: bytes) -> "PILImage":
+    Image, ImageDraw = _get_pil()
     img = Image.new(mode='L', size=(8*16, 8*32))
     d = ImageDraw.Draw(img)
     for x, y, pixel in chr_values(chr_bytes):