diff --git a/examples/image_bot.py b/examples/image_bot.py index 3e45c39..48816c8 100644 --- a/examples/image_bot.py +++ b/examples/image_bot.py @@ -52,9 +52,6 @@ async def on_image(ctx: EventContext) -> None: s = time.time() data = await ctx.bot.download_file(ctx) took_time = time.time() - s - if not isinstance(data, bytes): - await ctx.bot.send_text(ctx.room, "😧") - return # convert to Image with BytesIO(data) as buf: img = Image.open(buf) diff --git a/src/mab/bot/bot.py b/src/mab/bot/bot.py index 3a79e29..9649898 100644 --- a/src/mab/bot/bot.py +++ b/src/mab/bot/bot.py @@ -1,7 +1,7 @@ import asyncio import logging -from typing import Callable, Coroutine, Any +from typing import Callable, Coroutine, Any, overload from nio import AsyncClient, MatrixRoom, Event @@ -251,6 +251,18 @@ class MatrixBot: is_html=is_html, timeout=timeout ) + + @overload + async def download_file(self, + source: EventContext | Event, + *, + path: str | Path) -> Path: ... + + @overload + async def download_file(self, + source: EventContext | Event, + *, + path: None = None) -> bytes: ... async def download_file(self, source: EventContext | Event,