From c8dc2cbf094120de4e560a306d8eb1418663c62c Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Wed, 31 May 2023 21:57:54 +0800
Subject: [PATCH] browser: decouple dataset controller from dataset dock

---
 artiq/browser/datasets.py       | 45 ++++++++++++++++++---------------
 artiq/frontend/artiq_browser.py | 14 +++++-----
 2 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/artiq/browser/datasets.py b/artiq/browser/datasets.py
index 78bf6e8f7..76da6340e 100644
--- a/artiq/browser/datasets.py
+++ b/artiq/browser/datasets.py
@@ -23,8 +23,30 @@ class Model(DictSyncTreeSepModel):
         return short_format(v[1])
 
 
+class DatasetCtl:
+    def __init__(self, master_host, master_port):
+        self.master_host = master_host
+        self.master_port = master_port
+
+    async def set(self, key, value, persist=None):
+        logger.info("Uploading dataset '%s' to master...", key)
+        try:
+            remote = RPCClient()
+            await remote.connect_rpc(self.master_host, self.master_port,
+                                     "master_dataset_db")
+            try:
+                await remote.set(key, value, persist)
+            finally:
+                remote.close_rpc()
+        except:
+            logger.error("Failed uploading dataset '%s'",
+                         key, exc_info=True)
+        else:
+            logger.info("Finished uploading dataset '%s'", key)
+
+
 class DatasetsDock(QtWidgets.QDockWidget):
-    def __init__(self, dataset_sub, master_host, master_port):
+    def __init__(self, dataset_sub, dataset_ctl):
         QtWidgets.QDockWidget.__init__(self, "Datasets")
         self.setObjectName("Datasets")
         self.setFeatures(QtWidgets.QDockWidget.DockWidgetMovable |
@@ -64,8 +86,7 @@ class DatasetsDock(QtWidgets.QDockWidget):
         self.set_model(Model(dict()))
         dataset_sub.add_setmodel_callback(self.set_model)
 
-        self.master_host = master_host
-        self.master_port = master_port
+        self.dataset_ctl = dataset_ctl
 
     def _search_datasets(self):
         if hasattr(self, "table_model_filter"):
@@ -82,22 +103,6 @@ class DatasetsDock(QtWidgets.QDockWidget):
         self.table_model_filter.setSourceModel(self.table_model)
         self.table.setModel(self.table_model_filter)
 
-    async def _upload_dataset(self, name, value,):
-        logger.info("Uploading dataset '%s' to master...", name)
-        try:
-            remote = RPCClient()
-            await remote.connect_rpc(self.master_host, self.master_port,
-                                     "master_dataset_db")
-            try:
-                await remote.set(name, value)
-            finally:
-                remote.close_rpc()
-        except:
-            logger.error("Failed uploading dataset '%s'",
-                         name, exc_info=True)
-        else:
-            logger.info("Finished uploading dataset '%s'", name)
-
     def upload_clicked(self):
         idx = self.table.selectedIndexes()
         if idx:
@@ -105,7 +110,7 @@ class DatasetsDock(QtWidgets.QDockWidget):
             key = self.table_model.index_to_key(idx)
             if key is not None:
                 persist, value = self.table_model.backing_store[key]
-                asyncio.ensure_future(self._upload_dataset(key, value))
+                asyncio.ensure_future(self.dataset_ctl.set(key, value))
 
     def save_state(self):
         return bytes(self.table.header().saveState())
diff --git a/artiq/frontend/artiq_browser.py b/artiq/frontend/artiq_browser.py
index af469c05a..5cb7ef090 100755
--- a/artiq/frontend/artiq_browser.py
+++ b/artiq/frontend/artiq_browser.py
@@ -48,8 +48,8 @@ def get_argparser():
 
 
 class Browser(QtWidgets.QMainWindow):
-    def __init__(self, smgr, dataset_sub, browse_root,
-                 master_host, master_port, *, loop=None):
+    def __init__(self, smgr, dataset_sub, dataset_ctl, browse_root,
+                 *, loop=None):
         QtWidgets.QMainWindow.__init__(self)
         smgr.register(self)
 
@@ -81,12 +81,11 @@ class Browser(QtWidgets.QMainWindow):
         self.files.dataset_changed.connect(
             self.experiments.dataset_changed)
 
-        self.applets = applets.AppletsDock(self, dataset_sub, loop=loop)
+        self.applets = applets.AppletsDock(self, dataset_sub, dataset_ctl, loop=loop)
         smgr.register(self.applets)
         atexit_register_coroutine(self.applets.stop, loop=loop)
 
-        self.datasets = datasets.DatasetsDock(
-            dataset_sub, master_host, master_port)
+        self.datasets = datasets.DatasetsDock(dataset_sub, dataset_ctl)
         smgr.register(self.datasets)
         self.files.metadata_changed.connect(self.datasets.metadata_changed)
 
@@ -151,8 +150,9 @@ def main():
 
     smgr = state.StateManager(args.db_file)
 
-    browser = Browser(smgr, dataset_sub, args.browse_root,
-                      args.server, args.port, loop=loop)
+    dataset_ctl = datasets.DatasetCtl(args.server, args.port)
+    browser = Browser(smgr, dataset_sub, dataset_ctl, args.browse_root,
+                      loop=loop)
     widget_log_handler.callback = browser.log.model.append
 
     if os.name == "nt":