Dirty Memory Usage Snippet in Python

Today I had a memory issue in my code, with memory usage increasing exponentially faster than it should have. I needed to find where (i.e. in what objects) that memory was consumed. And ran into one of the discomforts in using Sage instead of bare Python: I could not use directly any of the memory profiling tools designed for Python, because Sage has its own Python install and I didn’t want to spend time understanding how to tweak it. So instead, I adapted a simple code snippet from this page to follow memory usage line by line. Here is the code:

def get_mem_usage(read_scale='MB'):
    import os
    scales = {'KB': 1024., 'MB': 1024.**2}
    handle = open('/proc/%d/status' % os.getpid())
    fstr = handle.read()
    handle.close()
    i = fstr.index('VmRSS')
    v = fstr[i:].split(None, 3)
    assert len(v) >= 3
    return float(v[1]) * scales[v[2].upper()] / scales[read_scale]

It’s even dirtier than the original and works only on Linux, but one can just copy-paste it anywhere in their code and use it right away. And that’s exactly what I wanted to do here. Just sharing the tip, hoping it helps :-)