progress_stream.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. import codecs
  2. import json
  3. import six
  4. class StreamOutputError(Exception):
  5. pass
  6. def stream_output(output, stream):
  7. is_terminal = hasattr(stream, 'isatty') and stream.isatty()
  8. if not six.PY3:
  9. stream = codecs.getwriter('utf-8')(stream)
  10. all_events = []
  11. lines = {}
  12. diff = 0
  13. for chunk in output:
  14. if six.PY3 and not isinstance(chunk, str):
  15. chunk = chunk.decode('utf-8')
  16. event = json.loads(chunk)
  17. all_events.append(event)
  18. if 'progress' in event or 'progressDetail' in event:
  19. image_id = event.get('id')
  20. if not image_id:
  21. continue
  22. if image_id in lines:
  23. diff = len(lines) - lines[image_id]
  24. else:
  25. lines[image_id] = len(lines)
  26. stream.write("\n")
  27. diff = 0
  28. if is_terminal:
  29. # move cursor up `diff` rows
  30. stream.write("%c[%dA" % (27, diff))
  31. print_output_event(event, stream, is_terminal)
  32. if 'id' in event and is_terminal:
  33. # move cursor back down
  34. stream.write("%c[%dB" % (27, diff))
  35. stream.flush()
  36. return all_events
  37. def print_output_event(event, stream, is_terminal):
  38. if 'errorDetail' in event:
  39. raise StreamOutputError(event['errorDetail']['message'])
  40. terminator = ''
  41. if is_terminal and 'stream' not in event:
  42. # erase current line
  43. stream.write("%c[2K\r" % 27)
  44. terminator = "\r"
  45. elif 'progressDetail' in event:
  46. return
  47. if 'time' in event:
  48. stream.write("[%s] " % event['time'])
  49. if 'id' in event:
  50. stream.write("%s: " % event['id'])
  51. if 'from' in event:
  52. stream.write("(from %s) " % event['from'])
  53. status = event.get('status', '')
  54. if 'progress' in event:
  55. stream.write("%s %s%s" % (status, event['progress'], terminator))
  56. elif 'progressDetail' in event:
  57. detail = event['progressDetail']
  58. total = detail.get('total')
  59. if 'current' in detail and total:
  60. percentage = float(detail['current']) / float(total) * 100
  61. stream.write('%s (%.1f%%)%s' % (status, percentage, terminator))
  62. else:
  63. stream.write('%s%s' % (status, terminator))
  64. elif 'stream' in event:
  65. stream.write("%s%s" % (event['stream'], terminator))
  66. else:
  67. stream.write("%s%s\n" % (status, terminator))