Add full profile sort, dumping to profile results

Change-Id: Ib256ae34de15d29ee9a48e3be86073610f8d1a65

1 files changed, 10 insertions, 2 deletions

diff --git a/lib/sqlalchemy/testing/profiling.py b/lib/sqlalchemy/testing/profiling.py
index 92bd452a5..16215dcd5 100644
--- a/lib/sqlalchemy/testing/profiling.py
+++ b/lib/sqlalchemy/testing/profiling.py
@@ -17,6 +17,7 @@ import contextlib
 import os
 import platform
 import pstats
+import re
 import sys
 
 from . import config
@@ -64,7 +65,7 @@ class ProfileStatsFile(object):
 
     """
 
-    def __init__(self, filename, sort="cumulative"):
+    def __init__(self, filename, sort="cumulative", dump=None):
         self.force_write = (
             config.options is not None and config.options.force_write_profiles
         )
@@ -76,6 +77,7 @@ class ProfileStatsFile(object):
         self.data = collections.defaultdict(
             lambda: collections.defaultdict(dict)
         )
+        self.dump = dump
         self.sort = sort
         self._read()
         if self.write:
@@ -292,8 +294,14 @@ def count_functions(variance=0.05):
         line_no, expected_count = expected
 
     print(("Pstats calls: %d Expected %s" % (callcount, expected_count)))
-    stats.sort_stats(_profile_stats.sort)
+    stats.sort_stats(*re.split(r"[, ]", _profile_stats.sort))
     stats.print_stats()
+    if _profile_stats.dump:
+        base, ext = os.path.splitext(_profile_stats.dump)
+        test_name = _current_test.split(".")[-1]
+        dumpfile = "%s_%s%s" % (base, test_name, ext or ".profile")
+        stats.dump_stats(dumpfile)
+        print("Dumped stats to file %s" % dumpfile)
     # stats.print_callers()
     if _profile_stats.force_write:
         _profile_stats.replace(callcount)