From 70110187ccdaad848d7bd8c792c31516803e77ba Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Sun, 25 Sep 2016 15:55:08 +0200
Subject: [PATCH] overload __iter__ in loop.Wait

This will automatically exit child tasks in case of close() or throw() on the waiting task, but only if run through `await` or `yield from`
---
 src/apps/homescreen/layout_homescreen.py | 6 +-----
 src/trezor/loop.py                       | 7 +++++++
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/apps/homescreen/layout_homescreen.py b/src/apps/homescreen/layout_homescreen.py
index 8a064a38..36f73544 100644
--- a/src/apps/homescreen/layout_homescreen.py
+++ b/src/apps/homescreen/layout_homescreen.py
@@ -15,8 +15,4 @@ async def animate_logo():
 
 
 async def layout_homescreen():
-    wait = loop.Wait([swipe_to_rotate(), animate_logo()])
-    try:
-        await wait
-    finally:
-        wait.exit()
+    await loop.Wait([swipe_to_rotate(), animate_logo()])
diff --git a/src/trezor/loop.py b/src/trezor/loop.py
index ad94eded..2105e0b6 100644
--- a/src/trezor/loop.py
+++ b/src/trezor/loop.py
@@ -194,3 +194,10 @@ class Wait(Syscall):
             if self.exit_others:
                 self.exit()
             schedule_task(self.callback, result)
+
+    def __iter__(self):
+        try:
+            return (yield self)
+        except:
+            self.exit()
+            raise