Rework hub and port enumeration

When a new hub instance is created, the user should be able to enumerate
all attached ports. This is now implemented in the Hub() constructor with
the optional argument `enumerate`. The argument defaults to False.

In order to avoid code deduplication, the port enumeration in discover_hubs()
has been modified to use the new functionality with Hub(path, enumerate=True).

Signed-off-by: Nicolai Buchwitz <nb@tipi-net.de>

Author: nbuchwitz

uhubctl/__init__.py

@@ -1,10 +1,11 @@
 import re
 import subprocess
+from typing import List, Optional
 
 UHUBCTL_BINARY = "uhubctl"
 
 
-def __uhubctl(args: list = []) -> list:
+def _uhubctl(args: list = []) -> list:
     cmd = UHUBCTL_BINARY.split(" ") + args
     result = subprocess.run(cmd, capture_output=True)
     stdout = result.stdout.decode()
@@ -18,62 +19,123 @@ def __uhubctl(args: list = []) -> list:
 
 
 def discover_hubs():
-    hubs = []
+    """
+    Return list of all by uhubctl supported USB hubs with their ports
 
-    pattern_hub = re.compile("Current status for hub ([\.\d-]+)")
-    pattern_port = re.compile("  Port (\d+): \d{4} (power|off)")
+    Returns:
+        List of hubs
 
-    for line in __uhubctl():
-        reg_hub = pattern_hub.match(line)
-        reg_port = pattern_port.match(line)
+    """
+    hubs = []
 
-        if reg_hub:
-            hub = Hub(reg_hub.group(1))
-            hubs.append(hub)
-        elif reg_port:
-            # assume that the port is the last detected one
-            hub = hubs[-1]
+    pattern = re.compile("Current status for hub ([\.\d-]+)")
 
-            if not hub:
-                continue
+    for line in _uhubctl():
+        regex = pattern.match(line)
 
-            port = Port(hub, reg_port.group(1))
-            hub.ports.append(port)
+        if regex:
+            hub = Hub(regex.group(1), enumerate=True)
+            hubs.append(hub)
 
     return hubs
 
 
 class Hub:
-    def __init__(self, path: str) -> None:
-        self.path = path
-        self.ports = []
+    def __init__(self, path: str, enumerate: bool = False) -> None:
+        """
+        Create new hub instance
+
+        Arguments:
+            path: USB hub path identifier
+            enumerate: Automatically enumerate ports
+        """
+        self.path: str = path
+        self.ports: List[Port] = []
+
+        if enumerate:
+            self.discover_ports()
 
     def add_port(self, port_number: int) -> 'Port':
+        """
+        Add port to hub by port number
+
+        Arguments:
+            port_number: Indentification number of port
+
+        Returns:
+            Port
+
+        """
         port = Port(self, port_number)
         self.ports.append(port)
 
         return port
 
     def add_ports(self, port_start: int, port_end: int):
+        """
+        Add multiple ports to hub
+
+        Arguments:
+            port_start: First port's indentification number
+            port_end: Last port's ndentification number
+        """
         for port_number in range(port_start, port_end):
             self.add_port(port_number)
 
+    def find_port(self, port_number: int) -> Optional['Port']:
+        """
+        Find port by port number
+
+        Arguments:
+            port_number: Identification number of port to find
+
+        Returns:
+            Port or None
+        """
+
+        for port in self.ports:
+            if port.port_number == int(port_number):
+                return port
+
+        return None
+
+    def discover_ports(self) -> None:
+        """
+        Discover ports for this hub instance
+        """
+        pattern = re.compile("  Port (\d+): \d{4} (power|off)")
+
+        for line in _uhubctl(["-l", self.path]):
+            regex = pattern.match(line)
+
+            if regex:
+                port = Port(self, regex.group(1))
+                self.ports.append(port)
+
     def __str__(self) -> str:
         return f"USB Hub {self.path}"
 
 
 class Port:
     def __init__(self, hub: Hub, port_number: int):
+        """
+        Create new port instance
+
+        Arguments:
+            hub: Hub to attach port to
+            port_number: Number of port to create
+
+        """
         self.hub = hub
-        self.port_number = port_number
+        self.port_number = int(port_number)
 
     @property
     def status(self) -> bool:
         status = None
         pattern = re.compile(f"  Port {self.port_number}: \d{{4}} (power|off)")
 
         args = ["-l", self.hub.path, "-p", self.port_number]
-        for line in __uhubctl(args):
+        for line in _uhubctl(args):
             reg = pattern.match(line)
 
             if reg:
@@ -93,7 +155,7 @@ def status(self, status: bool) -> None:
         else:
             args.append("off")
 
-        __uhubctl(args)
+        _uhubctl(args)
 
     def __str__(self) -> str:
         return f"USB Port {self.hub.path}.{self.port_number}"