import timeit
import __main__
from cStringIO import StringIO
from itertools import chain, repeat, takewhile
from copy import copy


N = 1000 #Number of executions
L = 1000 #Length of initial list
stringio = "StringIO('1'*L+'Z111')"


def test_f_iter_plus_Z():
    stringio_loc = eval(stringio)
    i = takewhile(lambda x: x!='',iter(lambda: stringio_loc.read(1),'Z'))
    buf = ''.join(i) + 'Z'


def test_f_iter_chain():
    stringio_loc = eval(stringio)
    i = takewhile(lambda x: x!='',iter(lambda: stringio_loc.read(1),'Z'))
    i = chain(i,repeat('Z',1))
    buf = ''.join(i)


def take_until_included(stringio_loc):
    for s in iter(lambda: stringio_loc.read(1),''):
        yield s
        if s=='Z':
            return


def test_f_generator():
    stringio_loc = eval(stringio)
    i = take_until_included(stringio_loc)
    buf = ''.join(i)


def test_f_chunks():
    stringio_loc = eval(stringio)
    datalist = []
    for chunk in iter(lambda: stringio_loc.read(256),''):
        if 'Z' not in chunk:
            datalist.append(chunk)
        else:
            datalist.append(chunk[:chunk.find('Z')+1])
            break
    return ''.join(datalist)
   
 
def test_f_chunks_generator():
    stringio_loc = eval(stringio)
    def gen(): 
        for chunk in iter(lambda: stringio_loc.read(256),''):
            if 'Z' not in chunk:
                yield chunk
            else:
                yield chunk[:chunk.find('Z')+1]
                break
    return ''.join(gen())  

   
names = dir(__main__)
for name in names:
    attr = getattr(__main__,name)
    if hasattr(attr,'__call__'):
        if name.startswith('test_f_'):
            init_name = name + '_init'
            if hasattr(__main__, init_name):
                init = getattr(__main__,init_name)
            else:
                init = 'from __main__ import {name}; from cStringIO import StringIO;'.format(name=name)
            t = timeit.Timer(stmt='{name}()'.format(name=name),
                             setup=init)

            time = t.timeit(N)
            print('{name}: {time}'.format(name=name,time=time))