diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/qemu/qmp/protocol.py | 35 |
1 files changed, 16 insertions, 19 deletions
diff --git a/python/qemu/qmp/protocol.py b/python/qemu/qmp/protocol.py index 208bdec5c8..958aeca08a 100644 --- a/python/qemu/qmp/protocol.py +++ b/python/qemu/qmp/protocol.py @@ -15,6 +15,7 @@ class. import asyncio from asyncio import StreamReader, StreamWriter +from contextlib import asynccontextmanager from enum import Enum from functools import wraps import logging @@ -22,6 +23,7 @@ import socket from ssl import SSLContext from typing import ( Any, + AsyncGenerator, Awaitable, Callable, Generic, @@ -337,9 +339,8 @@ class AsyncProtocol(Generic[T]): This exception will wrap a more concrete one. In most cases, the wrapped exception will be `OSError`. """ - await self._session_guard( - self._do_start_server(address, ssl), - 'Failed to establish connection') + async with self._session_guard('Failed to establish connection'): + await self._do_start_server(address, ssl) assert self.runstate == Runstate.CONNECTING @upper_half @@ -362,12 +363,10 @@ class AsyncProtocol(Generic[T]): """ if self._accepted is None: raise QMPError("Cannot call accept() before start_server().") - await self._session_guard( - self._do_accept(), - 'Failed to establish connection') - await self._session_guard( - self._establish_session(), - 'Failed to establish session') + async with self._session_guard('Failed to establish connection'): + await self._do_accept() + async with self._session_guard('Failed to establish session'): + await self._establish_session() assert self.runstate == Runstate.RUNNING @upper_half @@ -392,12 +391,10 @@ class AsyncProtocol(Generic[T]): protocol-level failure occurs while establishing a new session, the wrapped error may also be an `QMPError`. """ - await self._session_guard( - self._do_connect(address, ssl), - 'Failed to establish connection') - await self._session_guard( - self._establish_session(), - 'Failed to establish session') + async with self._session_guard('Failed to establish connection'): + await self._do_connect(address, ssl) + async with self._session_guard('Failed to establish session'): + await self._establish_session() assert self.runstate == Runstate.RUNNING @upper_half @@ -418,7 +415,8 @@ class AsyncProtocol(Generic[T]): # Section: Session machinery # -------------------------- - async def _session_guard(self, coro: Awaitable[None], emsg: str) -> None: + @asynccontextmanager + async def _session_guard(self, emsg: str) -> AsyncGenerator[None, None]: """ Async guard function used to roll back to `IDLE` on any error. @@ -435,10 +433,9 @@ class AsyncProtocol(Generic[T]): :raise ConnectError: When any other error is encountered in the guarded block. """ - # Note: After Python 3.6 support is removed, this should be an - # @asynccontextmanager instead of accepting a callback. try: - await coro + # Caller's code runs here. + yield except BaseException as err: self.logger.error("%s: %s", emsg, exception_summary(err)) self.logger.debug("%s:\n%s\n", emsg, pretty_traceback()) |