artiq_run: refactor, support use from within experiments

You can always (under posix) use #!/usr/bin/env artiq_run as
shebang for experiments and make them executable.
Now, you can also do this (portable):

if __name__ == "__main__":
    from artiq.frontend.artiq_run import run
    run()

to make an experiment executable. The CLI options are all inherited.
Also:

* removed --elf: can be inferred from filename
* did some refactoring and cleanup
* use logging for all messages, except the result printing (use -v to get
parameter changes and dummy scheduler actions)

artiq/frontend/artiq_run.py

@@ -5,32 +5,36 @@ import sys
 import time
 from operator import itemgetter
 from itertools import chain
+import logging
 
 import h5py
 
 from artiq.language.db import *
-from artiq.language.experiment import is_experiment
+from artiq.language.experiment import is_experiment, Experiment
 from artiq.protocols import pyon
 from artiq.protocols.file_db import FlatFileDB
 from artiq.master.worker_db import DBHub, ResultDB
 from artiq.tools import file_import, verbosity_args, init_logger
 
 
-class ELFRunner(AutoDB):
+logger = logging.getLogger(__name__)
+
+
+class ELFRunner(Experiment, AutoDB):
     class DBKeys:
         comm = Device()
+        file = Argument()
 
-    def run(self, filename):
+    def run(self):
-        with open(filename, "rb") as f:
+        with open(self.file, "rb") as f:
-            binary = f.read()
+            self.comm.load(f.read())
-        comm.load(binary)
+        self.comm.run("run")
-        comm.run("run")
+        self.comm.serve(dict(), dict())
-        comm.serve(dict(), dict())
 
 
 class SimpleParamLogger:
     def set(self, timestamp, name, value):
-        print("Parameter change: {} -> {}".format(name, value))
+        logger.info("Parameter change: {} = {}".format(name, value))
 
 
 class DummyWatchdog:
@@ -52,26 +56,26 @@ class DummyScheduler:
     def run_queued(self, run_params):
         rid = self.next_rid
         self.next_rid += 1
-        print("Queuing: {}, RID={}".format(run_params, rid))
+        logger.info("Queuing: %s, RID=%s", run_params, rid)
         return rid
 
     def cancel_queued(self, rid):
-        print("Cancelling RID {}".format(rid))
+        logger.info("Cancelling RID %s", rid)
 
     def run_timed(self, run_params, next_run):
         trid = self.next_trid
         self.next_trid += 1
         next_run_s = time.strftime("%m/%d %H:%M:%S", time.localtime(next_run))
-        print("Timing: {} at {}, TRID={}".format(run_params, next_run_s, trid))
+        logger.info("Timing: %s at %s, TRID=%s", run_params, next_run_s, trid)
         return trid
 
     def cancel_timed(self, trid):
-        print("Cancelling TRID {}".format(trid))
+        logger.info("Cancelling TRID %s", trid)
 
     watchdog = DummyWatchdog
 
 
-def get_argparser():
+def get_argparser(with_file):
     parser = argparse.ArgumentParser(
         description="Local experiment running tool")
 
@@ -81,13 +85,12 @@ def get_argparser():
     parser.add_argument("-p", "--pdb", default="pdb.pyon",
                         help="parameter database file")
 
-    parser.add_argument("-E", "--elf", default=False, action="store_true",
-                        help="run ELF binary")
     parser.add_argument("-e", "--experiment", default=None,
                         help="experiment to run")
     parser.add_argument("-o", "--hdf5", default=None,
                         help="write results to specified HDF5 file"
                              " (default: print them)")
+    if with_file:
         parser.add_argument("file",
                             help="file containing the experiment to run")
     parser.add_argument("arguments", nargs="*",
@@ -99,86 +102,77 @@ def get_argparser():
 def _parse_arguments(arguments):
     d = {}
     for argument in arguments:
-        name, value = argument.split("=")
+        name, eq, value = argument.partition("=")
         d[name] = pyon.decode(value)
     return d
 
 
-def main():
+def _get_experiment(module, experiment=None):
-    args = get_argparser().parse_args()
+    if experiment:
+        return getattr(module, experiment)
+
+    exps = [(k, v) for k, v in module.__dict__.items()
+            if is_experiment(v)]
+    if not exps:
+        logger.error("No experiments in module")
+    if len(exps) > 1:
+        logger.warning("Multiple experiments (%s), using first",
+                       ", ".join(k for (k, v) in exps))
+    return exps[0][1]
+
+
+def _build_experiment(dbh, args):
+    if hasattr(args, "file"):
+        if args.file.endswith(".elf"):
+            if args.arguments:
+                raise ValueError("arguments not supported for ELF kernels")
+            if args.experiment:
+                raise ValueError("experiment-by-name not supported "
+                                 "for ELF kernels")
+            return ELFRunner(dbh, file=args.file)
+        else:
+            module = file_import(args.file)
+        file = args.file
+    else:
+        module = sys.modules["__main__"]
+        file = getattr(module, "__file__")
+    exp = _get_experiment(module, args.experiment)
+    arguments = _parse_arguments(args.arguments)
+    return exp(dbh,
+               scheduler=DummyScheduler(),
+               run_params=dict(file=file,
+                               experiment=args.experiment,
+                               arguments=arguments),
+               **arguments)
+
+
+def run(with_file=False):
+    args = get_argparser(with_file).parse_args()
     init_logger(args)
 
     ddb = FlatFileDB(args.ddb)
     pdb = FlatFileDB(args.pdb)
     pdb.hooks.append(SimpleParamLogger())
     rdb = ResultDB(lambda description: None, lambda mod: None)
-    dbh = DBHub(ddb, pdb, rdb)
-    try:
-        if args.elf:
-            if args.arguments:
-                print("Run arguments are not supported in ELF mode")
-                sys.exit(1)
-            exp_inst = ELFRunner(dbh)
-            rdb.build()
-            exp_inst.run(args.file)
-        else:
-            module = file_import(args.file)
-            if args.experiment is None:
-                exps = [(k, v) for k, v in module.__dict__.items()
-                        if is_experiment(v)]
-                l = len(exps)
-                if l == 0:
-                    print("No experiments found in module")
-                    sys.exit(1)
-                elif l > 1:
-                    print("More than one experiment found in module:")
-                    for k, v in sorted(experiments, key=itemgetter(0)):
-                        if v.__doc__ is None:
-                            print("    {}".format(k))
-                        else:
-                            print("    {} ({})".format(
-                                k, v.__doc__.splitlines()[0].strip()))
-                    print("Use -u to specify which experiment to use.")
-                    sys.exit(1)
-                else:
-                    exp = exps[0][1]
-            else:
-                exp = getattr(module, args.experiment)
 
-            try:
+    with DBHub(ddb, pdb, rdb) as dbh:
-                arguments = _parse_arguments(args.arguments)
+        exp_inst = _build_experiment(dbh, args)
-            except:
-                print("Failed to parse run arguments")
-                sys.exit(1)
-
-            run_params = {
-                "file": args.file,
-                "experiment": args.experiment,
-                "arguments": arguments
-            }
-            exp_inst = exp(dbh,
-                           scheduler=DummyScheduler(),
-                           run_params=run_params,
-                           **run_params["arguments"])
         rdb.build()
         exp_inst.run()
         exp_inst.analyze()
 
     if args.hdf5 is not None:
-                f = h5py.File(args.hdf5, "w")
+        with h5py.File(args.hdf5, "w") as f:
-                try:
             rdb.write_hdf5(f)
-                finally:
+    elif rdb.data.read or rdb.realtime_data.read:
-                    f.close()
+        r = chain(rdb.realtime_data.read.items(), rdb.data.read.items())
-            else:
+        for k, v in sorted(r, key=itemgetter(0)):
-                if rdb.data.read or rdb.realtime_data.read:
-                    print("Results:")
-                    for k, v in sorted(chain(rdb.realtime_data.read.items(),
-                                             rdb.data.read.items()),
-                                       key=itemgetter(0)):
             print("{}: {}".format(k, v))
-    finally:
+
-        dbh.close_devices()
+
+def main():
+    return run(with_file=True)
+
 
 if __name__ == "__main__":
     main()