diff options
Diffstat (limited to 'lib/bb/asyncrpc/client.py')
-rw-r--r-- | lib/bb/asyncrpc/client.py | 13 |
1 files changed, 13 insertions, 0 deletions
diff --git a/lib/bb/asyncrpc/client.py b/lib/bb/asyncrpc/client.py index fa042bbe8..dcbe7e576 100644 --- a/lib/bb/asyncrpc/client.py +++ b/lib/bb/asyncrpc/client.py @@ -126,6 +126,12 @@ class AsyncClient(object): {'ping': {}} ) + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.close() + class Client(object): def __init__(self): @@ -176,3 +182,10 @@ class Client(object): if sys.version_info >= (3, 6): self.loop.run_until_complete(self.loop.shutdown_asyncgens()) self.loop.close() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + return False |