progress_stream.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. import json
  2. import os
  3. import codecs
  4. class StreamOutputError(Exception):
  5. pass
  6. def stream_output(output, stream):
  7. is_terminal = hasattr(stream, 'fileno') and os.isatty(stream.fileno())
  8. stream = codecs.getwriter('utf-8')(stream)
  9. all_events = []
  10. lines = {}
  11. diff = 0
  12. for chunk in output:
  13. event = json.loads(chunk)
  14. all_events.append(event)
  15. if 'progress' in event or 'progressDetail' in event:
  16. image_id = event.get('id')
  17. if not image_id:
  18. continue
  19. if image_id in lines:
  20. diff = len(lines) - lines[image_id]
  21. else:
  22. lines[image_id] = len(lines)
  23. stream.write("\n")
  24. diff = 0
  25. if is_terminal:
  26. # move cursor up `diff` rows
  27. stream.write("%c[%dA" % (27, diff))
  28. print_output_event(event, stream, is_terminal)
  29. if 'id' in event and is_terminal:
  30. # move cursor back down
  31. stream.write("%c[%dB" % (27, diff))
  32. stream.flush()
  33. return all_events
  34. def print_output_event(event, stream, is_terminal):
  35. if 'errorDetail' in event:
  36. raise StreamOutputError(event['errorDetail']['message'])
  37. terminator = ''
  38. if is_terminal and 'stream' not in event:
  39. # erase current line
  40. stream.write("%c[2K\r" % 27)
  41. terminator = "\r"
  42. pass
  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))