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

53 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-05-13 19:46 -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 bio = BytesIO() 

68 image.save(bio, format=_format) 

69 return bio.getvalue() 

70 

71 

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

73 return open_image( 

74 BytesIO(b), 

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

76 ) 

77 

78 

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

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

81 

82 

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

84 """ 

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

86 """ 

87 

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

89 

90 

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

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

93 

94 

95class CommonSize(Enum): 

96 """ 

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

98 """ 

99 

100 # 4:3 

101 VGA = (640, 480) 

102 SVGA = (800, 600) 

103 XGA = (1024, 768) 

104 SXGA = (1280, 1024) 

105 

106 # 16:9 

107 _480p = SD = (854, 480) 

108 _720p = HD = (1280, 720) 

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

110 _3K = (2880, 1620) 

111 _4K = (3840, 2160) 

112 _5K = (5120, 2880) 

113 _8K = (7680, 4320) 

114 _16K = (15360, 8640) 

115 

116 def __repr__(self) -> str: 

117 x, y = self.value 

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