diff --git a/mystbin/client.py b/mystbin/client.py index 44b9ddb..91371a0 100755 --- a/mystbin/client.py +++ b/mystbin/client.py @@ -23,7 +23,8 @@ from __future__ import annotations import datetime -from typing import List, Literal, Optional, Sequence, Union, overload +from types import TracebackType +from typing import List, Literal, Optional, Sequence, Type, Union, overload import aiohttp @@ -41,6 +42,17 @@ class Client: def __init__(self, *, token: Optional[str] = None, session: Optional[aiohttp.ClientSession] = None) -> None: self.http: HTTPClient = HTTPClient(token=token, session=session) + async def __aenter__(self): + return self + + async def __aexit__( + self, + exc_cls: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType] + ): + await self.close() + async def close(self) -> None: """|coro|