22
33from resolver_athena_client .generated .athena .models_pb2 import ImageFormat
44
5+ PNG_MAGIC_BYTES = b"\x89 PNG"
6+ JPEG_MAGIC_BYTES = b"\xFF \xD8 \xFF "
7+ GIF87A_MAGIC_BYTES = b"GIF87a"
8+ GIF89A_MAGIC_BYTES = b"GIF89a"
9+ BMP_MAGIC_BYTES = b"BM"
10+ WEBP_RIFF_MAGIC_BYTES = b"RIFF"
11+ WEBP_WEBP_MAGIC_BYTES = b"WEBP"
12+ TIFF_LE_MAGIC_BYTES = b"II*\x00 "
13+ TIFF_BE_MAGIC_BYTES = b"MM\x00 *"
14+
515
616def detect_image_format (data : bytes ) -> ImageFormat .ValueType : # noqa: PLR0911
717 """Detect image format from raw bytes using magic number signatures.
@@ -19,29 +29,45 @@ def detect_image_format(data: bytes) -> ImageFormat.ValueType: # noqa: PLR0911
1929 return ImageFormat .IMAGE_FORMAT_UNSPECIFIED
2030
2131 # Check magic numbers for common image formats
22- # PNG: starts with \x89PNG (need at least 4 bytes)
23- if len (data ) >= 4 and data [:4 ] == b"\x89 PNG" : # noqa: PLR2004
32+ # PNG: starts with PNG_MAGIC_BYTES
33+ png_len = len (PNG_MAGIC_BYTES )
34+ if len (data ) >= png_len and data [:png_len ] == PNG_MAGIC_BYTES :
2435 return ImageFormat .IMAGE_FORMAT_PNG
2536
26- # JPEG: starts with \xFF\xD8\xFF (need at least 3 bytes)
27- if len (data ) >= 3 and data [:3 ] == b"\xff \xd8 \xff " : # noqa: PLR2004
37+ # JPEG: starts with JPEG_MAGIC_BYTES
38+ jpeg_len = len (JPEG_MAGIC_BYTES )
39+ if len (data ) >= jpeg_len and data [:jpeg_len ] == JPEG_MAGIC_BYTES :
2840 return ImageFormat .IMAGE_FORMAT_JPEG
2941
30- # GIF: starts with GIF87a or GIF89a (need at least 6 bytes)
31- if len (data ) >= 6 and data [:6 ] in (b"GIF87a" , b"GIF89a" ): # noqa: PLR2004
42+ # GIF: starts with GIF87A_MAGIC_BYTES or GIF89A_MAGIC_BYTES
43+ gif_len = len (GIF87A_MAGIC_BYTES )
44+ if len (data ) >= gif_len and data [:gif_len ] in (
45+ GIF87A_MAGIC_BYTES ,
46+ GIF89A_MAGIC_BYTES ,
47+ ):
3248 return ImageFormat .IMAGE_FORMAT_GIF
3349
34- # BMP: starts with BM (need at least 2 bytes)
35- if len (data ) >= 2 and data [:2 ] == b"BM" : # noqa: PLR2004
50+ # BMP: starts with BMP_MAGIC_BYTES
51+ bmp_len = len (BMP_MAGIC_BYTES )
52+ if len (data ) >= bmp_len and data [:bmp_len ] == BMP_MAGIC_BYTES :
3653 return ImageFormat .IMAGE_FORMAT_BMP
3754
38- # WebP: starts with RIFF....WEBP (need at least 12 bytes)
39- if len (data ) >= 12 and data [:4 ] == b"RIFF" and data [8 :12 ] == b"WEBP" : # noqa: PLR2004
55+ # WebP: RIFF....WEBP (12 bytes minimum for full signature)
56+ webp_min_len = len (WEBP_RIFF_MAGIC_BYTES ) + len (WEBP_WEBP_MAGIC_BYTES ) + 4
57+ if (
58+ len (data ) >= webp_min_len
59+ and data [:4 ] == WEBP_RIFF_MAGIC_BYTES
60+ and data [8 :12 ] == WEBP_WEBP_MAGIC_BYTES
61+ ):
4062 return ImageFormat .IMAGE_FORMAT_WEBP
4163
42- # TIFF: starts with II* or MM* (need at least 4 bytes)
43- if len (data ) >= 4 and data [:4 ] in (b"II*\x00 " , b"MM\x00 *" ): # noqa: PLR2004
64+ # TIFF: little-endian or big-endian magic bytes
65+ tiff_len = len (TIFF_LE_MAGIC_BYTES )
66+ if len (data ) >= tiff_len and (
67+ data [:tiff_len ] == TIFF_LE_MAGIC_BYTES
68+ or data [:tiff_len ] == TIFF_BE_MAGIC_BYTES
69+ ):
4470 return ImageFormat .IMAGE_FORMAT_TIFF
4571
46- # If we can't detect the format, return UNSPECIFIED
72+ # Fallback when format cannot be determined
4773 return ImageFormat .IMAGE_FORMAT_UNSPECIFIED
0 commit comments