progress_stream.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. from __future__ import absolute_import
  2. from __future__ import unicode_literals
  3. from compose import utils
  4. class StreamOutputError(Exception):
  5. pass
  6. def write_to_stream(s, stream):
  7. try:
  8. stream.write(s)
  9. except UnicodeEncodeError:
  10. encoding = getattr(stream, 'encoding', 'ascii')
  11. stream.write(s.encode(encoding, errors='replace').decode(encoding))
  12. def stream_output(output, stream):
  13. is_terminal = hasattr(stream, 'isatty') and stream.isatty()
  14. stream = utils.get_output_stream(stream)
  15. lines = {}
  16. diff = 0
  17. for event in utils.json_stream(output):
  18. yield event
  19. is_progress_event = 'progress' in event or 'progressDetail' in event
  20. if not is_progress_event:
  21. print_output_event(event, stream, is_terminal)
  22. stream.flush()
  23. continue
  24. if not is_terminal:
  25. continue
  26. # if it's a progress event and we have a terminal, then display the progress bars
  27. image_id = event.get('id')
  28. if not image_id:
  29. continue
  30. if image_id not in lines:
  31. lines[image_id] = len(lines)
  32. write_to_stream("\n", stream)
  33. diff = len(lines) - lines[image_id]
  34. # move cursor up `diff` rows
  35. write_to_stream("%c[%dA" % (27, diff), stream)
  36. print_output_event(event, stream, is_terminal)
  37. if 'id' in event:
  38. # move cursor back down
  39. write_to_stream("%c[%dB" % (27, diff), stream)
  40. stream.flush()
  41. def print_output_event(event, stream, is_terminal):
  42. if 'errorDetail' in event:
  43. raise StreamOutputError(event['errorDetail']['message'])
  44. terminator = ''
  45. if is_terminal and 'stream' not in event:
  46. # erase current line
  47. write_to_stream("%c[2K\r" % 27, stream)
  48. terminator = "\r"
  49. elif 'progressDetail' in event:
  50. return
  51. if 'time' in event:
  52. write_to_stream("[%s] " % event['time'], stream)
  53. if 'id' in event:
  54. write_to_stream("%s: " % event['id'], stream)
  55. if 'from' in event:
  56. write_to_stream("(from %s) " % event['from'], stream)
  57. status = event.get('status', '')
  58. if 'progress' in event:
  59. write_to_stream("%s %s%s" % (status, event['progress'], terminator), stream)
  60. elif 'progressDetail' in event:
  61. detail = event['progressDetail']
  62. total = detail.get('total')
  63. if 'current' in detail and total:
  64. percentage = float(detail['current']) / float(total) * 100
  65. write_to_stream('%s (%.1f%%)%s' % (status, percentage, terminator), stream)
  66. else:
  67. write_to_stream('%s%s' % (status, terminator), stream)
  68. elif 'stream' in event:
  69. write_to_stream("%s%s" % (event['stream'], terminator), stream)
  70. else:
  71. write_to_stream("%s%s\n" % (status, terminator), stream)
  72. def get_digest_from_pull(events):
  73. digest = None
  74. for event in events:
  75. status = event.get('status')
  76. if not status or 'Digest' not in status:
  77. continue
  78. else:
  79. digest = status.split(':', 1)[1].strip()
  80. return digest
  81. def get_digest_from_push(events):
  82. for event in events:
  83. digest = event.get('aux', {}).get('Digest')
  84. if digest:
  85. return digest
  86. return None