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

60 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-05-15 20:19 -0700

1from __future__ import annotations 

2 

3import asyncio 

4import ssl 

5from base64 import b64decode, b64encode 

6from enum import Enum 

7from http import HTTPStatus 

8from io import BytesIO 

9 

10import backoff 

11import httpx 

12import truststore 

13from PIL.Image import Image 

14from PIL.Image import open as open_image 

15 

16MAX_TRIES: int = 3 

17MAX_TIMEOUT: int = 60 

18REQUEST_TIMEOUT: int = 30 

19 

20ctx = truststore.SSLContext(ssl.PROTOCOL_TLS_CLIENT) 

21 

22 

23@backoff.on_predicate( 

24 backoff.expo, 

25 max_tries=MAX_TRIES, 

26 max_time=MAX_TIMEOUT, 

27) 

28async def download_image_async( 

29 image_url: str, 

30 *, 

31 user_agent: str | None = None, 

32) -> Image | None: 

33 # https://www.python-httpx.org/advanced/ssl/ 

34 async with httpx.AsyncClient(verify=ctx).stream( 

35 "GET", 

36 image_url, 

37 follow_redirects=True, 

38 timeout=REQUEST_TIMEOUT, 

39 headers=( 

40 { 

41 "User-Agent": user_agent, 

42 } if user_agent 

43 else {} 

44 ), 

45 ) as response: 

46 if response.status_code == HTTPStatus.TOO_MANY_REQUESTS: 

47 # It also triggers backoff if necessary 

48 return None 

49 

50 response.raise_for_status() 

51 

52 return bytes_to_image(await response.aread()) 

53 

54 

55def download_image( 

56 image_url: str, 

57 *, 

58 user_agent: str | None = None, 

59) -> Image | None: 

60 return asyncio.run(download_image_async( 

61 image_url, 

62 user_agent=user_agent, 

63 )) 

64 

65 

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

67 _format = _format.upper() 

68 if _format == "JPG": 

69 _format = "JPEG" 

70 

71 bio = BytesIO() 

72 image.save(bio, format=_format) 

73 return bio.getvalue() 

74 

75 

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

77 if _format: 

78 _format = _format.upper() 

79 if _format == "JPG": 

80 _format = "JPEG" 

81 

82 return open_image( 

83 BytesIO(b), 

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

85 ) 

86 

87 

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

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

90 

91 

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

93 """ 

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

95 """ 

96 

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

98 

99 

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

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

102 

103 

104class CommonSize(Enum): 

105 """ 

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

107 """ 

108 

109 # 4:3 

110 VGA = (640, 480) 

111 SVGA = (800, 600) 

112 XGA = (1024, 768) 

113 SXGA = (1280, 1024) 

114 

115 # 16:9 

116 _480p = SD = (854, 480) 

117 _720p = HD = (1280, 720) 

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

119 _3K = (2880, 1620) 

120 _4K = (3840, 2160) 

121 _5K = (5120, 2880) 

122 _8K = (7680, 4320) 

123 _16K = (15360, 8640) 

124 

125 def __repr__(self) -> str: 

126 x, y = self.value 

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