marsam · June 30, 2012 19:55 · Jun 30, 2012
diff --git a/magic_memit.py b/magic_memit.py
@@ -0,0 +1,75 @@
+# 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)