Skip to content

Commit 4d5a451

Browse files
refactor: use named constants for image format magic bytes
- Remove __future__ annotations import from input_model.py - Replace inline byte literals with named constants in image_format_detector.py - Calculate lengths dynamically using len() instead of hardcoded values - Eliminates need for noqa comments on magic value comparisons
1 parent e3ca6e1 commit 4d5a451

2 files changed

Lines changed: 39 additions & 15 deletions

File tree

src/resolver_athena_client/client/image_format_detector.py

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,16 @@
22

33
from resolver_athena_client.generated.athena.models_pb2 import ImageFormat
44

5+
PNG_MAGIC_BYTES = b"\x89PNG"
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

616
def 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"\x89PNG": # 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

src/resolver_athena_client/client/models/input_model.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55
consistent input handling across the application.
66
"""
77

8-
from __future__ import annotations
9-
108
import hashlib
119
from typing import TYPE_CHECKING
1210

0 commit comments

Comments
 (0)