Make the pstats browser recognise SampledStats

The SampledStats results are stored in the exact same format as Stats,
but since the results don't represent call counts but sample counts, the
column headers are different to account for this.

This ensure that using the pstats browser instantiates the right object
to handle the correct columns, add a factory function which can
instantiate the correct class. As the Stats class can only handle either
a filename or an object which provides the 'stats' attribute in a
pre-parsed format, this provides a StatsLoaderShim to avoid marshalling
the data twice (once to check the marker and once in the Stats module if
we were to pass the file name).

Author: lkollar

Lib/profile/sample.py

@@ -62,8 +62,10 @@ def print_stats(self, sort=-1):
         pstats.SampledStats(self).strip_dirs().sort_stats(*sort).print_stats()
 
     def dump_stats(self, file):
+        stats_with_marker = dict(self.stats)
+        stats_with_marker[("__sampled__",)] = True
         with open(file, "wb") as f:
-            marshal.dump(self.stats, f)
+            marshal.dump(stats_with_marker, f)
 
     # Needed for compatibility with pstats.Stats
     def create_stats(self):

Lib/pstats.py

@@ -139,7 +139,11 @@ def load_stats(self, arg):
             return
         elif isinstance(arg, str):
             with open(arg, 'rb') as f:
-                self.stats = marshal.load(f)
+                stats = marshal.load(f)
+            if (('__sampled__',)) in stats:
+                stats.pop((('__sampled__',)))
+                self.__class__ = SampledStats
+            self.stats = stats
             try:
                 file_stats = os.stat(arg)
                 arg = time.ctime(file_stats.st_mtime) + "    " + arg
@@ -639,6 +643,24 @@ def f8(x):
 # Statistics browser added by ESR, April 2001
 #**************************************************************************
 
+class StatsLoaderShim:
+    """Compatibility shim implementing 'create_stats' needed by Stats classes
+    to handle already unmarshalled data."""
+    def __init__(self, raw_stats):
+        self.stats = raw_stats
+
+    def create_stats(self):
+        pass
+
+def stats_factory(raw_stats):
+    """Return a Stats or SampledStats instance based on the marker in raw_stats."""
+    if (('__sampled__',)) in raw_stats:
+        raw_stats = dict(raw_stats)  # avoid mutating caller's dict
+        raw_stats.pop((('__sampled__',)))
+        return SampledStats(StatsLoaderShim(raw_stats))
+    else:
+        return Stats(StatsLoaderShim(raw_stats))
+
 if __name__ == '__main__':
     import cmd
     try:
@@ -725,7 +747,15 @@ def help_quit(self):
         def do_read(self, line):
             if line:
                 try:
-                    self.stats = Stats(line)
+                    with open(line, 'rb') as f:
+                        raw_stats = marshal.load(f)
+                    self.stats = stats_factory(raw_stats)
+                    try:
+                        file_stats = os.stat(line)
+                        arg = time.ctime(file_stats.st_mtime) + "    " + line
+                    except Exception:
+                        arg = line
+                    self.stats.files = [arg]
                 except OSError as err:
                     print(err.args[1], file=self.stream)
                     return