From 19b652d4c06b6a1a135bb059973e32ddede0812c Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Fri, 22 Mar 2024 17:30:56 +0800
Subject: [PATCH] fix interactive args cancellation

---
 artiq/language/environment.py | 16 ++++++++++++++--
 artiq/master/databases.py     |  2 +-
 artiq/master/worker_impl.py   |  5 +++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/artiq/language/environment.py b/artiq/language/environment.py
index 8b4076a4c..fe5753184 100644
--- a/artiq/language/environment.py
+++ b/artiq/language/environment.py
@@ -12,7 +12,8 @@ from artiq.language.core import rpc
 __all__ = ["NoDefault", "DefaultMissing",
            "PYONValue", "BooleanValue", "EnumerationValue",
            "NumberValue", "StringValue",
-           "HasEnvironment", "Experiment", "EnvExperiment"]
+           "HasEnvironment", "Experiment", "EnvExperiment",
+           "CancelledArgsError"]
 
 
 class NoDefault:
@@ -26,6 +27,12 @@ class DefaultMissing(Exception):
     pass
 
 
+class CancelledArgsError(Exception):
+    """Raised by the ``interactive`` context manager when an interactive
+    arguments request is cancelled."""
+    pass
+
+
 class _SimpleArgProcessor:
     def __init__(self, default=NoDefault):
         # If default is a list, it means multiple defaults are specified, with
@@ -341,7 +348,10 @@ class HasEnvironment:
         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."""
+        the namespace contains the values of the arguments.
+
+        If the interactive arguments request is cancelled, raises
+        ``CancelledArgsError``."""
         interactive_arglist = []
         namespace = SimpleNamespace()
         def setattr_argument(key, processor=None, group=None, tooltip=None):
@@ -350,6 +360,8 @@ class HasEnvironment:
         yield namespace
         del namespace.setattr_argument
         argdict = self.__argument_mgr.get_interactive(interactive_arglist, title)
+        if argdict is None:
+            raise CancelledArgsError
         for key, value in argdict.items():
             setattr(namespace, key, value)
 
diff --git a/artiq/master/databases.py b/artiq/master/databases.py
index 7a4a6b153..6edcc966d 100644
--- a/artiq/master/databases.py
+++ b/artiq/master/databases.py
@@ -144,4 +144,4 @@ class InteractiveArgDB:
         if rid not in self.futures:
             raise ValueError("no experiment with this RID is "
                              "waiting for interactive arguments")
-        self.futures[rid].cancel()
+        self.futures[rid].set_result(None)
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 65c7f2a88..1f069a50c 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -222,8 +222,9 @@ class ArgumentManager(ProcessArgumentManager):
         arglist_desc = [(k, p.describe(), g, t)
                         for k, p, g, t in interactive_arglist]
         arguments = ArgumentManager._get_interactive(arglist_desc, title)
-        for key, processor, _, _ in interactive_arglist:
-            arguments[key] = processor.process(arguments[key])
+        if arguments is not None:
+            for key, processor, _, _ in interactive_arglist:
+                arguments[key] = processor.process(arguments[key])
         return arguments