From 60bec0b4d1887579659bf1a31f24cb3cbf89cb34 Mon Sep 17 00:00:00 2001 From: Jan Pochyla Date: Thu, 5 Apr 2018 12:34:39 +0200 Subject: [PATCH] src/trezor/wire: allow passing MessageType classes to register() --- src/trezor/wire/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/trezor/wire/__init__.py b/src/trezor/wire/__init__.py index 3f9f441c66..504c905b65 100644 --- a/src/trezor/wire/__init__.py +++ b/src/trezor/wire/__init__.py @@ -13,6 +13,8 @@ workflow_handlers = {} def register(mtype, handler, *args): '''Register `handler` to get scheduled after `mtype` message is received.''' + if isinstance(mtype, type) and issubclass(mtype, protobuf.MessageType): + mtype = mtype.MESSAGE_WIRE_TYPE if mtype in workflow_handlers: raise KeyError workflow_handlers[mtype] = (handler, args)