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

50 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-04-25 17:52 -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 

11from PIL.Image import Image 

12from PIL.Image import open as open_image 

13 

14MAX_TRIES: int = 3 

15MAX_TIMEOUT: int = 60 

16REQUEST_TIMEOUT: int = 30 

17 

18 

19@backoff.on_predicate( 

20 backoff.expo, 

21 max_tries=MAX_TRIES, 

22 max_time=MAX_TIMEOUT, 

23) 

24async def download_image_async( 

25 image_url: str, 

26 *, 

27 user_agent: str | None = None, 

28) -> Image | None: 

29 async with httpx.AsyncClient().stream( 

30 "GET", 

31 image_url, 

32 follow_redirects=True, 

33 timeout=REQUEST_TIMEOUT, 

34 headers=( 

35 { 

36 "User-Agent": user_agent, 

37 } if user_agent 

38 else {} 

39 ), 

40 ) as response: 

41 if response.status_code == HTTPStatus.TOO_MANY_REQUESTS: 

42 # It also triggers backoff if necessary 

43 return None 

44 

45 response.raise_for_status() 

46 

47 return bytes_to_image(await response.aread()) 

48 

49 

50def download_image( 

51 image_url: str, 

52 *, 

53 user_agent: str | None = None, 

54) -> Image | None: 

55 return asyncio.run(download_image_async( 

56 image_url, 

57 user_agent=user_agent, 

58 )) 

59 

60 

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

62 bio = BytesIO() 

63 image.save(bio, format=_format) 

64 return bio.getvalue() 

65 

66 

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

68 return open_image( 

69 BytesIO(b), 

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

71 ) 

72 

73 

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

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

76 

77 

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

79 """ 

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

81 """ 

82 

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

84 

85 

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

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

88 

89 

90class CommonSize(Enum): 

91 """ 

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

93 """ 

94 

95 # 4:3 

96 VGA = (640, 480) 

97 SVGA = (800, 600) 

98 XGA = (1024, 768) 

99 SXGA = (1280, 1024) 

100 

101 # 16:9 

102 _480p = SD = (854, 480) 

103 _720p = HD = (1280, 720) 

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

105 _3K = (2880, 1620) 

106 _4K = (3840, 2160) 

107 _5K = (5120, 2880) 

108 _8K = (7680, 4320) 

109 _16K = (15360, 8640) 

110 

111 def __repr__(self) -> str: 

112 x, y = self.value 

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