progress_stream.py 2.2 KB

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