flterm: use asyncio (Linux only for now)

sbourdeauducq · sbourdeauducq · commit 75b3ac991c20 · 2016-02-28T01:57:15.000+08:00
diff --git a/misoc/tools/flterm.py b/misoc/tools/flterm.py
@@ -1,27 +1,23 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python3.5
 
 import sys
 import os
 import time
-import serial
-import threading
+import asyncio
+from serial import aio as asyncserial
 import argparse
 
-if sys.platform == "win32":
-    import msvcrt
 
-    def init_getkey():
-        pass
+if sys.platform == "win32":
+    def init_getkey(callback):
+        raise NotImplementedError
 
     def deinit_getkey():
-        pass
-
-    def getkey():
-        return msvcrt.getch()
+        raise NotImplementedError
 else:
     import termios
 
-    def init_getkey():
+    def init_getkey(callback):
         global old_termios
 
         fd = sys.stdin.fileno()
@@ -30,11 +26,12 @@ def init_getkey():
         new[3] = new[3] & ~termios.ICANON & ~termios.ECHO
         termios.tcsetattr(fd, termios.TCSANOW, new)
 
+        loop = asyncio.get_event_loop()
+        loop.add_reader(sys.stdin.fileno(), callback)
+
     def deinit_getkey():
         termios.tcsetattr(sys.stdin.fileno(), termios.TCSANOW, old_termios)
 
-    def getkey():
-        return os.read(sys.stdin.fileno(), 1)
 
 
 sfl_magic_req = b"sL5DdSMmkekro\n"
@@ -115,39 +112,21 @@ class Flterm:
     def __init__(self, kernel_image, kernel_address):
         self.kernel_image = kernel_image
         self.kernel_address = kernel_address
+        self.magic_detect_buffer = b"\x00"*len(sfl_magic_req)
 
-        self.reader_alive = False
-        self.writer_alive = False
-
-        self.magic_detect_buffer = bytes(len(sfl_magic_req))
-
-    def open(self, port, baudrate):
-        if hasattr(self, "port"):
-            return
-        self.port = serial.serial_for_url(port, baudrate)
-
-    def close(self):
-        if not hasattr(self, "port"):
-            return
-        self.port.close()
-        del self.port
-
-    def send_frame(self, frame):
-        retry = 1
+    async def send_frame(self, frame):
         while retry:
-            self.port.write(frame.encode())
-            # Get the reply from the device
-            reply = self.port.read()
+            self.writer.write(frame.encode())
+            reply = await self.reader.read(1)
             if reply == sfl_ack_success:
-                retry = 0
+                return
             elif reply == sfl_ack_crcerror:
-                retry = 1
+                pass  # retry
             else:
                 print("[FLTERM] Got unknown reply '{}' from the device, aborting.".format(reply))
-                return 0
-        return 1
+                raise ValueError
 
-    def upload(self, filename, address):
+    async def upload(self, filename, address):
         with open(filename, "rb") as f:
             data = f.read()
         print("[FLTERM] Uploading {} ({} bytes)...".format(filename, len(data)))
@@ -162,7 +141,9 @@ def upload(self, filename, address):
             frame.cmd = sfl_cmd_load
             frame.payload = current_address.to_bytes(4, "big")
             frame.payload += frame_data
-            if self.send_frame(frame) == 0:
+            try:
+                await self.send_frame(frame)
+            except ValueError:
                 return
             current_address += len(frame_data)
             position += len(frame_data)
@@ -175,111 +156,85 @@ def upload(self, filename, address):
         print("[FLTERM] Upload complete ({0:.1f}KB/s).".format(length/(elapsed*1024)))
         return length
 
-    def boot(self):
+    async def boot(self):
         print("[FLTERM] Booting the device.")
         frame = SFLFrame()
         frame.cmd = sfl_cmd_jump
         frame.payload = self.kernel_address.to_bytes(4, "big") 
-        self.send_frame(frame)
+        await self.send_frame(frame)
 
-    def detect_magic(self, data):
-        if len(data):
-            self.magic_detect_buffer = self.magic_detect_buffer[1:] + data
-            return self.magic_detect_buffer == sfl_magic_req
-        else:
-            return False
-
-    def answer_magic(self):
+    async def answer_magic(self):
         print("[FLTERM] Received firmware download request from the device.")
-        if os.path.exists(self.kernel_image):
-            self.port.write(sfl_magic_ack)
-            self.upload(self.kernel_image, self.kernel_address)
-            self.boot()
-        print("[FLTERM] Done.");
-
-    def reader(self):
+        self.writer.write(sfl_magic_ack)
         try:
-            while self.reader_alive:
-                c = self.port.read()
-                if c == b"\r":
-                    sys.stdout.write(b"\n")
-                else:
-                    sys.stdout.buffer.write(c)
-                sys.stdout.flush()
-
-                if self.kernel_image is not None:
-                    if self.detect_magic(c):
-                        self.answer_magic()
-
-        except serial.SerialException:
-            self.reader_alive = False
-            raise
-
-    def start_reader(self):
-        self.reader_alive = True
-        self.reader_thread = threading.Thread(target=self.reader)
-        self.reader_thread.setDaemon(True)
-        self.reader_thread.start()
-
-    def stop_reader(self):
-        self.reader_alive = False
-        self.reader_thread.join()
-
-    def writer(self):
-        try:
-            while self.writer_alive:
-                self.port.write(getkey())
-        except:
-            self.writer_alive = False
-            raise
+            await self.upload(self.kernel_image, self.kernel_address)
+        except FileNotFoundError:
+            print("[FLTERM] File not found")
+        else:
+            await self.boot()
+        print("[FLTERM] Done.");
 
-    def start_writer(self):
-        self.writer_alive = True
-        self.writer_thread = threading.Thread(target=self.writer)
-        self.writer_thread.setDaemon(True)
-        self.writer_thread.start()
+    async def reader_coro(self):
+        while True:
+            c = await self.reader.read(1)
+            if c == b"\r":
+                sys.stdout.write(b"\n")
+            else:
+                sys.stdout.buffer.write(c)
+            sys.stdout.flush()
 
-    def stop_writer(self):
-        self.writer_alive = False
-        self.writer_thread.join()
+            if self.kernel_image is not None:
+                self.magic_detect_buffer = self.magic_detect_buffer[1:] + c
+                if self.magic_detect_buffer == sfl_magic_req:
+                    self.answer_magic()
 
-    def start(self):
-        print("[FLTERM] Starting....")
-        self.start_reader()
-        self.start_writer()
+    def getkey_callback(self):
+        self.writer.write(os.read(sys.stdin.fileno(), 1))
 
-    def stop(self):
-        self.reader_alive = False
-        self.writer_alive = False
+    async def open(self, port, baudrate):
+        self.reader, self.writer = await asyncserial.open_serial_connection(
+            port, baudrate=baudrate)
+        self.reader_task = asyncio.ensure_future(self.reader_coro())
+        init_getkey(self.getkey_callback)
 
-    def join(self, writer_only=False):
-        self.writer_thread.join()
-        if not writer_only:
-            self.reader_thread.join()
+    async def close(self):
+        self.reader_task.cancel()
+        try:
+            await asyncio.wait_for(self.reader_task, None)
+        except asyncio.CancelledError:
+            pass
+        self.writer.close()
+        deinit_getkey()
 
 
 def _get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("port", help="serial port")
     parser.add_argument("--speed", default=115200, help="serial baudrate")
     parser.add_argument("--kernel", default=None, help="kernel image")
-    parser.add_argument("--kernel-adr", type=lambda a: int(a, 0), default=0x40000000, help="kernel address")
+    parser.add_argument("--kernel-adr", type=lambda a: int(a, 0),
+                        default=0x40000000, help="kernel address")
     return parser.parse_args()
 
 
 def main():
-    args = _get_args()
-    flterm = Flterm(args.kernel, args.kernel_adr)
-    init_getkey()
+    if os.name == "nt":
+        loop = asyncio.ProactorEventLoop()
+        asyncio.set_event_loop(loop)
+    else:
+        loop = asyncio.get_event_loop()
     try:
+        args = _get_args()
+        flterm = Flterm(args.kernel, args.kernel_adr)
+        loop.run_until_complete(
+            flterm.open(args.port, args.speed))
         try:
-            flterm.open(args.port, args.speed)
-            flterm.start()
-            flterm.join(True)
+            loop.run_forever()
         finally:
-            flterm.close()
+            loop.run_until_complete(
+                flterm.close())
     finally:
-        deinit_getkey()
+        loop.close()
 
 
 if __name__ == "__main__":
