environment: refactor

m-labs · Apr 16, 2016 · caf7745 · caf7745
1 parent 12a8c76
commit caf7745
Showing 11 changed files with 97 additions and 115 deletions.
diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst
@@ -36,6 +36,8 @@ unreleased [2.x]
   (i.e. grouping by day and then by hour, instead of by day and then by minute)
 * GUI tools save their state file in the user's home directory instead of the
   current directory.
+* The ``parent`` keyword argument of ``HasEnvironment`` (and ``EnvExperiment``)
+  has been replaced. Pass the parent as first argument instead.
 
 
 unreleased [1.0rc3]

diff --git a/artiq/examples/master/repository/arguments_demo.py b/artiq/examples/master/repository/arguments_demo.py
@@ -54,8 +54,8 @@ def build(self):
         self.setattr_argument("enum", EnumerationValue(
             ["foo", "bar", "quux"], "foo"), "Group")
 
-        self.sc1 = SubComponent1(parent=self)
-        self.sc2 = SubComponent2(parent=self)
+        self.sc1 = SubComponent1(self)
+        self.sc2 = SubComponent2(self)
 
     def run(self):
         logging.error("logging test: error")

diff --git a/artiq/examples/master/repository/coredevice_examples/transport.py b/artiq/examples/master/repository/coredevice_examples/transport.py
@@ -17,8 +17,8 @@ def build(self):
 
         self.setattr_argument("wait_at_stop", NumberValue(100*us))
         self.setattr_argument("speed", NumberValue(1/(10*us)))
-        self.repeats = int(self.get_argument("repeats", NumberValue(100)))
-        self.bins = int(self.get_argument("bins", NumberValue(100)))
+        self.setattr_argument("repeats", NumberValue(100, step=1, ndecimals=0))
+        self.setattr_argument("bins", NumberValue(100, step=1, ndecimals=0))
 
         t = np.linspace(0, 10, 101)  # waveform time
         u = 1 - np.cos(np.pi*t/t[-1])

diff --git a/artiq/examples/master/repository/histograms.py b/artiq/examples/master/repository/histograms.py
@@ -7,9 +7,6 @@
 
 class Histograms(EnvExperiment):
     """Histograms demo"""
-    def build(self):
-        pass
-
     def run(self):
         nbins = 50
         npoints = 20

diff --git a/artiq/examples/master/repository/speed_benchmark.py b/artiq/examples/master/repository/speed_benchmark.py
@@ -102,7 +102,7 @@ def run_with_scheduler(self):
                               self.scheduler.priority, None, False)
 
     def run_without_scheduler(self, pause):
-        payload = globals()["_Payload" + self.payload](*self.managers())
+        payload = globals()["_Payload" + self.payload](self)
 
         start_time = time.monotonic()
         for i in range(int(self.nruns)):

diff --git a/artiq/frontend/artiq_run.py b/artiq/frontend/artiq_run.py
@@ -12,7 +12,7 @@
 
 from llvmlite_artiq import binding as llvm
 
-from artiq.language.environment import EnvExperiment
+from artiq.language.environment import EnvExperiment, ProcessArgumentManager
 from artiq.master.databases import DeviceDB, DatasetDB
 from artiq.master.worker_db import DeviceManager, DatasetManager
 from artiq.coredevice.core import CompileError, host_only
@@ -167,7 +167,8 @@ def _build_experiment(device_mgr, dataset_mgr, args):
         "arguments": arguments
     }
     device_mgr.virtual_devices["scheduler"].expid = expid
-    return exp(device_mgr, dataset_mgr, **arguments)
+    argument_mgr = ProcessArgumentManager(arguments)
+    return exp((device_mgr, dataset_mgr, argument_mgr))
 
 
 def run(with_file=False):

diff --git a/artiq/language/environment.py b/artiq/language/environment.py
@@ -7,8 +7,7 @@
 __all__ = ["NoDefault",
            "PYONValue", "BooleanValue", "EnumerationValue",
            "NumberValue", "StringValue",
-           "HasEnvironment",
-           "Experiment", "EnvExperiment", "is_experiment"]
+           "HasEnvironment", "Experiment", "EnvExperiment"]
 
 
 class NoDefault:
@@ -145,51 +144,63 @@ class StringValue(_SimpleArgProcessor):
     pass
 
 
-class HasEnvironment:
-    """Provides methods to manage the environment of an experiment (devices,
-    parameters, results, arguments)."""
-    def __init__(self, device_mgr=None, dataset_mgr=None, *, parent=None,
-                 default_arg_none=False, **kwargs):
+class TraceArgumentManager:
+    def __init__(self):
         self.requested_args = OrderedDict()
 
-        self.__device_mgr = device_mgr
-        self.__dataset_mgr = dataset_mgr
-        self.__parent = parent
-        self.__default_arg_none = default_arg_none
+    def get(self, key, processor, group):
+        self.requested_args[key] = processor, group
+        return None
+
+
+class ProcessArgumentManager:
+    def __init__(self, unprocessed_arguments):
+        self.unprocessed_arguments = unprocessed_arguments
+
+    def get(self, key, processor, group):
+        if key in self.unprocessed_arguments:
+            r = processor.process(self.unprocessed_arguments[key])
+        else:
+            r = processor.default()
+        return r
+
+
+class HasEnvironment:
+    """Provides methods to manage the environment of an experiment (arguments,
+    devices, datasets)."""
+    def __init__(self, managers_or_parent, *args, **kwargs):
+        if isinstance(managers_or_parent, tuple):
+            self.__device_mgr = managers_or_parent[0]
+            self.__dataset_mgr = managers_or_parent[1]
+            self.__argument_mgr = managers_or_parent[2]
+        else:
+            self.__device_mgr = managers_or_parent.__device_mgr
+            self.__dataset_mgr = managers_or_parent.__dataset_mgr
+            self.__argument_mgr = managers_or_parent.__argument_mgr
 
-        self.__kwargs = kwargs
         self.__in_build = True
-        self.build()
+        self.build(*args, **kwargs)
         self.__in_build = False
-        for key in self.__kwargs.keys():
-            if key not in self.requested_args:
-                raise TypeError("Got unexpected argument: " + key)
-        del self.__kwargs
 
     def build(self):
-        """Must be implemented by the user to request arguments.
+        """Should be implemented by the user to request arguments.
 
-        Other initialization steps such as requesting devices and parameters
-        or initializing real-time results may also be performed here.
+        Other initialization steps such as requesting devices may also be
+        performed here.
 
-        When the repository is scanned, any requested devices and parameters
-        are set to ``None``."""
-        raise NotImplementedError
-
-    def managers(self):
-        """Returns the device manager and the dataset manager, in this order.
-
-        This is the same order that the constructor takes them, allowing
-        sub-objects to be created with this idiom to pass the environment
-        around: ::
+        When the repository is scanned, any requested devices and arguments
+        are set to ``None``.
 
-            sub_object = SomeLibrary(*self.managers())
-        """
-        return self.__device_mgr, self.__dataset_mgr
+        Leftover positional and keyword arguments from the constructor are
+        forwarded to this method. This is intended for experiments that are
+        only meant to be executed programmatically (not from the GUI)."""
+        pass
 
-    def get_argument(self, key, processor=None, group=None):
+    def get_argument(self, key, processor, group=None):
         """Retrieves and returns the value of an argument.
 
+        This function should only be called from ``build``.
+
         :param key: Name of the argument.
         :param processor: A description of how to process the argument, such
             as instances of ``BooleanValue`` and ``NumberValue``.
@@ -199,22 +210,7 @@ def get_argument(self, key, processor=None, group=None):
         if not self.__in_build:
             raise TypeError("get_argument() should only "
                             "be called from build()")
-        if self.__parent is not None and key not in self.__kwargs:
-            return self.__parent.get_argument(key, processor, group)
-        if processor is None:
-            processor = PYONValue()
-        self.requested_args[key] = processor, group
-        try:
-            argval = self.__kwargs[key]
-        except KeyError:
-            try:
-                return processor.default()
-            except DefaultMissing:
-                if self.__default_arg_none:
-                    return None
-                else:
-                    raise
-        return processor.process(argval)
+        return self.__argument_mgr.get(key, processor, group)
 
     def setattr_argument(self, key, processor=None, group=None):
         """Sets an argument as attribute. The names of the argument and of the
@@ -223,16 +219,10 @@ def setattr_argument(self, key, processor=None, group=None):
 
     def get_device_db(self):
         """Returns the full contents of the device database."""
-        if self.__parent is not None:
-            return self.__parent.get_device_db()
         return self.__device_mgr.get_device_db()
 
     def get_device(self, key):
         """Creates and returns a device driver."""
-        if self.__parent is not None:
-            return self.__parent.get_device(key)
-        if self.__device_mgr is None:
-            raise ValueError("Device manager not present")
         return self.__device_mgr.get(key)
 
     def setattr_device(self, key):
@@ -254,11 +244,6 @@ def set_dataset(self, key, value,
         :param save: the data is saved into the local storage of the current
             run (archived as a HDF5 file).
         """
-        if self.__parent is not None:
-            self.__parent.set_dataset(key, value, broadcast, persist, save)
-            return
-        if self.__dataset_mgr is None:
-            raise ValueError("Dataset manager not present")
         self.__dataset_mgr.set(key, value, broadcast, persist, save)
 
     def mutate_dataset(self, key, index, value):
@@ -267,10 +252,6 @@ def mutate_dataset(self, key, index, value):
 
         If the dataset was created in broadcast mode, the modification is
         immediately transmitted."""
-        if self.__parent is not None:
-            self.__parent.mutate_dataset(key, index, value)
-        if self.__dataset_mgr is None:
-            raise ValueError("Dataset manager not present")
         self.__dataset_mgr.mutate(key, index, value)
 
     def get_dataset(self, key, default=NoDefault):
@@ -283,10 +264,6 @@ def get_dataset(self, key, default=NoDefault):
         If the dataset does not exist, returns the default value. If no default
         is provided, raises ``KeyError``.
         """
-        if self.__parent is not None:
-            return self.__parent.get_dataset(key, default)
-        if self.__dataset_mgr is None:
-            raise ValueError("Dataset manager not present")
         try:
             return self.__dataset_mgr.get(key)
         except KeyError:
@@ -302,7 +279,7 @@ def setattr_dataset(self, key, default=NoDefault):
 
 
 class Experiment:
-    """Base class for experiments.
+    """Base class for top-level experiments.
 
     Deriving from this class enables automatic experiment discovery in
     Python modules.
@@ -348,15 +325,15 @@ def analyze(self):
 
 
 class EnvExperiment(Experiment, HasEnvironment):
-    """Base class for experiments that use the ``HasEnvironment`` environment
-    manager.
+    """Base class for top-level experiments that use the ``HasEnvironment``
+    environment manager.
 
     Most experiment should derive from this class."""
     pass
 
 
 def is_experiment(o):
-    """Checks if a Python object is an instantiable user experiment."""
+    """Checks if a Python object is a top-level experiment class."""
     return (isclass(o)
         and issubclass(o, Experiment)
         and o is not Experiment

diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
@@ -12,7 +12,8 @@
 from artiq.protocols.packed_exceptions import raise_packed_exc
 from artiq.tools import multiline_log_config, file_import
 from artiq.master.worker_db import DeviceManager, DatasetManager
-from artiq.language.environment import is_experiment
+from artiq.language.environment import (is_experiment, TraceArgumentManager,
+                                        ProcessArgumentManager)
 from artiq.language.core import set_watchdog_factory, TerminationRequested
 from artiq.coredevice.core import CompileError, host_only, _render_diagnostic
 from artiq import __version__ as artiq_version
@@ -138,11 +139,11 @@ def examine(device_mgr, dataset_mgr, file):
                 name = exp_class.__doc__.splitlines()[0].strip()
                 if name[-1] == ".":
                     name = name[:-1]
-            exp_inst = exp_class(device_mgr, dataset_mgr,
-                                 default_arg_none=True)
+            argument_mgr = TraceArgumentManager()
+            exp_inst = exp_class((device_mgr, dataset_mgr, argument_mgr))
             arginfo = OrderedDict(
                 (k, (proc.describe(), group))
-                for k, (proc, group) in exp_inst.requested_args.items())
+                for k, (proc, group) in argument_mgr.requested_args.items())
             register_experiment(class_name, name, arginfo)
 
 
@@ -213,8 +214,8 @@ def main():
                                        time.strftime("%H", start_time))
                 os.makedirs(dirname, exist_ok=True)
                 os.chdir(dirname)
-                exp_inst = exp(
-                    device_mgr, dataset_mgr, **expid["arguments"])
+                argument_mgr = ProcessArgumentManager(expid["arguments"])
+                exp_inst = exp((device_mgr, dataset_mgr, argument_mgr))
                 put_object({"action": "completed"})
             elif action == "prepare":
                 exp_inst.prepare()

diff --git a/artiq/test/coredevice/test_portability.py b/artiq/test/coredevice/test_portability.py
@@ -6,19 +6,21 @@
 from artiq.test.hardware_testbench import ExperimentCase
 
 
-def _run_on_host(k_class, **arguments):
-    dmgr = dict()
-    dmgr["core"] = sim_devices.Core(dmgr)
-    k_inst = k_class(dmgr, **arguments)
+def _run_on_host(k_class, *args, **kwargs):
+    device_mgr = dict()
+    device_mgr["core"] = sim_devices.Core(device_mgr)
+
+    k_inst = k_class((device_mgr, None, None),
+                     *args, **kwargs)
     k_inst.run()
     return k_inst
 
 
 class _Primes(EnvExperiment):
-    def build(self):
+    def build(self, output_list, maximum):
         self.setattr_device("core")
-        self.setattr_argument("output_list")
-        self.setattr_argument("maximum")
+        self.output_list = output_list
+        self.maximum = maximum
 
     def _add_output(self, x):
         self.output_list.append(x)
@@ -72,10 +74,10 @@ def run(self):
 
 
 class _PulseLogger(EnvExperiment):
-    def build(self):
+    def build(self, parent_test, name):
         self.setattr_device("core")
-        self.setattr_argument("parent_test")
-        self.setattr_argument("name")
+        self.parent_test = parent_test
+        self.name = name
 
     def _append(self, t, l, f):
         if not hasattr(self.parent_test, "first_timestamp"):
@@ -98,12 +100,12 @@ def pulse(self, f, duration):
 
 
 class _Pulses(EnvExperiment):
-    def build(self):
+    def build(self, output_list):
         self.setattr_device("core")
-        self.setattr_argument("output_list")
+        self.output_list = output_list
 
         for name in "a", "b", "c", "d":
-            pl = _PulseLogger(*self.managers(),
+            pl = _PulseLogger(self,
                               parent_test=self,
                               name=name)
             setattr(self, name, pl)
@@ -125,9 +127,9 @@ class _MyException(Exception):
 
 
 class _Exceptions(EnvExperiment):
-    def build(self):
+    def build(self, trace):
         self.setattr_device("core")
-        self.setattr_argument("trace")
+        self.trace = trace
 
     def _trace(self, i):
         self.trace.append(i)
@@ -172,9 +174,9 @@ def run(self):
 
 
 class _RPCExceptions(EnvExperiment):
-    def build(self):
+    def build(self, catch):
         self.setattr_device("core")
-        self.setattr_argument("catch", PYONValue(False))
+        self.catch = catch
 
         self.success = False
 

diff --git a/artiq/test/coredevice/test_rtio.py b/artiq/test/coredevice/test_rtio.py
@@ -147,11 +147,11 @@ def run(self):
 
 
 class LoopbackCount(EnvExperiment):
-    def build(self):
+    def build(self, npulses):
         self.setattr_device("core")
         self.setattr_device("loop_in")
         self.setattr_device("loop_out")
-        self.setattr_argument("npulses")
+        self.npulses = npulses
 
     def set_count(self, count):
         self.set_dataset("count", count)
@@ -320,9 +320,9 @@ def test_handover(self):
 
 
 class RPCTiming(EnvExperiment):
-    def build(self):
+    def build(self, repeats=100):
         self.setattr_device("core")
-        self.setattr_argument("repeats", PYONValue(100))
+        self.repeats = repeats
 
     def nop(self):
         pass

diff --git a/artiq/test/hardware_testbench.py b/artiq/test/hardware_testbench.py
@@ -108,25 +108,27 @@ def setUp(self):
     def tearDown(self):
         self.device_mgr.close_devices()
 
-    def create(self, cls, **kwargs):
+    def create(self, cls, *args, **kwargs):
         try:
-            exp = cls(self.device_mgr, self.dataset_mgr, **kwargs)
+            exp = cls(
+                (self.device_mgr, self.dataset_mgr, None),
+                *args, **kwargs)
             exp.prepare()
             return exp
         except KeyError as e:
             # skip if ddb does not match requirements
             raise unittest.SkipTest(
                 "device_db entry `{}` not found".format(*e.args))
 
-    def execute(self, cls, **kwargs):
+    def execute(self, cls, *args, **kwargs):
         expid = {
             "file": sys.modules[cls.__module__].__file__,
             "class_name": cls.__name__,
-            "arguments": kwargs
+            "arguments": dict()
         }
         self.device_mgr.virtual_devices["scheduler"].expid = expid
         try:
-            exp = self.create(cls, **kwargs)
+            exp = self.create(cls, *args, **kwargs)
             exp.run()
             exp.analyze()
             return exp