From 883b392fdb5c3bd76396ba184962fb536e4b7b30 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Mon, 26 Sep 2016 16:23:48 +0200
Subject: [PATCH] adopt wire dispatching, log removed modules

---
 src/trezor/messages/__init__.py    |  2 +-
 src/trezor/utils.py                | 21 +++++++++++----------
 src/trezor/wire/wire_dispatcher.py | 22 +++++++++++-----------
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/src/trezor/messages/__init__.py b/src/trezor/messages/__init__.py
index 10d61a35b..91109bfff 100644
--- a/src/trezor/messages/__init__.py
+++ b/src/trezor/messages/__init__.py
@@ -9,5 +9,5 @@ def get_protobuf_type_name(wire_type):
 
 def get_protobuf_type(wire_type):
     name = get_protobuf_type_name(wire_type)
-    module = __import__('.%s' % name, globals(), locals(), (name,), 1)
+    module = __import__('.%s' % name, None, None, (name,), 1)
     return getattr(module, name)
diff --git a/src/trezor/utils.py b/src/trezor/utils.py
index b07d9446e..ac4152582 100644
--- a/src/trezor/utils.py
+++ b/src/trezor/utils.py
@@ -1,4 +1,7 @@
 import sys
+import gc
+
+from trezor import log
 
 type_gen = type((lambda: (yield))())
 type_genfunc = type((lambda: (yield)))
@@ -11,19 +14,25 @@ def _unimport_func(func):
             ret = func(*args, **kwargs)
         finally:
             for to_remove in set(sys.modules) - mods:
+                print('removing module', to_remove)
+                # log.info(__name__, 'removing module %s', to_remove)
                 del sys.modules[to_remove]
+            gc.collect()
         return ret
     return inner
 
 
 def _unimport_genfunc(genfunc):
-    def inner(*args, **kwargs):
+    async def inner(*args, **kwargs):
         mods = set(sys.modules)
         try:
-            ret = yield from genfunc(*args, **kwargs)
+            ret = await genfunc(*args, **kwargs)
         finally:
             for to_remove in set(sys.modules) - mods:
+                print('removing module', to_remove)
+                # log.info(__name__, 'removing module %s', to_remove)
                 del sys.modules[to_remove]
+            gc.collect()
         return ret
     return inner
 
@@ -35,14 +44,6 @@ def unimport(func):
         return _unimport_func(func)
 
 
-def coroutine(func):
-    def inner(*args, **kwargs):
-        gen = func(*args, **kwargs)
-        gen.send(None)
-        return gen
-    return inner
-
-
 def chunks(l, n):
     for i in range(0, len(l), n):
         yield l[i:i + n]
diff --git a/src/trezor/wire/wire_dispatcher.py b/src/trezor/wire/wire_dispatcher.py
index 729c555dc..554e06a3a 100644
--- a/src/trezor/wire/wire_dispatcher.py
+++ b/src/trezor/wire/wire_dispatcher.py
@@ -3,9 +3,9 @@ from .wire_codec import parse_report, REP_MARKER_OPEN, REP_MARKER_CLOSE
 
 
 def dispatch_reports_by_session(handlers,
-                                open_handler,
-                                close_handler,
-                                fallback_handler):
+                                open_callback,
+                                close_callback,
+                                unknown_callback):
     '''
     Consumes reports adhering to the wire codec and dispatches the report
     payloads by between the passed handlers.
@@ -16,21 +16,21 @@ def dispatch_reports_by_session(handlers,
 
         if marker == REP_MARKER_OPEN:
             log.debug(__name__, 'request for new session')
-            open_handler.send(session_id)
+            open_callback()
             continue
 
         elif marker == REP_MARKER_CLOSE:
             log.debug(__name__, 'request for closing session %d', session_id)
-            close_handler.send(session_id)
+            close_callback(session_id)
             continue
 
-        elif session_id in handlers:
-            log.debug(__name__, 'report on session %d', session_id)
-            handler = handlers[session_id]
-
-        else:
+        elif session_id not in handlers:
             log.debug(__name__, 'report on unknown session %d', session_id)
-            handler = fallback_handler
+            unknown_callback(session_id, report_data)
+            continue
+
+        log.debug(__name__, 'report on session %d', session_id)
+        handler = handlers[session_id]
 
         try:
             handler.send(report_data)