Move PybricksHub client creation into constructor

Moves client creation into the PybricksHub constructor. A
future change will move this into a subclass specific for
BLE or USB, so configuring this is more appropriately done
in the constructor instead of the connect method.

The disconnect logic is moved into a class method so that
it can be shared by both BLE and USB implementations.

Temporarily comments out a message that prints the device
name as this is no longer available in the connect method.

Signed-off-by: Nate Karstens <[email protected]>

Author: nkarstens

pybricksdev/cli/__init__.py

@@ -183,22 +183,21 @@ async def run(self, args: argparse.Namespace):
                 print("--name is required for SSH connections", file=sys.stderr)
                 exit(1)
 
-            hub = EV3Connection()
             device_or_address = socket.gethostbyname(args.name)
+            hub = EV3Connection(device_or_address)
         elif args.conntype == "ble":
             # It is a Pybricks Hub with BLE. Device name or address is given.
-            hub = PybricksHub()
             print(f"Searching for {args.name or 'any hub with Pybricks service'}...")
             device_or_address = await find_device(args.name)
+            hub = PybricksHub(device_or_address)
 
         elif args.conntype == "usb":
             hub = REPLHub()
-            device_or_address = None
         else:
             raise ValueError(f"Unknown connection type: {args.conntype}")
 
         # Connect to the address and run the script
-        await hub.connect(device_or_address)
+        await hub.connect()
         try:
             with _get_script_path(args.file) as script_path:
                 await hub.run(script_path, args.wait)

pybricksdev/connections/ev3dev.py

@@ -18,10 +18,13 @@ class EV3Connection:
     _USER = "robot"
     _PASSWORD = "maker"
 
+    def __init__(self, address: str):
+        self._address = address
+
     def abs_path(self, path):
         return os.path.join(self._HOME, path)
 
-    async def connect(self, address):
+    async def connect(self):
         """Connects to ev3dev using SSH with a known IP address.
 
         Arguments:
@@ -33,9 +36,9 @@ async def connect(self, address):
                 Connect failed.
         """
 
-        print("Connecting to", address, "...", end=" ")
+        print("Connecting to", self._address, "...", end=" ")
         self.client = await asyncssh.connect(
-            address, username=self._USER, password=self._PASSWORD
+            self._address, username=self._USER, password=self._PASSWORD
         )
         print("Connected.", end=" ")
         self.client.sftp = await self.client.start_sftp_client()

pybricksdev/connections/pybricks.py

@@ -78,7 +78,7 @@ class PybricksHub:
     has not been connected yet or the connected hub has Pybricks profile < v1.2.0.
     """
 
-    def __init__(self):
+    def __init__(self, device: BLEDevice):
         self.connection_state_observable = BehaviorSubject(ConnectionState.DISCONNECTED)
         self.status_observable = BehaviorSubject(StatusFlag(0))
         self._stdout_subject = Subject()
@@ -120,6 +120,11 @@ def __init__(self):
         # File handle for logging
         self.log_file = None
 
+        def handle_disconnect(_: BleakClient):
+            self._handle_disconnect()
+
+        self.client = BleakClient(device, disconnected_callback=handle_disconnect)
+
     @property
     def stdout_observable(self) -> Observable[bytes]:
         """
@@ -227,18 +232,20 @@ def _pybricks_service_handler(self, _: int, data: bytes) -> None:
             if self._enable_line_handler:
                 self._handle_line_data(payload)
 
-    async def connect(self, device: BLEDevice):
-        """Connects to a device that was discovered with :meth:`pybricksdev.ble.find_device`
+    def _handle_disconnect(self):
+        logger.info("Disconnected!")
+        self.connection_state_observable.on_next(ConnectionState.DISCONNECTED)
 
-        Args:
-            device: The device to connect to.
+    async def connect(self):
+        """Connects to a device that was discovered with :meth:`pybricksdev.ble.find_device`
 
         Raises:
             BleakError: if connecting failed (or old firmware without Device
                 Information Service)
             RuntimeError: if Pybricks Protocol version is not supported
         """
-        logger.info(f"Connecting to {device.name}")
+        # TODO: Fix this
+        # logger.info(f"Connecting to {device.name}")
 
         if self.connection_state_observable.value != ConnectionState.DISCONNECTED:
             raise RuntimeError(
@@ -252,12 +259,6 @@ async def connect(self, device: BLEDevice):
                 self.connection_state_observable.on_next, ConnectionState.DISCONNECTED
             )
 
-            def handle_disconnect(_: BleakClient):
-                logger.info("Disconnected!")
-                self.connection_state_observable.on_next(ConnectionState.DISCONNECTED)
-
-            self.client = BleakClient(device, disconnected_callback=handle_disconnect)
-
             await self.client.connect()
 
             stack.push_async_callback(self.disconnect)