From 7688f380b122603cfdcd2541a2ad6d8014e5d623 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Bourdeauducq?= <sb@m-labs.hk>
Date: Mon, 26 Feb 2024 19:30:31 +0800
Subject: [PATCH] environment, artiq_run: introduce interactive arguments

---
 RELEASE_NOTES.rst             |  1 +
 artiq/frontend/artiq_run.py   | 23 +++++++++++++++++++++--
 artiq/language/environment.py | 31 +++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst
index d8d5e3029..fa801a11f 100644
--- a/RELEASE_NOTES.rst
+++ b/RELEASE_NOTES.rst
@@ -41,6 +41,7 @@ Highlights:
       + CTRL+SHIFT+C cascades experiment windows
 * Datasets can now be associated with units and scale factors, and displayed accordingly in the dashboard
   including applets, like widgets such as ``NumberValue`` already did in earlier ARTIQ versions.
+* Experiments can now request arguments interactively from the user at any time.
 * Persistent datasets are now stored in a LMDB database for improved performance.
 * Python's built-in types (such as ``float``, or ``List[...]``) can now be used in type annotations on
   kernel functions.
diff --git a/artiq/frontend/artiq_run.py b/artiq/frontend/artiq_run.py
index ae3e5f1c7..19d83e787 100755
--- a/artiq/frontend/artiq_run.py
+++ b/artiq/frontend/artiq_run.py
@@ -13,7 +13,7 @@ import h5py
 
 from llvmlite import binding as llvm
 
-from sipyco import common_args
+from sipyco import common_args, pyon
 
 from artiq import __version__ as artiq_version
 from artiq.language.environment import EnvExperiment, ProcessArgumentManager
@@ -166,9 +166,28 @@ def get_argparser(with_file=True):
     return parser
 
 
+class ArgumentManager(ProcessArgumentManager):
+    def get_interactive(self, interactive_arglist):
+        result = dict()
+        for key, processor, group, tooltip in interactive_arglist:
+            success = False
+            while not success:
+                user_input = input("{}:{} (group={}, tooltip={}): ".format(
+                    key, type(processor).__name__, group, tooltip))
+                try:
+                    user_input_deser = pyon.decode(user_input)
+                    value = processor.process(user_input_deser)
+                except:
+                    logger.error("failed to process user input, retrying", exc_info=True)
+                else:
+                    success = True
+            result[key] = value
+        return result
+
+
 def _build_experiment(device_mgr, dataset_mgr, args):
     arguments = parse_arguments(args.arguments)
-    argument_mgr = ProcessArgumentManager(arguments)
+    argument_mgr = ArgumentManager(arguments)
     managers = (device_mgr, dataset_mgr, argument_mgr, {})
     if hasattr(args, "file"):
         is_tar = tarfile.is_tarfile(args.file)
diff --git a/artiq/language/environment.py b/artiq/language/environment.py
index ba42b2791..797f1eeeb 100644
--- a/artiq/language/environment.py
+++ b/artiq/language/environment.py
@@ -1,5 +1,7 @@
 from collections import OrderedDict
 from inspect import isclass
+from contextlib import contextmanager
+from types import SimpleNamespace
 
 from sipyco import pyon
 
@@ -212,6 +214,9 @@ class TraceArgumentManager:
         self.requested_args[key] = processor, group, tooltip
         return None
 
+    def get_interactive(self, interactive_arglist):
+        raise NotImplementedError
+
 
 class ProcessArgumentManager:
     def __init__(self, unprocessed_arguments):
@@ -233,6 +238,10 @@ class ProcessArgumentManager:
             raise AttributeError("Supplied argument(s) not queried in experiment: " +
                                  ", ".join(unprocessed))
 
+    def get_interactive(self, interactive_arglist):
+        raise NotImplementedError
+
+
 class HasEnvironment:
     """Provides methods to manage the environment of an experiment (arguments,
     devices, datasets)."""
@@ -322,6 +331,28 @@ class HasEnvironment:
         kernel_invariants = getattr(self, "kernel_invariants", set())
         self.kernel_invariants = kernel_invariants | {key}
 
+    @contextmanager
+    def interactive(self):
+        """Request arguments from the user interactively.
+
+        This context manager returns a namespace object on which the method
+        `setattr_argument` should be called, with the usual semantics.
+
+        When the context manager terminates, the experiment is blocked
+        and the user is presented with the requested argument widgets.
+        After the user enters values, the experiment is resumed and
+        the namespace contains the values of the arguments."""
+        interactive_arglist = []
+        namespace = SimpleNamespace()
+        def setattr_argument(key, processor=None, group=None, tooltip=None):
+            interactive_arglist.append((key, processor, group, tooltip))
+        namespace.setattr_argument = setattr_argument
+        yield namespace
+        del namespace.setattr_argument
+        argdict = self.__argument_mgr.get_interactive(interactive_arglist)
+        for key, value in argdict.items():
+            setattr(namespace, key, value)
+
     def get_device_db(self):
         """Returns the full contents of the device database."""
         return self.__device_mgr.get_device_db()