dispatch on wire_type instead of message_type

src/apps/homescreen/__init__.py

@@ -1,11 +1,17 @@
 from trezor.dispatcher import register
-from trezor.messages.Initialize import Initialize
+from trezor.utils import unimport_func
 
 
-def dispatch(message):
+@unimport_func
+def dispatch_Initialize(mtype, mbuf):
+    from trezor.messages.Initialize import Initialize
+
+    message = Initialize.loads(mbuf)
+
     from .layout_homescreen import layout_homescreen
     return layout_homescreen(message)
 
 
 def boot():
-    register(Initialize, dispatch)
+    Initialize = 0
+    register(Initialize, dispatch_Initialize)

src/apps/wallet/__init__.py

@@ -1,12 +1,17 @@
 from trezor.dispatcher import register
-from trezor.messages.GetPublicKey import GetPublicKey
+from trezor.utils import unimport_func
 
 
-def dispatch(message):
+@unimport_func
-    if message.message_type is GetPublicKey:
+def dispatch_GetPublicKey(mtype, mbuf):
-        from .layout_get_public_key import layout_get_public_key
+    from trezor.messages.GetPublicKey import GetPublicKey
-        return layout_get_public_key(message)
+
+    message = GetPublicKey.loads(mbuf)
+
+    from .layout_get_public_key import layout_get_public_key
+    return layout_get_public_key(message)
 
 
 def boot():
-    register(GetPublicKey, dispatch)
+    GetPublicKey = 11
+    register(GetPublicKey, dispatch_GetPublicKey)

src/trezor/dispatcher.py

@@ -5,16 +5,15 @@ from . import layout
 message_handlers = {}
 
 
-def register(message_type, handler):
+def register(mtype, handler):
-    message_handlers[message_type] = handler
+    message_handlers[mtype] = handler
 
 
-def unregister(message_type):
+def unregister(mtype):
-    del message_handlers[message_type]
+    del message_handlers[mtype]
 
 
 def dispatch():
-    mtypes = message_handlers.keys()
+    mtype, mbuf = yield from wire.read_wire_msg()
-    message = yield from wire.read_msg(*mtypes)
+    handler = message_handlers[mtype]
-    handler = message_handlers[message.message_type]
+    layout.change(handler(mtype, mbuf))
-    layout.change(handler(message))