"""IPython magic"""
import io
import os
from tempfile import mkstemp
import pyq
from getopt import getopt
import sys
StringIO = io.StringIO # LGTM.com does not like double import
STD_STREAM = [sys.stdin, sys.stdout, sys.stderr]
try:
string_types = (str, unicode)
except NameError:
string_types = (str,)
Q_NONE = pyq.q('::')
def logical_lines(lines):
"""Merge lines into chunks according to q rules"""
if isinstance(lines, string_types):
lines = StringIO(lines)
buf = []
for line in lines:
if buf and not line.startswith(' '):
chunk = ''.join(buf).strip()
if chunk:
yield chunk
buf[:] = []
buf.append(line)
chunk = ''.join(buf).strip()
if chunk:
yield chunk
def _forward_outputs(outs):
for fd in (1, 2):
if fd in outs:
os.lseek(fd, 0, os.SEEK_SET)
with io.open(fd, closefd=False) as f:
STD_STREAM[fd].writelines(f)
os.ftruncate(fd, 0)
def q(line, cell=None, _ns=None):
"""Run q code.
Options:
-l (dir|script) - pre-load database or script
-h host:port - execute on the given host
-o var - send output to a variable named var.
-i var1,..,varN - input variables
-1/-2 - redirect stdout/stderr
"""
if cell is None:
return pyq.q(line)
if _ns is None:
_ns = vars(sys.modules['__main__'])
input = output = None
preload = []
outs = {}
try:
h = pyq.q('0i')
if line:
for opt, value in getopt(line.split(), "h:l:o:i:12")[0]:
if opt == '-l':
preload.append(value)
elif opt == '-h':
h = pyq.K(str(':' + value))
elif opt == '-o':
output = str(value) # (see #673)
elif opt == '-i':
input = str(value).split(',')
elif opt in ('-1', '-2'):
outs[int(opt[1])] = None
if outs:
if int(h) != 0:
raise ValueError("Cannot redirect remote std stream")
for fd in outs:
tmpfd, tmpfile = mkstemp()
try:
pyq.q(r'\%d %s' % (fd, tmpfile))
finally:
os.unlink(tmpfile)
os.close(tmpfd)
r = None
for script in preload:
h(pyq.kp(r"\l " + script))
if input is not None:
for chunk in logical_lines(cell):
func = "{[%s]%s}" % (';'.join(input), chunk)
args = tuple(_ns[i] for i in input)
if r != Q_NONE:
r.show()
r = h((pyq.kp(func),) + args)
if outs:
_forward_outputs(outs)
else:
for chunk in logical_lines(cell):
if r != Q_NONE:
r.show()
r = h(pyq.kp(chunk))
if outs:
_forward_outputs(outs)
except pyq.kerr as e:
print("'%s" % e)
else:
if output is not None:
if output.startswith('q.'):
pyq.q('@[`.;;:;]', output[2:], r)
else:
_ns[output] = r
else:
if r != Q_NONE:
return r
def _q_formatter(x, p, _):
x_show = x.show(output=str)
p.text(x_show.strip())
def load_ipython_extension(ipython):
"""Register %q and %%q magics and pretty display for K objects"""
ipython.register_magic_function(q, 'line_cell')
fmr = ipython.display_formatter.formatters['text/plain']
fmr.for_type(pyq.K, _q_formatter)
def unload_ipython_extension(ipython):
pass
