1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-23 06:48:16 +00:00

src/trezor/wire: allow passing MessageType classes to register()

This commit is contained in:
Jan Pochyla 2018-04-05 12:34:39 +02:00
parent 2731742d34
commit 60bec0b4d1

View File

@ -13,6 +13,8 @@ workflow_handlers = {}
def register(mtype, handler, *args): def register(mtype, handler, *args):
'''Register `handler` to get scheduled after `mtype` message is received.''' '''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: if mtype in workflow_handlers:
raise KeyError raise KeyError
workflow_handlers[mtype] = (handler, args) workflow_handlers[mtype] = (handler, args)