#!/usr/bin/env python

"""Check the CPU and memory usage of some string concatenation approaches.

See:
http://etalog.blogspot.com/2009/07/faster-string-concatenation-in-python.html
"""

import os
import sys
import subprocess
from array import array
from cStringIO import StringIO
from time import clock

LOOP_CT = 1000000
NUMS = xrange(LOOP_CT)
PID = str(os.getpid())

def memsize():
    out = StringIO()
    ps = subprocess.Popen(['ps', 'h', '-o', '%z', PID], stdout=subprocess.PIPE)
    memsize = int(ps.communicate()[0])
    return memsize

# ---- strcat implementations ----

def method1():
    """1. string addition"""
    start = clock()
    out_str = ''
    for num in NUMS:
        out_str += str(num)
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method2():
    """2. %-interpolation"""
    start = clock()
    out_str = ('%d'*LOOP_CT) % tuple(NUMS)
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method3():
    """3. array object"""
    start = clock()
    char_array = array('c')
    for num in NUMS:
        char_array.fromstring(str(num))
    out_str = char_array.tostring()
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method4():
    """4. cStringIO object"""
    start = clock()
    file_str = StringIO()
    for num in NUMS:
        file_str.write(str(num))
    out_str = file_str.getvalue()
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method5():
    """5. join + for loop"""
    start = clock()
    str_list = []
    for num in NUMS:
        str_list.append(str(num))
    out_str = ''.join(str_list)
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method6():
    """6. join + list comp"""
    start = clock()
    out_str = ''.join([str(num) for num in NUMS])
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method7():
    """7. join + gen expr"""
    start = clock()
    out_str = ''.join(str(num) for num in NUMS)
    cpu = clock() - start
    return (out_str, cpu, memsize())

def method8():
    """8. join + str map"""
    start = clock()
    out_str = ''.join(map(str, NUMS))
    cpu = clock() - start
    return (out_str, cpu, memsize())

# ----------------------------------------------------------

def call_method(choice):
    try:
        func = globals()['method'+choice]
    except KeyError:
        print 'Hey, %s is not an option' % choice
        sys.exit(1)
    outstr, cpu, memsize = func()
    print func.__doc__, '\t',
    print "CPU (s): %.2f\tMem (K): %.1f" % (cpu, memsize/1024.0)

if __name__ == '__main__':
    assert len(sys.argv) == 2, "usage: strcat.py method_number"
    call_method(sys.argv[1])