From 3db1bf89fae9d4481812f04e404e4780f46713b5 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Wed, 16 Aug 2017 15:02:03 +0200
Subject: [PATCH] remove trezor.main, simplify trezor.workflow

---
 src/main.py                 | 12 +++++--
 src/trezor/loop.py          |  6 ++--
 src/trezor/main.py          | 13 --------
 src/trezor/wire/__init__.py |  7 ++++-
 src/trezor/workflow.py      | 63 +++++++++++++++++++------------------
 5 files changed, 52 insertions(+), 49 deletions(-)
 delete mode 100644 src/trezor/main.py

diff --git a/src/main.py b/src/main.py
index 38ee70a3..ec8c753b 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,8 +1,15 @@
 from micropython import const
 
+from trezor import config
 from trezor import io
+from trezor import log
+from trezor import loop
 from trezor import wire
-from trezor import main
+from trezor import workflow
+
+config.init()
+
+log.level = log.DEBUG
 
 # initialize the USB stack
 usb_wire = io.HID(
@@ -97,4 +104,5 @@ wire.setup(usb_wire)
 from apps.homescreen.homescreen import layout_homescreen
 
 # run main even loop and specify which screen is default
-main.run(default_workflow=layout_homescreen)
+workflow.startdefault(layout_homescreen)
+loop.run()
\ No newline at end of file
diff --git a/src/trezor/loop.py b/src/trezor/loop.py
index b9ebd79b..88826788 100644
--- a/src/trezor/loop.py
+++ b/src/trezor/loop.py
@@ -4,7 +4,7 @@ the form of python coroutines (either plain generators or `async` functions) are
 stepped through until completion, and can get asynchronously blocked by
 `yield`ing or `await`ing a syscall.
 
-See `schedule_task`, `run_forever`, and syscalls `Sleep`, `Select`, `Signal`
+See `schedule_task`, `run`, and syscalls `Sleep`, `Select`, `Signal`
 and `Wait`.
 '''
 
@@ -41,7 +41,7 @@ if __debug__:
 def schedule_task(task, value=None, deadline=None):
     '''
     Schedule task to be executed with `value` on given `deadline` (in
-    microseconds).  Does not start the event loop itself, see `run_forever`.
+    microseconds).  Does not start the event loop itself, see `run`.
     '''
     if deadline is None:
         deadline = utime.ticks_us()
@@ -75,7 +75,7 @@ def _unpause_task(task):
             _paused_tasks[iface].remove(task)
 
 
-def run_forever():
+def run():
     '''
     Loop forever, stepping through scheduled tasks and awaiting I/O events
     inbetween.  Use `schedule_task` first to add a coroutine to the task queue.
diff --git a/src/trezor/main.py b/src/trezor/main.py
deleted file mode 100644
index b8167198..00000000
--- a/src/trezor/main.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from trezor import config
-from trezor import loop
-from trezor import workflow
-from trezor import log
-
-config.init()
-
-log.level = log.DEBUG
-
-
-def run(default_workflow):
-    workflow.start_default(default_workflow)
-    loop.run_forever()
diff --git a/src/trezor/wire/__init__.py b/src/trezor/wire/__init__.py
index 8f1e5ea3..7c578513 100644
--- a/src/trezor/wire/__init__.py
+++ b/src/trezor/wire/__init__.py
@@ -112,7 +112,12 @@ async def session_handler(iface, sid):
             except KeyError:
                 handler, args = unexpected_msg, ()
 
-            await handler(ctx, reader, *args)
+            w = handler(ctx, reader, *args)
+            try:
+                workflow.onstart(w)
+                await w
+            finally:
+                workflow.onclose(w)
 
         except UnexpectedMessageError as exc:
             # retry with opened reader from the exception
diff --git a/src/trezor/workflow.py b/src/trezor/workflow.py
index 6b881dde..d52eee16 100644
--- a/src/trezor/workflow.py
+++ b/src/trezor/workflow.py
@@ -1,40 +1,43 @@
-from trezor import log, loop, ui
+from trezor import log
+from trezor import loop
+from trezor import ui
 
-_started = []
-_default = None
-_default_genfunc = None
+started = []
+default = None
+default_handler = None
 
 
-def start_default(genfunc):
-    global _default
-    global _default_genfunc
-    _default_genfunc = genfunc
-    _default = _default_genfunc()
-    log.info(__name__, 'start default %s', _default)
-    loop.schedule_task(_default)
+def onstart(w):
+    closedefault()
+    started.append(w)
     ui.display.backlight(ui.BACKLIGHT_NORMAL)
+    log.debug(__name__, 'onstart: %s', w)
 
 
-def close_default():
-    global _default
-    if _default is not None:
-        log.info(__name__, 'close default %s', _default)
-        _default.close()
-        _default = None
+def onclose(w):
+    started.remove(w)
+    log.debug(__name__, 'onclose: %s', w)
+
+    if not started and default_handler:
+        startdefault(default_handler)
 
 
-def start(workflow):
-    close_default()
-    _started.append(workflow)
-    log.info(__name__, 'start %s', workflow)
-    loop.schedule_task(_watch(workflow))
-    ui.display.backlight(ui.BACKLIGHT_NORMAL)
+def closedefault():
+    global default
+
+    if default:
+        default.close()
+        default = None
+        log.debug(__name__, 'closedefault')
 
 
-async def _watch(workflow):
-    try:
-        return await workflow
-    finally:
-        _started.remove(workflow)
-        if not _started and _default_genfunc is not None:
-            start_default(_default_genfunc)
+def startdefault(handler):
+    global default
+    global default_handler
+
+    if not default:
+        default_handler = handler
+        default = handler()
+        loop.schedule_task(default)
+        ui.display.backlight(ui.BACKLIGHT_NORMAL)
+        log.debug(__name__, 'startdefault')
\ No newline at end of file