examples/discover-loop.py: rework to use 'with' statement

Now that we have resource tracking we can use the 'with' statement
for discovery and clean up connections and discovery controllers
on exit.

Signed-off-by: Hannes Reinecke <[email protected]>

Author: hreinecke

examples/discover-loop.py

@@ -16,46 +16,43 @@ def disc_supp_str(dlp_supp_opts):
     }
     return [txt for msk, txt in d.items() if dlp_supp_opts & msk]
 
-root = nvme.root()
-host = nvme.host(root)
-
-subsysnqn = nvme.NVME_DISC_SUBSYS_NAME
-transport = 'tcp'
-traddr = '127.0.0.1'
-trsvcid = '4420'
+def discover(host, ctrl):
+    try:
+        ctrl.connect(host)
+    except Exception as e:
+        print(f'Failed to connect: {e}')
+        return
 
-ctrl = nvme.ctrl(root, subsysnqn=subsysnqn, transport=transport, traddr=traddr, trsvcid=trsvcid)
+    print(f'{ctrl.name} connected to subsys {ctrl.subsystem}')
 
-try:
-    ctrl.connect(host)
-except Exception as e:
-    sys.exit(f'Failed to connect: {e}')
+    slp = ctrl.supported_log_pages()
+    try:
+        dlp_supp_opts = slp[nvme.NVME_LOG_LID_DISCOVER] >> 16
+    except (TypeError, IndexError):
+        dlp_supp_opts = 0
 
-print(f'{ctrl.name} connected to subsys {ctrl.subsystem}')
+    print(f"LID {nvme.NVME_LOG_LID_DISCOVER}h (Discovery), supports: {disc_supp_str(dlp_supp_opts)}")
 
-slp = ctrl.supported_log_pages()
+    try:
+        lsp = nvme.NVMF_LOG_DISC_LSP_PLEO if dlp_supp_opts & nvme.NVMF_LOG_DISC_LID_PLEOS else 0
+        disc_log = ctrl.discover(lsp=lsp)
+    except Exception as e:
+        print(f'Failed to discover: {e}')
+        return
 
-try:
-    dlp_supp_opts = slp[nvme.NVME_LOG_LID_DISCOVER] >> 16
-except (TypeError, IndexError):
-    dlp_supp_opts = 0
+    for dlpe in disc_log:
+        print(f'log entry {dlpe["portid"]}: {dlpe["subtype"]} {dlpe["subnqn"]}')
 
-print(f"LID {nvme.NVME_LOG_LID_DISCOVER}h (Discovery), supports: {disc_supp_str(dlp_supp_opts)}")
-
-try:
-    lsp = nvme.NVMF_LOG_DISC_LSP_PLEO if dlp_supp_opts & nvme.NVMF_LOG_DISC_LID_PLEOS else 0
-    disc_log = ctrl.discover(lsp=lsp)
-except Exception as e:
-    print(f'Failed to discover: {e}')
-    disc_log = []
+root = nvme.root()
+host = nvme.host(root)
 
-for dlpe in disc_log:
-    print(f'log entry {dlpe["portid"]}: {dlpe["subtype"]} {dlpe["subnqn"]}')
+subsysnqn = nvme.NVME_DISC_SUBSYS_NAME
+transport = 'tcp'
+traddr = '127.0.0.1'
+trsvcid = '4420'
 
-try:
-    ctrl.disconnect()
-except Exception as e:
-    sys.exit(f'Failed to disconnect: {e}')
+with nvme.ctrl(root, subsysnqn=subsysnqn, transport=transport, traddr=traddr, trsvcid=trsvcid) as ctrl:
+    discover(host, ctrl)
 
 for s in host.subsystems():
     for c in s.controllers():