import gc
import random
import os
import time

try:
    long
except NameError:
    # py3k
    long = int

try:
    xrange
except NameError:
    # py3k
    xrange = range

### http://code.activestate.com/recipes/286222/ ###

_proc_status = '/proc/%d/status' % os.getpid()

_scale = {'kB': 1024.0, 'mB': 1024.0*1024.0,
          'KB': 1024.0, 'MB': 1024.0*1024.0}

def _VmB(VmKey):
    '''Private.
    '''
    global _proc_status, _scale
     # get pseudo file  /proc/<pid>/status
    try:
        t = open(_proc_status)
        v = t.read()
        t.close()
    except:
        return 0.0  # non-Linux?
     # get VmKey line e.g. 'VmRSS:  9999  kB\n ...'
    i = v.index(VmKey)
    v = v[i:].split(None, 3)  # whitespace
    if len(v) < 3:
        return 0.0  # invalid format?
     # convert Vm value to bytes
    return float(v[1]) * _scale[v[2]]


def memory(since=0.0):
    '''Return memory usage in bytes.
    '''
    return _VmB('VmSize:') - since


def resident(since=0.0):
    '''Return resident memory usage in bytes.
    '''
    return _VmB('VmRSS:') - since


def stacksize(since=0.0):
    '''Return stack size in bytes.
    '''
    return _VmB('VmStk:') - since

### http://code.activestate.com/recipes/286222/ end ###

def getmem(msg=None):
    mem = memory()
    print("Memory used (%s): %iMb" % (msg, round(mem / (1024*1024))))
    return mem

def allocate(n=2000000, t=long):
    x = [t(i) for i in xrange(n)]
    getmem("peak")

def memtest(n=2000000, t=long):
    print("")
    print("memtest %i %s" % (n, t.__name__))
    print("====================")
    a = xrange(n)
    x = [t(i) for i in a]
    getmem("peak")
    del a, x
    #allocate(n, t)    
    memunfreed()

mem_init = getmem("begin")

def _dict(i):
    """Dictionary constructor."""
    return {str(i): i}

def _tuple(i):
    """Tuple constructor."""
    return (i,)

def _set(i):
    """Tuple constructor."""
    return set([i])

def memunfreed():
    global mem_init
    gc.collect()
    mem_unfreed = getmem("end") - mem_init
    print("Unfreed memory: %iMb" % round(mem_unfreed / (1024*1024)))

class Test(object):
    """No custom attributes."""
    def __init__(self, x):
        pass

class Test2(object):
    """Use slots for attributes."""
    __slots__ = ["x"]
    def __init__(self, x):
        self.x = None

class Test3(object):
    """Assign an attribute."""
    def __init__(self, x):
        self.x = None

memtest(t=long)
memtest(t=str)
memtest(t=float)
memtest(t=int)
memtest(t=Test)
memtest(t=Test2)
memtest(t=_tuple)
memtest(t=_set)
memtest(t=_dict)
memtest(t=Test3)