fake_shell.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. import cStringIO
  2. import json
  3. import logging
  4. import sys
  5. import traceback
  6. logging.basicConfig()
  7. logger = logging.getLogger('fake_shell')
  8. sys_stdin = sys.stdin
  9. sys_stdout = sys.stdout
  10. sys_stderr = sys.stderr
  11. fake_stdin = cStringIO.StringIO()
  12. fake_stdout = cStringIO.StringIO()
  13. fake_stderr = cStringIO.StringIO()
  14. sys.stdin = fake_stdin
  15. sys.stdout = fake_stdout
  16. sys.stderr = fake_stderr
  17. global_dict = {}
  18. execution_count = 0
  19. def execute_reply(status, content):
  20. global execution_count
  21. execution_count += 1
  22. return {
  23. 'msg_type': 'execute_reply',
  24. 'content': dict(
  25. content,
  26. status=status,
  27. execution_count=execution_count - 1
  28. )
  29. }
  30. def execute_reply_ok(data):
  31. return execute_reply('ok', {
  32. 'data': data,
  33. })
  34. def execute_reply_error(exc_type, exc_value, tb):
  35. logger.error('execute_reply', exc_info=True)
  36. return execute_reply('error', {
  37. 'ename': unicode(exc_type.__name__),
  38. 'evalue': unicode(exc_value),
  39. 'traceback': traceback.format_exception(exc_type, exc_value, tb, -1),
  40. })
  41. def execute(code):
  42. try:
  43. code = compile(code, '<stdin>', 'single')
  44. exec code in global_dict
  45. except:
  46. return execute_reply_error(*sys.exc_info())
  47. stdout = fake_stdout.getvalue()
  48. stderr = fake_stderr.getvalue()
  49. output = ''
  50. if stdout:
  51. output += stdout
  52. if stderr:
  53. output += stderr
  54. return execute_reply_ok({
  55. 'text/plain': output.rstrip(),
  56. })
  57. def execute_request(content):
  58. try:
  59. code = content['code']
  60. except KeyError:
  61. exc_type, exc_value, tb = sys.exc_info()
  62. return execute_reply_error(exc_type, exc_value, [])
  63. if code.startswith('%'):
  64. parts = code[1:].split(' ', 1)
  65. if len(parts) == 1:
  66. magic, rest = parts[0], ()
  67. else:
  68. magic, rest = parts[0], (parts[1],)
  69. try:
  70. handler = magic_router[magic]
  71. except KeyError:
  72. exc_type, exc_value, tb = sys.exc_info()
  73. return execute_reply_error(exc_type, exc_value, [])
  74. else:
  75. return handler(*rest)
  76. else:
  77. return execute(code)
  78. def table_magic(name):
  79. try:
  80. value = global_dict[name]
  81. except KeyError:
  82. exc_type, exc_value, tb = sys.exc_info()
  83. return execute_reply_error(exc_type, exc_value, [])
  84. max_list_cols = 0
  85. dict_headers = set()
  86. if isinstance(value, list):
  87. for row in value:
  88. if isinstance(row, dict):
  89. dict_headers.update(row.iterkeys())
  90. elif isinstance(row, list):
  91. max_list_cols = max(max_list_cols, len(row))
  92. else:
  93. return execute_reply_error(Exception, 'row is not a list or dict', [])
  94. elif isinstance(value, dict):
  95. dict_headers = value.keys()
  96. value = [value]
  97. else:
  98. return execute_reply_error(Exception, 'value is not a list or dict', [])
  99. headers = [i for i in xrange(max_list_cols)]
  100. dict_header_offset = len(headers)
  101. dict_header_index = {}
  102. for i, key in enumerate(sorted(dict_headers)):
  103. headers.append(key)
  104. dict_header_index[key] = dict_header_offset + i
  105. table = []
  106. for row in value:
  107. table_row = [None] * len(headers)
  108. table.append(table_row)
  109. if isinstance(row, list):
  110. for i, col in enumerate(row):
  111. table_row[i] = col
  112. else:
  113. for key, col in row.iteritems():
  114. i = dict_header_index[key]
  115. table_row[i] = col
  116. return execute_reply_ok({
  117. 'application/vnd.livy.table.v1+json': {
  118. 'headers': headers,
  119. 'data': table,
  120. }
  121. })
  122. magic_router = {
  123. 'table': table_magic,
  124. }
  125. msg_type_router = {
  126. 'execute_request': execute_request,
  127. }
  128. try:
  129. while True:
  130. fake_stdout.truncate(0)
  131. line = sys_stdin.readline()
  132. if line == '':
  133. break
  134. elif line == '\n':
  135. continue
  136. try:
  137. msg = json.loads(line)
  138. except ValueError:
  139. logger.error('failed to parse message', exc_info=True)
  140. continue
  141. try:
  142. msg_type = msg['msg_type']
  143. except KeyError:
  144. logger.error('missing message type', exc_info=True)
  145. continue
  146. try:
  147. content = msg['content']
  148. except KeyError:
  149. logger.error('missing content', exc_info=True)
  150. continue
  151. try:
  152. handler = msg_type_router[msg_type]
  153. except KeyError:
  154. logger.error('unknown message type: %s', msg_type)
  155. continue
  156. response = handler(content)
  157. try:
  158. response = json.dumps(response)
  159. except ValueError, e:
  160. response = json.dumps({
  161. 'msg_type': 'inspect_reply',
  162. 'execution_count': execution_count - 1,
  163. 'content': {
  164. 'status': 'error',
  165. 'ename': 'ValueError',
  166. 'evalue': 'cannot json-ify %s' % response,
  167. 'traceback': [],
  168. }
  169. })
  170. print >> sys_stdout, response
  171. sys_stdout.flush()
  172. finally:
  173. sys.stdin = sys_stdin
  174. sys.stdout = sys_stdout
  175. sys.stderr = sys_stderr