From b32062d8551dc18dc8dafdcb7f85553d9f77f5cb Mon Sep 17 00:00:00 2001
From: atse <atse@m-labs.hk>
Date: Mon, 31 Jul 2023 12:33:00 +0800
Subject: [PATCH] More elegant connection stopping

---
 pytec/pytec/aioclient.py | 24 +++++++++++++++---------
 pytec/tec_qt.py          |  9 +++++----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/pytec/pytec/aioclient.py b/pytec/pytec/aioclient.py
index 0095c69..cf016ce 100644
--- a/pytec/pytec/aioclient.py
+++ b/pytec/pytec/aioclient.py
@@ -5,6 +5,9 @@ import logging
 class CommandError(Exception):
     pass
 
+class StoppedConnecting(Exception):
+    pass
+
 class Client:
     def __init__(self):
         self._reader = None
@@ -14,25 +17,28 @@ class Client:
         self._report_mode_on = False
 
     async def connect(self, host='192.168.1.26', port=23, timeout=None):
-        """Connect to the TEC with host and port, throws TimeoutError if
-        unable to connect. Returns True if not cancelled with disconnect.
+        """Connect to the Thermostat with host and port.
+        Throws StoppedConnecting if disconnect was called while connecting.
+        Throws asyncio.TimeoutError if timeout was exceeded.
 
         Example::
-            client = aioclient.Client()
-            connected = await client.connect()
-            if connected:
-                return
+            client = Client()
+            try:
+                await client.connect()
+            except StoppedConnecting:
+                print("Stopped connecting")
         """
-        self._connecting_task = asyncio.create_task(asyncio.wait_for(asyncio.open_connection(host, port), timeout))
+        self._connecting_task = asyncio.create_task(
+            asyncio.wait_for(asyncio.open_connection(host, port), timeout)
+        )
         try:
             self._reader, self._writer = await self._connecting_task
         except asyncio.CancelledError:
-            return False
+            raise StoppedConnecting
         finally:
             self._connecting_task = None
 
         await self._check_zero_limits()
-        return True
 
     def is_connecting(self):
         """Returns True if client is connecting"""
diff --git a/pytec/tec_qt.py b/pytec/tec_qt.py
index 470fb28..f3af821 100644
--- a/pytec/tec_qt.py
+++ b/pytec/tec_qt.py
@@ -11,7 +11,7 @@ import sys
 import argparse
 import logging
 import asyncio
-from pytec.aioclient import Client
+from pytec.aioclient import Client, StoppedConnecting
 import qasync
 from qasync import asyncSlot, asyncClose
 
@@ -483,15 +483,16 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
                 self.ip_set_line.setEnabled(False)
                 self.port_set_spin.setEnabled(False)
 
-                connected = await self.tec_client.connect(host=ip, port=port, timeout=30)
-                if not connected:
+                try:
+                    await self.tec_client.connect(host=ip, port=port, timeout=30)
+                except StoppedConnecting:
                     return
                 await self._on_connection_changed(True)
             else:
                 await self._on_connection_changed(False)
                 await self.tec_client.disconnect()
 
-        except (OSError, TimeoutError) as e:
+        except (OSError, TimeoutError, asyncio.TimeoutError) as e: # TODO: Remove asyncio.TimeoutError in Python 3.11
             logging.error(f"Failed communicating to {ip}:{port}: {e}")
             await self._on_connection_changed(False)
             await self.tec_client.disconnect()