diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/qemu/qmp/protocol.py | 53 |
1 files changed, 32 insertions, 21 deletions
diff --git a/python/qemu/qmp/protocol.py b/python/qemu/qmp/protocol.py index 3d5eb553aa..4d8a39f014 100644 --- a/python/qemu/qmp/protocol.py +++ b/python/qemu/qmp/protocol.py @@ -18,6 +18,7 @@ from asyncio import StreamReader, StreamWriter from contextlib import asynccontextmanager from enum import Enum from functools import wraps +from inspect import iscoroutinefunction import logging import socket from ssl import SSLContext @@ -130,6 +131,25 @@ def require(required_state: Runstate) -> Callable[[F], F]: :param required_state: The `Runstate` required to invoke this method. :raise StateError: When the required `Runstate` is not met. """ + def _check(proto: 'AsyncProtocol[Any]') -> None: + name = type(proto).__name__ + if proto.runstate == required_state: + return + + if proto.runstate == Runstate.CONNECTING: + emsg = f"{name} is currently connecting." + elif proto.runstate == Runstate.DISCONNECTING: + emsg = (f"{name} is disconnecting." + " Call disconnect() to return to IDLE state.") + elif proto.runstate == Runstate.RUNNING: + emsg = f"{name} is already connected and running." + elif proto.runstate == Runstate.IDLE: + emsg = f"{name} is disconnected and idle." + else: + assert False + + raise StateError(emsg, proto.runstate, required_state) + def _decorator(func: F) -> F: # _decorator is the decorator that is built by calling the # require() decorator factory; e.g.: @@ -140,29 +160,20 @@ def require(required_state: Runstate) -> Callable[[F], F]: @wraps(func) def _wrapper(proto: 'AsyncProtocol[Any]', *args: Any, **kwargs: Any) -> Any: - # _wrapper is the function that gets executed prior to the - # decorated method. - - name = type(proto).__name__ - - if proto.runstate != required_state: - if proto.runstate == Runstate.CONNECTING: - emsg = f"{name} is currently connecting." - elif proto.runstate == Runstate.DISCONNECTING: - emsg = (f"{name} is disconnecting." - " Call disconnect() to return to IDLE state.") - elif proto.runstate == Runstate.RUNNING: - emsg = f"{name} is already connected and running." - elif proto.runstate == Runstate.IDLE: - emsg = f"{name} is disconnected and idle." - else: - assert False - raise StateError(emsg, proto.runstate, required_state) - # No StateError, so call the wrapped method. + _check(proto) return func(proto, *args, **kwargs) - # Return the decorated method; - # Transforming Func to Decorated[Func]. + @wraps(func) + async def _async_wrapper(proto: 'AsyncProtocol[Any]', + *args: Any, **kwargs: Any) -> Any: + _check(proto) + return await func(proto, *args, **kwargs) + + # Return the decorated method; F => Decorated[F] + # Use an async version when applicable, which + # preserves async signature generation in sphinx. + if iscoroutinefunction(func): + return cast(F, _async_wrapper) return cast(F, _wrapper) # Return the decorator instance from the decorator factory. Phew! |