| 1 | from __future__ import annotations |
| 2 | |
| 3 | import base64 |
| 4 | from io import BytesIO |
| 5 | from typing import Any |
| 6 | |
| 7 | from PIL import Image |
| 8 | |
| 9 | |
| 10 | def image_from_response_part(part: Any) -> Image.Image | None: |
| 11 | inline_data = getattr(part, "inline_data", None) |
| 12 | if inline_data is None and isinstance(part, dict): |
| 13 | inline_data = part.get("inline_data") |
| 14 | if inline_data is None: |
| 15 | return None |
| 16 | |
| 17 | as_image = getattr(part, "as_image", None) |
| 18 | if callable(as_image): |
| 19 | image = as_image() |
| 20 | if isinstance(image, Image.Image): |
| 21 | return image |
| 22 | |
| 23 | data = _value(inline_data, "data") |
| 24 | if data is None: |
| 25 | return None |
| 26 | if isinstance(data, str): |
| 27 | if data.startswith("data:") and "," in data: |
| 28 | data = data.split(",", 1)[1] |
| 29 | data = base64.b64decode(data) |
| 30 | if isinstance(data, bytearray): |
| 31 | data = bytes(data) |
| 32 | if not isinstance(data, bytes): |
| 33 | return None |
| 34 | return Image.open(BytesIO(data)).convert("RGB") |
| 35 | |
| 36 | |
| 37 | def _value(obj: Any, key: str) -> Any: |
| 38 | if isinstance(obj, dict): |
| 39 | return obj.get(key) |
| 40 | return getattr(obj, key, None) |
| 41 |