Merged in eteq/line_profiler (pull request rkern#4)

Author: rkern

line_profiler.py

@@ -10,7 +10,6 @@
     from cStringIO import StringIO
 except ImportError:
     from io import StringIO
-
 import inspect
 import linecache
 import optparse
@@ -50,6 +49,7 @@ def is_generator(f):
     isgen = (f.__code__.co_flags & CO_GENERATOR) != 0
     return isgen
 
+
 class LineProfiler(CLineProfiler):
     """ A profiler that records the execution times of individual lines.
     """
@@ -111,11 +111,11 @@ def dump_stats(self, filename):
         with open(filename, 'wb') as f:
             pickle.dump(lstats, f, pickle.HIGHEST_PROTOCOL)
 
-    def print_stats(self, stream=None):
+    def print_stats(self, stream=None, stripzeros=False):
         """ Show the gathered statistics.
         """
         lstats = self.get_stats()
-        show_text(lstats.timings, lstats.unit, stream=stream)
+        show_text(lstats.timings, lstats.unit, stream=stream, stripzeros=stripzeros)
 
     def run(self, cmd):
         """ Profile a single executable statment in the main namespace.
@@ -143,24 +143,46 @@ def runcall(self, func, *args, **kw):
         finally:
             self.disable_by_count()
 
+    def add_module(self, mod):
+        """ Add all the functions in a module and its classes.
+        """
+        from inspect import isclass, isfunction
+
+        nfuncsadded = 0
+        for item in mod.__dict__.values():
+            if isclass(item):
+                for k, v in item.__dict__.items():
+                    if isfunction(v):
+                        self.add_function(v)
+                        nfuncsadded += 1
+            elif isfunction(item):
+                self.add_function(item)
+                nfuncsadded += 1
 
-def show_func(filename, start_lineno, func_name, timings, unit, stream=None):
+        return nfuncsadded
+
+
+def show_func(filename, start_lineno, func_name, timings, unit, stream=None, stripzeros=False):
     """ Show results for a single function.
     """
     if stream is None:
         stream = sys.stdout
 
-    stream.write("File: %s\n" % filename)
-    stream.write("Function: %s at line %s\n" % (func_name, start_lineno))
     template = '%6s %9s %12s %8s %8s  %-s'
     d = {}
     total_time = 0.0
     linenos = []
     for lineno, nhits, time in timings:
         total_time += time
         linenos.append(lineno)
+
+    if stripzeros and total_time == 0:
+        return
+
     stream.write("Total time: %g s\n" % (total_time * unit))
     if os.path.exists(filename) or filename.startswith("<ipython-input-"):
+        stream.write("File: %s\n" % filename)
+        stream.write("Function: %s at line %s\n" % (func_name, start_lineno))
         if os.path.exists(filename):
             # Clear the cache to ensure that we get up-to-date results.
             linecache.clearcache()
@@ -195,15 +217,15 @@ def show_func(filename, start_lineno, func_name, timings, unit, stream=None):
         stream.write("\n")
     stream.write("\n")
 
-def show_text(stats, unit, stream=None):
+def show_text(stats, unit, stream=None, stripzeros=False):
     """ Show text for the given timings.
     """
     if stream is None:
         stream = sys.stdout
 
     stream.write('Timer unit: %g s\n\n' % unit)
     for (fn, lineno, name), timings in sorted(stats.items()):
-        show_func(fn, lineno, name, stats[fn, lineno, name], unit, stream=stream)
+        show_func(fn, lineno, name, stats[fn, lineno, name], unit, stream=stream, stripzeros=stripzeros)
 
 # A %lprun magic for IPython.
 def magic_lprun(self, parameter_s=''):
@@ -228,7 +250,9 @@ def magic_lprun(self, parameter_s=''):
     in the interpreter at the In[] prompt or via %run currently cannot be
     displayed.  Write these functions out to a separate file and import them.
 
-    One or more -f options are required to get any useful results.
+    -m <module>: Get all the functions/methods in a module
+
+    One or more -f or -m options are required to get any useful results.
 
     -D <filename>: dump the raw statistics out to a pickle file on disk. The
     usual extension for this is ".lprof". These statistics may be viewed later
@@ -238,6 +262,8 @@ def magic_lprun(self, parameter_s=''):
     out to a text file.
 
     -r: return the LineProfiler object after it has completed profiling.
+
+    -s: strip out all entries from the print-out that have zeros.
     """
     # Local imports to avoid hard dependency.
     from distutils.version import LooseVersion
@@ -253,9 +279,9 @@ def magic_lprun(self, parameter_s=''):
         from IPython.core.error import UsageError
 
     # Escape quote markers.
-    opts_def = Struct(D=[''], T=[''], f=[])
+    opts_def = Struct(D=[''], T=[''], f=[], m=[])
     parameter_s = parameter_s.replace('"', r'\"').replace("'", r"\'")
-    opts, arg_str = self.parse_options(parameter_s, 'rf:D:T:', list_all=True)
+    opts, arg_str = self.parse_options(parameter_s, 'rsf:m:D:T:', list_all=True)
     opts.merge(opts_def)
 
     global_ns = self.shell.user_global_ns
@@ -272,6 +298,15 @@ def magic_lprun(self, parameter_s=''):
 
     profile = LineProfiler(*funcs)
 
+    # Get the modules, too
+    for modname in opts.m:
+        try:
+            mod = __import__(modname, fromlist=[''])
+            profile.add_module(mod)
+        except Exception as e:
+            raise UsageError('Could not find module %r.\n%s: %s' % (name,
+                e.__class__.__name__, e))
+
     # Add the profiler to the builtins for @profile.
     if PY3:
         import builtins
@@ -301,7 +336,7 @@ def magic_lprun(self, parameter_s=''):
 
     # Trap text output.
     stdout_trap = StringIO()
-    profile.print_stats(stdout_trap)
+    profile.print_stats(stdout_trap, stripzeros='s' in opts)
     output = stdout_trap.getvalue()
     output = output.rstrip()