artiq/artiq/tools.py

171 lines
4.3 KiB
Python
Raw Normal View History

from operator import itemgetter
import importlib.machinery
import linecache
2015-01-28 21:44:15 +08:00
import logging
2015-02-18 23:52:31 +08:00
import sys
import asyncio
import time
import os.path
from artiq.language.experiment import is_experiment
from artiq.protocols import pyon
def parse_arguments(arguments):
d = {}
for argument in arguments:
name, eq, value = argument.partition("=")
d[name] = pyon.decode(value)
return d
2015-02-27 15:17:11 +08:00
def format_arguments(arguments):
fmtargs = []
for k, v in sorted(arguments.items(), key=itemgetter(0)):
2015-02-27 15:17:11 +08:00
fmtargs.append(k + "=" + repr(v))
if fmtargs:
2015-02-27 15:17:11 +08:00
return ", ".join(fmtargs)
else:
2015-02-27 15:17:11 +08:00
return ""
def file_import(filename):
linecache.checkcache(filename)
modname = filename
i = modname.rfind("/")
if i > 0:
modname = modname[i+1:]
i = modname.find(".")
if i > 0:
modname = modname[:i]
modname = "file_import_" + modname
path = os.path.dirname(os.path.realpath(filename))
sys.path.insert(0, path)
loader = importlib.machinery.SourceFileLoader(modname, filename)
module = loader.load_module()
sys.path.remove(path)
return module
2015-01-28 21:44:15 +08:00
def get_experiment(module, experiment=None):
if experiment:
return getattr(module, experiment)
exps = [(k, v) for k, v in module.__dict__.items()
if is_experiment(v)]
if not exps:
raise ValueError("No experiments in module")
if len(exps) > 1:
raise ValueError("More than one experiment found in module")
return exps[0][1]
2015-01-28 21:44:15 +08:00
def verbosity_args(parser):
group = parser.add_argument_group("verbosity")
2015-02-04 19:09:37 +08:00
group.add_argument("-v", "--verbose", default=0, action="count",
help="increase logging level")
group.add_argument("-q", "--quiet", default=0, action="count",
help="decrease logging level")
2015-01-28 21:44:15 +08:00
def simple_network_args(parser, default_port):
group = parser.add_argument_group("network")
group.add_argument("--bind", default="::1",
help="hostname or IP address to bind to")
group.add_argument("-p", "--port", default=default_port, type=int,
2015-03-24 03:02:14 +08:00
help="TCP port to listen to (default: %(default)d)")
2015-01-28 21:44:15 +08:00
def init_logger(args):
logging.basicConfig(level=logging.WARNING + args.quiet*10 - args.verbose*10)
@asyncio.coroutine
def asyncio_process_wait_timeout(process, timeout):
# In Python < 3.5, asyncio.wait_for(process.wait(), ...
# causes a futures.InvalidStateError inside asyncio if and when the
# process terminates after the timeout.
# Work around this problem.
@asyncio.coroutine
def process_wait_returncode_timeout():
while True:
if process.returncode is not None:
break
yield from asyncio.sleep(0.1)
yield from asyncio.wait_for(process_wait_returncode_timeout(),
timeout=timeout)
2015-05-17 16:11:00 +08:00
@asyncio.coroutine
def asyncio_process_wait(process):
r = True
while r:
f, p = yield from asyncio.wait([process.stdout.read(1024)])
r = f.pop().result()
2015-05-17 16:11:00 +08:00
@asyncio.coroutine
def asyncio_wait_or_cancel(fs, **kwargs):
fs = [asyncio.async(f) for f in fs]
try:
d, p = yield from asyncio.wait(fs, **kwargs)
except:
for f in fs:
f.cancel()
raise
for f in p:
f.cancel()
yield from asyncio.wait([f])
2015-05-17 16:11:00 +08:00
return fs
def asyncio_queue_peek(q):
"""Like q.get_nowait(), but does not remove the item from the queue."""
if q._queue:
return q._queue[0]
else:
raise asyncio.QueueEmpty
2015-05-28 17:20:58 +08:00
class TaskObject:
def start(self):
self.task = asyncio.async(self._do())
@asyncio.coroutine
def stop(self):
self.task.cancel()
yield from asyncio.wait([self.task])
del self.task
@asyncio.coroutine
def _do(self):
raise NotImplementedError
2015-05-28 17:20:58 +08:00
class WaitSet:
def __init__(self):
self._s = set()
self._ev = asyncio.Event()
def _update_ev(self):
if self._s:
self._ev.clear()
else:
self._ev.set()
def add(self, e):
self._s.add(e)
self._update_ev()
def discard(self, e):
self._s.discard(e)
self._update_ev()
@asyncio.coroutine
def wait_empty(self):
yield from self._ev.wait()