From 1d31bee775be3878445730b4223b456060592ada Mon Sep 17 00:00:00 2001
From: Danny Staple <danny@orionrobots.co.uk>
Date: Sun, 13 Nov 2022 18:46:59 +0000
Subject: [PATCH] Make a little robust with lessons learned later

---
 .../computer/display_from_robot.py            | 47 ++++++++++---------
 ch-13/1-modelling-space/robot/code.py         | 11 +++--
 2 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/ch-13/1-modelling-space/computer/display_from_robot.py b/ch-13/1-modelling-space/computer/display_from_robot.py
index 8875729..a077197 100644
--- a/ch-13/1-modelling-space/computer/display_from_robot.py
+++ b/ch-13/1-modelling-space/computer/display_from_robot.py
@@ -1,6 +1,5 @@
 import asyncio
 import json
-
 import matplotlib.pyplot as plt
 
 from robot_ble_connection import BleConnection
@@ -10,7 +9,7 @@ class RobotDisplay:
     def __init__(self):
         self.ble_connection = BleConnection(self.handle_data)
         self.line = ""
-
+        self.arena = {}
         self.display_closed = False
 
     def handle_close(self, _):
@@ -18,41 +17,43 @@ class RobotDisplay:
 
     def handle_data(self, data):
         self.line += data.decode("utf-8")
-        if not self.line.endswith("\n"):
-            return
-        print(f"Received data: {self.line}")
-        message = json.loads(self.line)
-        self.line = ""
-        if "arena" in message:
-            self.update_arena(message)
+        while "\n" in self.line:
+            line, self.line = self.line.split("\n", 1)
+            print(f"Received data: ```{line}```")
+            try:
+                message = json.loads(line)
+            except ValueError:
+                print("Error parsing JSON")
+                return
+            if "arena" in message:
+                self.arena = message
 
-    def update_arena(self, arena):
+    def draw(self):
         plt.gca().clear()
-        for line in arena["arena"]:
-            plt.gca().plot(
-                [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="black"
-            )
-        for line in arena["target_zone"]:
-            plt.gca().plot(
-                [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="red"
-            )
+        if self.arena:
+            for line in self.arena["arena"]:
+                plt.gca().plot(
+                    [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="black"
+                )
+            for line in self.arena["target_zone"]:
+                plt.gca().plot(
+                    [line[0][0], line[1][0]], [line[0][1], line[1][1]], color="red"
+                )
 
     async def main(self):
         plt.ion()
-
+        await self.ble_connection.connect()
         try:
-            await self.ble_connection.connect()
             request = json.dumps({"command": "arena"}).encode()
             print(f"Sending request for arena: {request}")
             await self.ble_connection.send_uart_data(request)
             plt.gcf().canvas.mpl_connect("close_event", self.handle_close)
 
             while not self.display_closed:
-                plt.pause(0.05)
+                self.draw()
                 plt.draw()
+                plt.pause(0.05)
                 await asyncio.sleep(0.01)
-
-            plt.show()
         finally:
             await self.ble_connection.close()
 
diff --git a/ch-13/1-modelling-space/robot/code.py b/ch-13/1-modelling-space/robot/code.py
index 960da10..c9cb35f 100644
--- a/ch-13/1-modelling-space/robot/code.py
+++ b/ch-13/1-modelling-space/robot/code.py
@@ -9,10 +9,12 @@ async def command_handler():
   while True:
     if robot.uart.in_waiting:
       print("Receiving data...")
-      data = robot.uart.readline().decode()
-      print(f"Received data: {data}")
-      request = json.loads(data)
-      print(f"Received command: {request}")
+      try:
+        data = robot.uart.readline().decode()
+        request = json.loads(data)
+      except (UnicodeError, ValueError):
+        print("Invalid data")
+        continue
       # {"command": "arena"}
       if request["command"] == "arena":
          response = {
@@ -20,5 +22,6 @@ async def command_handler():
             "target_zone": arena.target_zone,
          }
          robot.uart.write((json.dumps(response)+"\n").encode())
+    await asyncio.sleep(0.1)
 
 asyncio.run(command_handler())