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
« 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
6import backoff
7import httpx
8from PIL.Image import Image
9from PIL.Image import open as open_image
11MAX_TRIES: int = 3
12MAX_TIMEOUT: int = 60
13REQUEST_TIMEOUT: int = 30
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
42 response.raise_for_status()
44 return bytes_to_image(await response.aread())
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 ))
58def image_to_bytes(image: Image, _format: str = "PNG") -> bytes:
59 bio = BytesIO()
60 image.save(bio, format=_format)
61 return bio.getvalue()
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 )
71def image_to_base64_str(image: Image, _format: str = "PNG") -> str:
72 return b64encode(image_to_bytes(image, _format)).decode()
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 """
80 return f"data:image/{_format.lower()};base64,{image_to_base64_str(image)}"
83def base64_str_to_image(s: str) -> Image:
84 return open_image(b64decode(s.encode()))