From 86257a70e3bd1c345056529e5b9fc28a8b0881b8 Mon Sep 17 00:00:00 2001 From: nikita Date: Sun, 13 Sep 2026 03:39:23 +0300 Subject: [PATCH] Added overloads for MatrixBot.download_file --- examples/image_bot.py | 3 --- src/mab/bot/bot.py | 14 +++++++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) 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,