python lib: close USB context properly (#1606)

* python lib: close USB context properly

* fix jungle

* rm that

---------

Co-authored-by: Comma Device <device@comma.ai>

board/jungle/__init__.py

@@ -50,7 +50,7 @@ class PandaJungle(Panda):
 
   @classmethod
   def spi_connect(cls, serial, ignore_version=False):
-    return None, None, None, None
+    return None, None, None, None, None
 
   def flash(self, fn=None, code=None, reconnect=True):
     if not fn:

python/__init__.py

@@ -280,6 +280,8 @@ class Panda:
     if self._handle_open:
       self._handle.close()
       self._handle_open = False
+      if self._context is not None:
+        self._context.close()
 
   def connect(self, claim=True, wait=False):
     self.close()
@@ -287,9 +289,9 @@ class Panda:
     self._handle = None
     while self._handle is None:
       # try USB first, then SPI
-      self._handle, serial, self.bootstub, bcd = self.usb_connect(self._connect_serial, claim=claim)
+      self._context, self._handle, serial, self.bootstub, bcd = self.usb_connect(self._connect_serial, claim=claim)
       if self._handle is None:
-        self._handle, serial, self.bootstub, bcd = self.spi_connect(self._connect_serial)
+        self._context, self._handle, serial, self.bootstub, bcd = self.spi_connect(self._connect_serial)
       if not wait:
         break
 
@@ -365,7 +367,7 @@ class Panda:
         err = f"panda protocol mismatch: expected {handle.PROTOCOL_VERSION}, got {spi_version}. reflash panda"
         raise PandaProtocolMismatch(err)
 
-    return handle, spi_serial, bootstub, None
+    return None, handle, spi_serial, bootstub, None
 
   @classmethod
   def usb_connect(cls, serial, claim=True):
@@ -407,7 +409,7 @@ class Panda:
     else:
       context.close()
 
-    return usb_handle, usb_serial, bootstub, bcd
+    return context, usb_handle, usb_serial, bootstub, bcd
 
   @classmethod
   def list(cls): # noqa: A003
@@ -436,7 +438,7 @@ class Panda:
 
   @classmethod
   def spi_list(cls):
-    _, serial, _, _ = cls.spi_connect(None, ignore_version=True)
+    _, _, serial, _, _ = cls.spi_connect(None, ignore_version=True)
     if serial is not None:
       return [serial, ]
     return []

python/dfu.py

@@ -14,9 +14,9 @@ class PandaDFU:
   def __init__(self, dfu_serial: Optional[str]):
     # try USB, then SPI
     handle: Optional[BaseSTBootloaderHandle]
-    handle = PandaDFU.usb_connect(dfu_serial)
+    self._context, handle = PandaDFU.usb_connect(dfu_serial)
     if handle is None:
-      handle = PandaDFU.spi_connect(dfu_serial)
+      self._context, handle = PandaDFU.spi_connect(dfu_serial)
 
     if handle is None:
       raise Exception(f"failed to open DFU device {dfu_serial}")
@@ -24,8 +24,21 @@ class PandaDFU:
     self._handle: BaseSTBootloaderHandle = handle
     self._mcu_type: McuType = self._handle.get_mcu_type()
 
+  def __enter__(self):
+    return self
+
+  def __exit__(self, *args):
+    self.close()
+
+  def close(self):
+    if self._handle is not None:
+      self._handle.close()
+      self._handle = None
+      if self._context is not None:
+        self._context.close()
+
   @staticmethod
-  def usb_connect(dfu_serial: Optional[str]) -> Optional[STBootloaderUSBHandle]:
+  def usb_connect(dfu_serial: Optional[str]):
     handle = None
     context = usb1.USBContext()
     context.open()
@@ -40,10 +53,10 @@ class PandaDFU:
           handle = STBootloaderUSBHandle(device, device.open())
           break
 
-    return handle
+    return context, handle
 
   @staticmethod
-  def spi_connect(dfu_serial: Optional[str]) -> Optional[STBootloaderSPIHandle]:
+  def spi_connect(dfu_serial: Optional[str]):
     handle = None
     this_dfu_serial = None
 
@@ -56,7 +69,7 @@ class PandaDFU:
     if dfu_serial is not None and dfu_serial != this_dfu_serial:
       handle = None
 
-    return handle
+    return None, handle
 
   @staticmethod
   def list() -> List[str]: # noqa: A003
@@ -82,7 +95,7 @@ class PandaDFU:
   @staticmethod
   def spi_list() -> List[str]:
     try:
-      h = PandaDFU.spi_connect(None)
+      _, h = PandaDFU.spi_connect(None)
       if h is not None:
         dfu_serial = PandaDFU.st_serial_to_dfu_serial(h.get_uid(), h.get_mcu_type())
         return [dfu_serial, ]