from __future__ import print_function, unicode_literals
from argparse import ArgumentParser
from datetime import datetime
from errno import EPIPE
import sys

from .kinesis_logs_reader import KinesisLogsReader


def print_stream(stream_name, start_time, stop_after):
    reader = KinesisLogsReader(stream_name, start_time=start_time)
    for i, fields in enumerate(reader, 1):
        if i == 1:
            keys = sorted(fields.keys())
            print(*keys, sep='\t')
        try:
            print(*[fields[k] for k in keys], sep='\t')
        except IOError as e:
            if e.errno == EPIPE:
                pass
        if i == stop_after:
            break


def main(argv=None):
    argv = argv or sys.argv[1:]

    argument_parser = ArgumentParser('Read logs from AWS Kinesis')
    argument_parser.add_argument(
        'stream_name', type=str, help='Name of the Kinesis stream to read'
    )
    argument_parser.add_argument(
        '--start-time',
        type=str,
        help='Time from which to start reading (default: beginning of stream)'
    )
    argument_parser.add_argument(
        '--count',
        type=int,
        default=0,
        help='Number of records to return (default: 0, which is "no limit")'
    )
    argument_parser.add_argument(
        '--time-format',
        type=str,
        default='%Y-%m-%d %H:%M:%S',
        help='Format string for the --start-time argument',
    )
    args = argument_parser.parse_args(argv)

    if args.start_time:
        start_time = datetime.strptime(args.start_time, args.time_format)
    else:
        start_time = None

    print_stream(args.stream_name, start_time, args.count)


if __name__ == '__main__':
    main()
    sys.stdout.close()

@@ -139,16 +139,44 @@
 t EPIPE%0A
+from itertools import chain%0A
 import s
@@ -358,78 +358,92 @@
 me)%0A
+%0A
     
-for i, fields in enumerate(reader, 1):%0A        if i == 1:%0A        
+# Peek into the first row to determine the header%0A    first_row = next(reader)%0A
     
@@ -458,20 +458,23 @@
 orted(fi
-elds
+rst_row
 .keys())
@@ -474,24 +474,16 @@
 keys())%0A
-        
     prin
@@ -492,32 +492,186 @@
 *keys, sep='%5Ct')
+%0A%0A    # Join the first row with the rest of the rows and print them%0A    iterable = chain(%5Bfirst_row%5D, reader)%0A    for i, fields in enumerate(iterable, 1):
 %0A        try:%0A  

Check for first row just once