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

32 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-04-07 18:55 -0700

1import asyncio 

2from base64 import b64decode, b64encode 

3from http import HTTPStatus 

4from io import BytesIO 

5 

6import backoff 

7import httpx 

8from PIL.Image import Image 

9from PIL.Image import open as open_image 

10 

11MAX_TRIES: int = 3 

12MAX_TIMEOUT: int = 60 

13REQUEST_TIMEOUT: int = 30 

14 

15 

16@backoff.on_predicate( 

17 backoff.expo, 

18 max_tries=MAX_TRIES, 

19 max_time=MAX_TIMEOUT, 

20) 

21async def download_image_async( 

22 image_url: str, 

23 *, 

24 user_agent: str | None = None, 

25) -> Image | None: 

26 async with httpx.AsyncClient().stream( 

27 "GET", 

28 image_url, 

29 follow_redirects=True, 

30 timeout=REQUEST_TIMEOUT, 

31 headers=( 

32 { 

33 "User-Agent": user_agent, 

34 } if user_agent 

35 else {} 

36 ), 

37 ) as response: 

38 if response.status_code == HTTPStatus.TOO_MANY_REQUESTS: 

39 # It also triggers backoff if necessary 

40 return None 

41 

42 response.raise_for_status() 

43 

44 return bytes_to_image(await response.aread()) 

45 

46 

47def download_image( 

48 image_url: str, 

49 *, 

50 user_agent: str | None = None, 

51) -> Image | None: 

52 return asyncio.run(download_image_async( 

53 image_url, 

54 user_agent=user_agent, 

55 )) 

56 

57 

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

59 bio = BytesIO() 

60 image.save(bio, format=_format) 

61 return bio.getvalue() 

62 

63 

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

65 return open_image( 

66 BytesIO(b), 

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

68 ) 

69 

70 

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

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

73 

74 

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

76 """ 

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

78 """ 

79 

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

81 

82 

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

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