bpo-29237: Create enum for pstats sorting options (GH-5103)

1 files changed, 38 insertions, 8 deletions

diff --git a/Lib/pstats.py b/Lib/pstats.py
index b7a2054..1b57d26 100644
--- a/Lib/pstats.py
+++ b/Lib/pstats.py
@@ -25,9 +25,32 @@ import os
 import time
 import marshal
 import re
+from enum import Enum
 from functools import cmp_to_key
 
-__all__ = ["Stats"]
+__all__ = ["Stats", "SortKey"]
+
+
+class SortKey(str, Enum):
+    CALLS = 'calls', 'ncalls'
+    CUMULATIVE = 'cumulative', 'cumtime'
+    FILENAME = 'filename', 'module'
+    LINE = 'line'
+    NAME = 'name'
+    NFL = 'nfl'
+    PCALLS = 'pcalls'
+    STDNAME = 'stdname'
+    TIME = 'time', 'tottime'
+
+    def __new__(cls, *values):
+        obj = str.__new__(cls)
+
+        obj._value_ = values[0]
+        for other_value in values[1:]:
+            cls._value2member_map_[other_value] = obj
+        obj._all_values = values
+        return obj
+
 
 class Stats:
     """This class is used for creating reports from data generated by the
@@ -49,13 +72,14 @@ class Stats:
 
     The sort_stats() method now processes some additional options (i.e., in
     addition to the old -1, 0, 1, or 2 that are respectively interpreted as
-    'stdname', 'calls', 'time', and 'cumulative').  It takes an arbitrary number
-    of quoted strings to select the sort order.
+    'stdname', 'calls', 'time', and 'cumulative').  It takes either an
+    arbitrary number of quoted strings or SortKey enum 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.
+    For example sort_stats('time', 'name') or sort_stats(SortKey.TIME,
+    SortKey.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 return self, so you can string together commands like:
         Stats('foo', 'goo').strip_dirs().sort_stats('calls').\
@@ -161,7 +185,6 @@ class Stats:
               "ncalls"    : (((1,-1),              ), "call count"),
               "cumtime"   : (((3,-1),              ), "cumulative time"),
               "cumulative": (((3,-1),              ), "cumulative time"),
-              "file"      : (((4, 1),              ), "file name"),
               "filename"  : (((4, 1),              ), "file name"),
               "line"      : (((5, 1),              ), "line number"),
               "module"    : (((4, 1),              ), "file name"),
@@ -202,12 +225,19 @@ class Stats:
                        0:  "calls",
                        1:  "time",
                        2:  "cumulative"}[field[0]] ]
+        elif len(field) >= 2:
+            for arg in field[1:]:
+                if type(arg) != type(field[0]):
+                    raise TypeError("Can't have mixed argument type")
 
         sort_arg_defs = self.get_sort_arg_defs()
+
         sort_tuple = ()
         self.sort_type = ""
         connector = ""
         for word in field:
+            if isinstance(word, SortKey):
+                word = word.value
             sort_tuple = sort_tuple + sort_arg_defs[word][0]
             self.sort_type += connector + sort_arg_defs[word][1]
             connector = ", "