From 4729978c60ecf837629e8698f6b799dd6f150a09 Mon Sep 17 00:00:00 2001
From: Alex Root Junior <jroot.junior@gmail.com>
Date: Mon, 22 Apr 2024 23:42:47 +0300
Subject: [PATCH] Add context manager support for bot client (#1468)

* Add context manager support for bot client

The bot client now supports the context manager protocol, providing automatic resource management. This enhancement helps to automatically close the session when leaving the context, which cleans up resources better. The documentation and tests have been updated accordingly to illustrate this new feature. Moreover, an example of usage without a dispatcher has been provided to clarify its use in simple cases.

* Added changelog
---
 CHANGES/1468.feature.rst               | 13 ++++++++++
 aiogram/client/bot.py                  | 13 ++++++++++
 docs/index.rst                         |  9 +++++++
 examples/without_dispatcher.py         | 36 ++++++++++++++++++++++++++
 tests/test_api/test_client/test_bot.py | 13 ++++++++++
 5 files changed, 84 insertions(+)
 create mode 100644 CHANGES/1468.feature.rst
 create mode 100644 examples/without_dispatcher.py

diff --git a/CHANGES/1468.feature.rst b/CHANGES/1468.feature.rst
new file mode 100644
index 00000000..de78716c
--- /dev/null
+++ b/CHANGES/1468.feature.rst
@@ -0,0 +1,13 @@
+Added context manager interface to Bot instance, from now you can use:
+
+.. code-block:: python
+
+    async with Bot(...) as bot:
+        ...
+
+instead of
+
+.. code-block:: python
+
+    async with Bot(...).context():
+        ...
diff --git a/aiogram/client/bot.py b/aiogram/client/bot.py
index bf556d1d..52ba326a 100644
--- a/aiogram/client/bot.py
+++ b/aiogram/client/bot.py
@@ -5,6 +5,7 @@ import io
 import pathlib
 import warnings
 from contextlib import asynccontextmanager
+from types import TracebackType
 from typing import (
     Any,
     AsyncGenerator,
@@ -12,6 +13,7 @@ from typing import (
     BinaryIO,
     List,
     Optional,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -300,6 +302,17 @@ class Bot:
         self.__token = token
         self._me: Optional[User] = None
 
+    async def __aenter__(self) -> "Bot":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        await self.session.close()
+
     @property
     def parse_mode(self) -> Optional[str]:
         warnings.warn(
diff --git a/docs/index.rst b/docs/index.rst
index 6be454e7..a923cab1 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -5,6 +5,15 @@ Simple usage
 
 .. literalinclude:: ../examples/echo_bot.py
 
+
+Usage without dispatcher
+------------------------
+
+Just only interact with Bot API, without handling events
+
+.. literalinclude:: ../examples/without_dispatcher.py
+
+
 Contents
 ========
 
diff --git a/examples/without_dispatcher.py b/examples/without_dispatcher.py
new file mode 100644
index 00000000..87e1d8e6
--- /dev/null
+++ b/examples/without_dispatcher.py
@@ -0,0 +1,36 @@
+import asyncio
+from argparse import ArgumentParser
+
+from aiogram import Bot
+from aiogram.client.default import DefaultBotProperties
+from aiogram.enums import ParseMode
+
+
+def create_parser() -> ArgumentParser:
+    parser = ArgumentParser()
+    parser.add_argument("--token", help="Telegram Bot API Token")
+    parser.add_argument("--chat-id", type=int, help="Target chat id")
+    parser.add_argument("--message", "-m", help="Message text to sent", default="Hello, World!")
+
+    return parser
+
+
+async def main():
+    parser = create_parser()
+    ns = parser.parse_args()
+
+    token = ns.token
+    chat_id = ns.chat_id
+    message = ns.message
+
+    async with Bot(
+        token=token,
+        default=DefaultBotProperties(
+            parse_mode=ParseMode.HTML,
+        ),
+    ) as bot:
+        await bot.send_message(chat_id=chat_id, text=message)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/tests/test_api/test_client/test_bot.py b/tests/test_api/test_client/test_bot.py
index 7a836b04..5db1845a 100644
--- a/tests/test_api/test_client/test_bot.py
+++ b/tests/test_api/test_client/test_bot.py
@@ -14,6 +14,7 @@ from aiogram.methods import GetFile, GetMe
 from aiogram.types import File, PhotoSize
 from tests.deprecated import check_deprecated
 from tests.mocked_bot import MockedBot
+from tests.test_api.test_client.test_session.test_base_session import CustomSession
 
 
 @pytest.fixture()
@@ -42,6 +43,18 @@ class TestBot:
         assert isinstance(bot.session, AiohttpSession)
         assert bot.id == 42
 
+    async def test_bot_context_manager_over_session(self):
+        session = CustomSession()
+        with patch(
+            "tests.test_api.test_client.test_session.test_base_session.CustomSession.close",
+            new_callable=AsyncMock,
+        ) as mocked_close:
+            async with Bot(token="42:TEST", session=session) as bot:
+                assert bot.id == 42
+                assert bot.session is session
+
+            mocked_close.assert_awaited_once()
+
     def test_init_default(self):
         with check_deprecated(
             max_version="3.7.0",