Coverage for src/extratools_image/__init__.py: 71%

52 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-04-28 20:29 -0700

1from __future__ import annotations 

2 

3import asyncio 

4from base64 import b64decode, b64encode 

5from enum import Enum 

6from http import HTTPStatus 

7from io import BytesIO 

8 

9import backoff 

10import httpx 

11import truststore 

12from PIL.Image import Image 

13from PIL.Image import open as open_image 

14 

15truststore.inject_into_ssl() 

16 

17MAX_TRIES: int = 3 

18MAX_TIMEOUT: int = 60 

19REQUEST_TIMEOUT: int = 30 

20 

21 

22@backoff.on_predicate( 

23 backoff.expo, 

24 max_tries=MAX_TRIES, 

25 max_time=MAX_TIMEOUT, 

26) 

27async def download_image_async( 

28 image_url: str, 

29 *, 

30 user_agent: str | None = None, 

31) -> Image | None: 

32 async with httpx.AsyncClient().stream( 

33 "GET", 

34 image_url, 

35 follow_redirects=True, 

36 timeout=REQUEST_TIMEOUT, 

37 headers=( 

38 { 

39 "User-Agent": user_agent, 

40 } if user_agent 

41 else {} 

42 ), 

43 ) as response: 

44 if response.status_code == HTTPStatus.TOO_MANY_REQUESTS: 

45 # It also triggers backoff if necessary 

46 return None 

47 

48 response.raise_for_status() 

49 

50 return bytes_to_image(await response.aread()) 

51 

52 

53def download_image( 

54 image_url: str, 

55 *, 

56 user_agent: str | None = None, 

57) -> Image | None: 

58 return asyncio.run(download_image_async( 

59 image_url, 

60 user_agent=user_agent, 

61 )) 

62 

63 

64def image_to_bytes(image: Image, _format: str = "PNG") -> bytes: 

65 bio = BytesIO() 

66 image.save(bio, format=_format) 

67 return bio.getvalue() 

68 

69 

70def bytes_to_image(b: bytes, _format: str | None = None) -> Image: 

71 return open_image( 

72 BytesIO(b), 

73 formats=((_format,) if _format else None), 

74 ) 

75 

76 

77def image_to_base64_str(image: Image, _format: str = "PNG") -> str: 

78 return b64encode(image_to_bytes(image, _format)).decode() 

79 

80 

81def image_to_data_url(image: Image, _format: str = "PNG") -> str: 

82 """ 

83 Following https://developer.mozilla.org/en-US/docs/Web/URI/Reference/Schemes/data 

84 """ 

85 

86 return f"data:image/{_format.lower()};base64,{image_to_base64_str(image)}" 

87 

88 

89def base64_str_to_image(s: str) -> Image: 

90 return open_image(b64decode(s.encode())) 

91 

92 

93class CommonSize(Enum): 

94 """ 

95 https://en.wikipedia.org/wiki/Display_resolution_standards 

96 """ 

97 

98 # 4:3 

99 VGA = (640, 480) 

100 SVGA = (800, 600) 

101 XGA = (1024, 768) 

102 SXGA = (1280, 1024) 

103 

104 # 16:9 

105 _480p = SD = (854, 480) 

106 _720p = HD = (1280, 720) 

107 _2K = _1080p = FULL_HD = (1920, 1080) 

108 _3K = (2880, 1620) 

109 _4K = (3840, 2160) 

110 _5K = (5120, 2880) 

111 _8K = (7680, 4320) 

112 _16K = (15360, 8640) 

113 

114 def __repr__(self) -> str: 

115 x, y = self.value 

116 return f"{self.name.rstrip('_').replace('_', '')} ({x}x{y})"