Add full profile sort, dumping to profile results

Change-Id: Ib256ae34de15d29ee9a48e3be86073610f8d1a65

2 files changed, 17 insertions, 2 deletions

diff --git a/lib/sqlalchemy/testing/plugin/plugin_base.py b/lib/sqlalchemy/testing/plugin/plugin_base.py
index f7d0dd3ea..bf168efce 100644
--- a/lib/sqlalchemy/testing/plugin/plugin_base.py
+++ b/lib/sqlalchemy/testing/plugin/plugin_base.py
@@ -117,6 +117,12 @@ def setup_options(make_option):
         help="Type of sort for profiling standard output",
     )
     make_option(
+        "--profile-dump",
+        type="string",
+        dest="profiledump",
+        help="Filename where a single profile run will be dumped",
+    )
+    make_option(
         "--postgresql-templatedb",
         type="string",
         help="name of template database to use for PostgreSQL "
@@ -495,6 +501,7 @@ def _setup_profiling(options, file_config):
     profiling._profile_stats = profiling.ProfileStatsFile(
         file_config.get("sqla_testing", "profile_file"),
         sort=options.profilesort,
+        dump=options.profiledump,
     )
 
 
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)