memit: magic memory usage benching for IPython

magic_memit.py

# Author: Vlad Niculae <vlad@vene.ro>
# Makes use of memory_profiler from Fabian Pedregosa
# available at https://github.com/fabianp/memory_profiler

from IPython.core.magic import Magics, line_magic, magics_class


@magics_class
class MemMagics(Magics):
    @line_magic
    def memit(self, line='', setup='pass'):
        """Measure memory usage of a Python statement

        Usage, in line mode:
          %memit [-r<R>] statement

        Options:
        -r<R>: repeat the loop iteration <R> times and take the best result.
        Default: 3

        Examples
        --------
        ::

          In [1]: %run mem_magic.py

          In [2]: import numpy as np

          In [3]: %memit np.zeros(1e7)
          best of 3: 76.402344 MB per loop
          Out[3]: 76.40234375

          In [4]: %memit np.ones(1e6)
          best of 3: 7.820312 MB per loop
          Out[4]: 7.8203125

          In [5]: %memit -r 10 np.empty(1e8)
          best of 10: 0.101562 MB per loop
          Out[5]: 0.1015625

        """

        import multiprocessing as pr

        opts, stmt = self.parse_options(line, 'n:r:tcp:',
                                        posix=False, strict=False)
        repeat = int(getattr(opts, "r", 3))
        ns = self.shell.user_ns

        def _get_usage(stmt, setup, ns, q):
            from memory_profiler import memory_usage as _mu
            try:
                exec setup in ns
                _mu0 = _mu()[0]
                exec stmt in ns
                q.put(_mu()[0] - _mu0)
            except:
                q.put(-1)

        q = pr.Queue()
        for _ in xrange(repeat):
            p = pr.Process(target=_get_usage, args=(stmt, setup, ns, q))
            p.start()
            p.join()

        usages = [q.get() for _ in xrange(repeat)]
        usage = min(usages)
        assert usage != 0
        print u"best of %d: %f MB per loop" % (repeat, usage)
        return usage


if __name__ == '__main__':
    ip = get_ipython()
    ip.register_magics(MemMagics)