progress_stream.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  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. event = json.loads(chunk)
  15. all_events.append(event)
  16. if 'progress' in event or 'progressDetail' in event:
  17. image_id = event.get('id')
  18. if not image_id:
  19. continue
  20. if image_id in lines:
  21. diff = len(lines) - lines[image_id]
  22. else:
  23. lines[image_id] = len(lines)
  24. stream.write("\n")
  25. diff = 0
  26. if is_terminal:
  27. # move cursor up `diff` rows
  28. stream.write("%c[%dA" % (27, diff))
  29. print_output_event(event, stream, is_terminal)
  30. if 'id' in event and is_terminal:
  31. # move cursor back down
  32. stream.write("%c[%dB" % (27, diff))
  33. stream.flush()
  34. return all_events
  35. def print_output_event(event, stream, is_terminal):
  36. if 'errorDetail' in event:
  37. raise StreamOutputError(event['errorDetail']['message'])
  38. terminator = ''
  39. if is_terminal and 'stream' not in event:
  40. # erase current line
  41. stream.write("%c[2K\r" % 27)
  42. terminator = "\r"
  43. elif 'progressDetail' in event:
  44. return
  45. if 'time' in event:
  46. stream.write("[%s] " % event['time'])
  47. if 'id' in event:
  48. stream.write("%s: " % event['id'])
  49. if 'from' in event:
  50. stream.write("(from %s) " % event['from'])
  51. status = event.get('status', '')
  52. if 'progress' in event:
  53. stream.write("%s %s%s" % (status, event['progress'], terminator))
  54. elif 'progressDetail' in event:
  55. detail = event['progressDetail']
  56. total = detail.get('total')
  57. if 'current' in detail and total:
  58. percentage = float(detail['current']) / float(total) * 100
  59. stream.write('%s (%.1f%%)%s' % (status, percentage, terminator))
  60. else:
  61. stream.write('%s%s' % (status, terminator))
  62. elif 'stream' in event:
  63. stream.write("%s%s" % (event['stream'], terminator))
  64. else:
  65. stream.write("%s%s\n" % (status, terminator))