From 29863f84ca83038c24ccf779ce75eb14fe5c07d8 Mon Sep 17 00:00:00 2001 From: matejcik Date: Wed, 26 Jul 2023 11:10:30 +0200 Subject: [PATCH] refactor(core): simplify context.wait() to only accept one task Nobody was using the multi-task version and this is easier to properly type-annotate. --- core/src/trezor/wire/context.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/core/src/trezor/wire/context.py b/core/src/trezor/wire/context.py index 4dc151397c..08eaab3474 100644 --- a/core/src/trezor/wire/context.py +++ b/core/src/trezor/wire/context.py @@ -38,6 +38,8 @@ if TYPE_CHECKING: LoadedMessageType = TypeVar("LoadedMessageType", bound=protobuf.MessageType) + T = TypeVar("T") + class UnexpectedMessage(Exception): """A message was received that is not part of the current workflow. @@ -159,19 +161,19 @@ class Context: CURRENT_CONTEXT: Context | None = None -def wait(*tasks: Awaitable) -> Any: +def wait(task: Awaitable[T]) -> Awaitable[T]: """ - Wait until one of the passed tasks finishes, and return the result, while servicing - the wire context. + Wait until the passed in task finishes, and return the result, while servicing the + wire context. Used to make sure the device is responsive on USB while waiting for user - interaction. If a message is received before any of the passed in tasks finish, it - raises an `UnexpectedMessage` exception, returning control to the session handler. + interaction. If a message is received before the task finishes, it raises an + `UnexpectedMessage` exception, returning control to the session handler. """ if CURRENT_CONTEXT is None: - return loop.race(*tasks) + return task else: - return loop.race(CURRENT_CONTEXT.read(()), *tasks) + return loop.race(CURRENT_CONTEXT.read(()), task) async def call(