Allow pstats.Stats creator to specify an alternate to stdout.

3 files changed, 91 insertions, 74 deletions

diff --git a/Doc/lib/libprofile.tex b/Doc/lib/libprofile.tex
index 97c7191..9ff5ba0 100644
--- a/Doc/lib/libprofile.tex
+++ b/Doc/lib/libprofile.tex
@@ -391,17 +391,17 @@ Analysis of the profiler data is done using this class from the
 % (This \stmodindex use may be hard to change ;-( )
 \stmodindex{pstats}
 
-\begin{classdesc}{Stats}{filename\optional{, \moreargs}}
+\begin{classdesc}{Stats}{filename\optional{, \moreargs\optional{, stream=sys.stdout}}}
 This class constructor creates an instance of a ``statistics object''
 from a \var{filename} (or set of filenames).  \class{Stats} objects are
-manipulated by methods, in order to print useful reports.
-
-The file selected by the above constructor must have been created by
-the corresponding version of \module{profile} or \module{cProfile}.
-To be specific, there is
-\emph{no} file compatibility guaranteed with future versions of this
-profiler, and there is no compatibility with files produced by other
-profilers.
+manipulated by methods, in order to print useful reports.  You may specify
+an alternate output stream by giving the keyword argument, \code{stream}.
+
+The file selected by the above constructor must have been created by the
+corresponding version of \module{profile} or \module{cProfile}.  To be
+specific, there is \emph{no} file compatibility guaranteed with future
+versions of this profiler, and there is no compatibility with files produced
+by other profilers.
 %(such as the old system profiler).
 
 If several files are provided, all the statistics for identical
diff --git a/Lib/pstats.py b/Lib/pstats.py
index 930cc6d..c3a8828 100644
--- a/Lib/pstats.py
+++ b/Lib/pstats.py
@@ -32,6 +32,7 @@
 # CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 
+import sys
 import os
 import time
 import marshal
@@ -58,18 +59,31 @@ class Stats:
     printed.
 
     The sort_stats() method now processes some additional options (i.e., in
-    addition to the old -1, 0, 1, or 2).  It takes an arbitrary number of quoted
-    strings to select the sort order.  For example sort_stats('time', 'name')
-    sorts on the major key of "internal function time", and on the minor
-    key of 'the name of the function'.  Look at the two tables in sort_stats()
-    and get_sort_arg_defs(self) for more examples.
+    addition to the old -1, 0, 1, or 2).  It takes an arbitrary number of
+    quoted strings to select the sort order.  For example sort_stats('time',
+    'name') sorts on the major key of 'internal function time', and on the
+    minor key of 'the name of the function'.  Look at the two tables in
+    sort_stats() and get_sort_arg_defs(self) for more examples.
 
-    All methods now return "self",  so you can string together commands like:
+    All methods return self,  so you can string together commands like:
         Stats('foo', 'goo').strip_dirs().sort_stats('calls').\
                             print_stats(5).print_callers(5)
     """
 
-    def __init__(self, *args):
+    def __init__(self, *args, **kwds):
+        # I can't figure out how to explictly specify a stream keyword arg
+        # with *args:
+        #   def __init__(self, *args, stream=sys.stdout): ...
+        # so I use **kwds and sqauwk if something unexpected is passed in.
+        self.stream = sys.stdout
+        if "stream" in kwds:
+            self.stream = kwds["stream"]
+            del kwds["stream"]
+        if kwds:
+            keys = kwds.keys()
+            keys.sort()
+            extras = ", ".join(["%s=%s" % (k, kwds[k]) for k in keys])
+            raise ValueError, "unrecognized keyword args: %s" % extras
         if not len(args):
             arg = None
         else:
@@ -96,9 +110,9 @@ class Stats:
             trouble = 0
         finally:
             if trouble:
-                print "Invalid timing data",
-                if self.files: print self.files[-1],
-                print
+                print >> self.stream, "Invalid timing data",
+                if self.files: print >> self.stream, self.files[-1],
+                print >> self.stream
 
     def load_stats(self, arg):
         if not arg:  self.stats = {}
@@ -320,7 +334,7 @@ class Stats:
 
         if not list:
             return 0, list
-        print msg
+        print >> self.stream, msg
         if count < len(self.stats):
             width = 0
             for func in list:
@@ -330,24 +344,24 @@ class Stats:
 
     def print_stats(self, *amount):
         for filename in self.files:
-            print filename
-        if self.files: print
+            print >> self.stream, filename
+        if self.files: print >> self.stream
         indent = ' ' * 8
         for func in self.top_level:
-            print indent, func_get_function_name(func)
+            print >> self.stream, indent, func_get_function_name(func)
 
-        print indent, self.total_calls, "function calls",
+        print >> self.stream, indent, self.total_calls, "function calls",
         if self.total_calls != self.prim_calls:
-            print "(%d primitive calls)" % self.prim_calls,
-        print "in %.3f CPU seconds" % self.total_tt
-        print
+            print >> self.stream, "(%d primitive calls)" % self.prim_calls,
+        print >> self.stream, "in %.3f CPU seconds" % self.total_tt
+        print >> self.stream
         width, list = self.get_print_list(amount)
         if list:
             self.print_title()
             for func in list:
                 self.print_line(func)
-            print
-            print
+            print >> self.stream
+            print >> self.stream
         return self
 
     def print_callees(self, *amount):
@@ -361,8 +375,8 @@ class Stats:
                     self.print_call_line(width, func, self.all_callees[func])
                 else:
                     self.print_call_line(width, func, {})
-            print
-            print
+            print >> self.stream
+            print >> self.stream
         return self
 
     def print_callers(self, *amount):
@@ -372,12 +386,12 @@ class Stats:
             for func in list:
                 cc, nc, tt, ct, callers = self.stats[func]
                 self.print_call_line(width, func, callers, "<-")
-            print
-            print
+            print >> self.stream
+            print >> self.stream
         return self
 
     def print_call_heading(self, name_size, column_title):
-        print "Function ".ljust(name_size) + column_title
+        print >> self.stream, "Function ".ljust(name_size) + column_title
         # print sub-header only if we have new-style callers
         subheader = False
         for cc, nc, tt, ct, callers in self.stats.itervalues():
@@ -386,12 +400,12 @@ class Stats:
                 subheader = isinstance(value, tuple)
                 break
         if subheader:
-            print " "*name_size + "    ncalls  tottime  cumtime"
+            print >> self.stream, " "*name_size + "    ncalls  tottime  cumtime"
 
     def print_call_line(self, name_size, source, call_dict, arrow="->"):
-        print func_std_string(source).ljust(name_size) + arrow,
+        print >> self.stream, func_std_string(source).ljust(name_size) + arrow,
         if not call_dict:
-            print
+            print >> self.stream
             return
         clist = call_dict.keys()
         clist.sort()
@@ -411,30 +425,30 @@ class Stats:
             else:
                 substats = '%s(%r) %s' % (name, value, f8(self.stats[func][3]))
                 left_width = name_size + 3
-            print indent*left_width + substats
+            print >> self.stream, indent*left_width + substats
             indent = " "
 
     def print_title(self):
-        print '   ncalls  tottime  percall  cumtime  percall', \
-              'filename:lineno(function)'
+        print >> self.stream, '   ncalls  tottime  percall  cumtime  percall',
+        print >> self.stream, 'filename:lineno(function)'
 
     def print_line(self, func):  # hack : should print percentages
         cc, nc, tt, ct, callers = self.stats[func]
         c = str(nc)
         if nc != cc:
             c = c + '/' + str(cc)
-        print c.rjust(9),
-        print f8(tt),
+        print >> self.stream, c.rjust(9),
+        print >> self.stream, f8(tt),
         if nc == 0:
-            print ' '*8,
+            print >> self.stream, ' '*8,
         else:
-            print f8(tt/nc),
-        print f8(ct),
+            print >> self.stream, f8(tt/nc),
+        print >> self.stream, f8(ct),
         if cc == 0:
-            print ' '*8,
+            print >> self.stream, ' '*8,
         else:
-            print f8(ct/cc),
-        print func_std_string(func)
+            print >> self.stream, f8(ct/cc),
+        print >> self.stream, func_std_string(func)
 
 class TupleComp:
     """This class provides a generic function for comparing any two tuples.
@@ -549,7 +563,7 @@ if __name__ == '__main__':
                 try:
                     frac = float(term)
                     if frac > 1 or frac < 0:
-                        print "Fraction argument mus be in [0, 1]"
+                        print >> self.stream, "Fraction argument must be in [0, 1]"
                         continue
                     processed.append(frac)
                     continue
@@ -559,93 +573,93 @@ if __name__ == '__main__':
             if self.stats:
                 getattr(self.stats, fn)(*processed)
             else:
-                print "No statistics object is loaded."
+                print >> self.stream, "No statistics object is loaded."
             return 0
         def generic_help(self):
-            print "Arguments may be:"
-            print "* An integer maximum number of entries to print."
-            print "* A decimal fractional number between 0 and 1, controlling"
-            print "  what fraction of selected entries to print."
-            print "* A regular expression; only entries with function names"
-            print "  that match it are printed."
+            print >> self.stream, "Arguments may be:"
+            print >> self.stream, "* An integer maximum number of entries to print."
+            print >> self.stream, "* A decimal fractional number between 0 and 1, controlling"
+            print >> self.stream, "  what fraction of selected entries to print."
+            print >> self.stream, "* A regular expression; only entries with function names"
+            print >> self.stream, "  that match it are printed."
 
         def do_add(self, line):
             self.stats.add(line)
             return 0
         def help_add(self):
-            print "Add profile info from given file to current statistics object."
+            print >> self.stream, "Add profile info from given file to current statistics object."
 
         def do_callees(self, line):
             return self.generic('print_callees', line)
         def help_callees(self):
-            print "Print callees statistics from the current stat object."
+            print >> self.stream, "Print callees statistics from the current stat object."
             self.generic_help()
 
         def do_callers(self, line):
             return self.generic('print_callers', line)
         def help_callers(self):
-            print "Print callers statistics from the current stat object."
+            print >> self.stream, "Print callers statistics from the current stat object."
             self.generic_help()
 
         def do_EOF(self, line):
-            print ""
+            print >> self.stream, ""
             return 1
         def help_EOF(self):
-            print "Leave the profile brower."
+            print >> self.stream, "Leave the profile brower."
 
         def do_quit(self, line):
             return 1
         def help_quit(self):
-            print "Leave the profile brower."
+            print >> self.stream, "Leave the profile brower."
 
         def do_read(self, line):
             if line:
                 try:
                     self.stats = Stats(line)
                 except IOError, args:
-                    print args[1]
+                    print >> self.stream, args[1]
                     return
                 self.prompt = line + "% "
             elif len(self.prompt) > 2:
                 line = self.prompt[-2:]
             else:
-                print "No statistics object is current -- cannot reload."
+                print >> self.stream, "No statistics object is current -- cannot reload."
             return 0
         def help_read(self):
-            print "Read in profile data from a specified file."
+            print >> self.stream, "Read in profile data from a specified file."
 
         def do_reverse(self, line):
             self.stats.reverse_order()
             return 0
         def help_reverse(self):
-            print "Reverse the sort order of the profiling report."
+            print >> self.stream, "Reverse the sort order of the profiling report."
 
         def do_sort(self, line):
             abbrevs = self.stats.get_sort_arg_defs()
             if line and not filter(lambda x,a=abbrevs: x not in a,line.split()):
                 self.stats.sort_stats(*line.split())
             else:
-                print "Valid sort keys (unique prefixes are accepted):"
+                print >> self.stream, "Valid sort keys (unique prefixes are accepted):"
                 for (key, value) in Stats.sort_arg_dict_default.iteritems():
-                    print "%s -- %s" % (key, value[1])
+                    print >> self.stream, "%s -- %s" % (key, value[1])
             return 0
         def help_sort(self):
-            print "Sort profile data according to specified keys."
-            print "(Typing `sort' without arguments lists valid keys.)"
+            print >> self.stream, "Sort profile data according to specified keys."
+            print >> self.stream, "(Typing `sort' without arguments lists valid keys.)"
         def complete_sort(self, text, *args):
             return [a for a in Stats.sort_arg_dict_default if a.startswith(text)]
 
         def do_stats(self, line):
             return self.generic('print_stats', line)
         def help_stats(self):
-            print "Print statistics from the current stat object."
+            print >> self.stream, "Print statistics from the current stat object."
             self.generic_help()
 
         def do_strip(self, line):
             self.stats.strip_dirs()
             return 0
         def help_strip(self):
-            print "Strip leading path information from filenames in the report."
+            print >> self.stream, "Strip leading path information from filenames in the report."
 
         def postcmd(self, stop, line):
             if stop:
@@ -653,14 +667,14 @@ if __name__ == '__main__':
             return None
 
     import sys
-    print "Welcome to the profile statistics browser."
+    print >> self.stream, "Welcome to the profile statistics browser."
     if len(sys.argv) > 1:
         initprofile = sys.argv[1]
     else:
         initprofile = None
     try:
         ProfileBrowser(initprofile).cmdloop()
-        print "Goodbye."
+        print >> self.stream, "Goodbye."
     except KeyboardInterrupt:
         pass
 
diff --git a/Misc/NEWS b/Misc/NEWS
index 0cda340..03dc27c 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -120,6 +120,9 @@ Library
 
 - Fix exception when doing glob.glob('anything*/')
 
+- The pstats.Stats class accepts an optional stream keyword argument to
+  direct output to an alternate file-like object.
+
 Build
 -----