use argparse instead of getopt in timeit

Author: picnixz

Lib/test/test_timeit.py

@@ -252,11 +252,9 @@ def run_main(self, seconds_per_increment=1.0, switches=None, timer=None):
         return s.getvalue()
 
     def test_main_bad_switch(self):
-        s = self.run_main(switches=['--bad-switch'])
-        self.assertEqual(s, dedent("""\
-            option --bad-switch not recognized
-            use -h/--help for command line help
-            """))
+        with self.assertRaises(SystemExit), captured_stderr() as s:
+            self.run_main(switches=['--bad-switch'])
+        self.assertIn("unrecognized arguments: --bad-switch", s.getvalue())
 
     def test_main_seconds(self):
         s = self.run_main(seconds_per_increment=5.5)
@@ -294,11 +292,6 @@ def test_main_negative_reps(self):
         s = self.run_main(seconds_per_increment=60.0, switches=['-r-5'])
         self.assertEqual(s, "1 loop, best of 1: 60 sec per loop\n")
 
-    @unittest.skipIf(sys.flags.optimize >= 2, "need __doc__")
-    def test_main_help(self):
-        s = self.run_main(switches=['-h'])
-        self.assertEqual(s, timeit.__doc__)
-
     def test_main_verbose(self):
         s = self.run_main(switches=['-v'])
         self.assertEqual(s, dedent("""\
@@ -345,11 +338,10 @@ def test_main_with_time_unit(self):
         self.assertEqual(unit_usec,
                 "100 loops, best of 5: 3e+03 usec per loop\n")
         # Test invalid unit input
-        with captured_stderr() as error_stringio:
-            invalid = self.run_main(seconds_per_increment=0.003,
-                    switches=['-u', 'parsec'])
-        self.assertEqual(error_stringio.getvalue(),
-                    "Unrecognized unit. Please select nsec, usec, msec, or sec.\n")
+        with self.assertRaises(SystemExit), captured_stderr() as s:
+            self.run_main(seconds_per_increment=0.003,
+                          switches=['-u', 'parsec'])
+        self.assertIn("invalid choice: 'parsec'", s.getvalue())
 
     def test_main_exception(self):
         with captured_stderr() as error_stringio:

Lib/timeit.py

@@ -6,34 +6,7 @@
 
 Library usage: see the Timer class.
 
-Command line usage:
-    python timeit.py [-n N] [-r N] [-s S] [-p] [-h] [--] [statement]
-
-Options:
-  -n/--number N: how many times to execute 'statement' (default: see below)
-  -r/--repeat N: how many times to repeat the timer (default 5)
-  -s/--setup S: statement to be executed once initially (default 'pass').
-                Execution time of this setup statement is NOT timed.
-  -p/--process: use time.process_time() (default is time.perf_counter())
-  -v/--verbose: print raw timing results; repeat for more digits precision
-  -u/--unit: set the output time unit (nsec, usec, msec, or sec)
-  -h/--help: print this usage message and exit
-  --: separate options from statement, use when statement starts with -
-  statement: statement to be timed (default 'pass')
-
-A multi-line statement may be given by specifying each line as a
-separate argument; indented lines are possible by enclosing an
-argument in quotes and using leading spaces.  Multiple -s options are
-treated similarly.
-
-If -n is not given, a suitable number of loops is calculated by trying
-increasing numbers from the sequence 1, 2, 5, 10, 20, 50, ... until the
-total time is at least 0.2 seconds.
-
-Note: there is a certain baseline overhead associated with executing a
-pass statement.  It differs between versions.  The code here doesn't try
-to hide it, but you should be aware of it.  The baseline overhead can be
-measured by invoking the program without arguments.
+Command line usage: python -m timeit --help
 
 Classes:
 
@@ -240,6 +213,71 @@ def repeat(stmt="pass", setup="pass", timer=default_timer,
     return Timer(stmt, setup, timer, globals).repeat(repeat, number)
 
 
+def _make_parser():
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="""\
+Tool avoiding a number of common traps for measuring execution times.
+See also Tim Peters' introduction to the Algorithms chapter in the
+Python Cookbook, published by O'Reilly.""",
+        epilog="""\
+A multi-line statement may be given by specifying each line as a
+separate argument; indented lines are possible by enclosing an
+argument in quotes and using leading spaces.  Multiple -s options are
+treated similarly.
+
+If -n is not given, a suitable number of loops is calculated by trying
+increasing numbers from the sequence 1, 2, 5, 10, 20, 50, ... until the
+total time is at least 0.2 seconds.
+
+Note: there is a certain baseline overhead associated with executing a
+pass statement.  It differs between versions.  The code here doesn't try
+to hide it, but you should be aware of it.  The baseline overhead can be
+measured by invoking the program without arguments.""",
+        formatter_class=argparse.RawTextHelpFormatter,
+        allow_abbrev=False,
+    )
+    # use a group to avoid rendering a 'positional arguments' section
+    group = parser.add_argument_group()
+    group.add_argument(
+        "-n", "--number", metavar="N", default=0, type=int,
+        help="how many times to execute 'statement' (default: see below)",
+    )
+    group.add_argument(
+        "-r", "--repeat", type=int, metavar="N", default=default_repeat,
+        help="how many times to repeat the timer (default: 5)",
+    )
+    group.add_argument(
+        "-s", "--setup", action="append", metavar="S", default=[],
+        help="statement to be executed once initially (default: 'pass').\n"
+             "Execution time of this setup statement is NOT timed.",
+    )
+    group.add_argument(
+        "-u", "--unit", choices=["nsec", "usec", "msec", "sec"], metavar="U",
+        help="set the output time unit (nsec, usec, msec, or sec)",
+    )
+    group.add_argument(
+        "-p", "--process", action="store_true",
+        help="use time.process_time() (default: time.perf_counter())",
+    )
+    group.add_argument(
+        "-v", "--verbose", action="count", default=0,
+        help="print raw timing results; repeat for more digits precision",
+    )
+    # add a dummy option to document '--'
+    group.add_argument(
+        "--", dest="_", action="store_const", const=None,
+        help="separate options from statement; "
+             "use when statement starts with -",
+    )
+    group.add_argument(
+        "statement", nargs="*",
+        help="statement to be timed (default: 'pass')",
+    )
+    return parser
+
+
 def main(args=None, *, _wrap_timer=None):
     """Main program, used when run as a script.
 
@@ -257,53 +295,16 @@ def main(args=None, *, _wrap_timer=None):
     is not None, it must be a callable that accepts a timer function
     and returns another timer function (used for unit testing).
     """
-    if args is None:
-        args = sys.argv[1:]
-    import getopt
-    try:
-        opts, args = getopt.getopt(args, "n:u:s:r:pvh",
-                                   ["number=", "setup=", "repeat=",
-                                    "process", "verbose", "unit=", "help"])
-    except getopt.error as err:
-        print(err)
-        print("use -h/--help for command line help")
-        return 2
-
-    timer = default_timer
-    stmt = "\n".join(args) or "pass"
-    number = 0  # auto-determine
-    setup = []
-    repeat = default_repeat
-    verbose = 0
-    time_unit = None
-    units = {"nsec": 1e-9, "usec": 1e-6, "msec": 1e-3, "sec": 1.0}
-    precision = 3
-    for o, a in opts:
-        if o in ("-n", "--number"):
-            number = int(a)
-        if o in ("-s", "--setup"):
-            setup.append(a)
-        if o in ("-u", "--unit"):
-            if a in units:
-                time_unit = a
-            else:
-                print("Unrecognized unit. Please select nsec, usec, msec, or sec.",
-                      file=sys.stderr)
-                return 2
-        if o in ("-r", "--repeat"):
-            repeat = int(a)
-            if repeat <= 0:
-                repeat = 1
-        if o in ("-p", "--process"):
-            timer = time.process_time
-        if o in ("-v", "--verbose"):
-            if verbose:
-                precision += 1
-            verbose += 1
-        if o in ("-h", "--help"):
-            print(__doc__, end="")
-            return 0
-    setup = "\n".join(setup) or "pass"
+    parser = _make_parser()
+    args = parser.parse_args(args)
+
+    setup = "\n".join(args.setup) or "pass"
+    stmt = "\n".join(args.statement) or "pass"
+    timer = time.process_time if args.process else default_timer
+    number = args.number  # will be deduced if 0
+    repeat = max(1, args.repeat)
+    verbose = bool(args.verbose)
+    precision = 3 + max(0, args.verbose - 1)
 
     # Include the current directory, so that local imports work (sys.path
     # contains the directory of this script, rather than the current
@@ -338,14 +339,16 @@ def callback(number, time_taken):
         t.print_exc()
         return 1
 
+    units = {"nsec": 1e-9, "usec": 1e-6, "msec": 1e-3, "sec": 1.0}
+    scales = [(scale, unit_name) for unit_name, scale in units.items()]
+    scales.sort(reverse=True)
+
     def format_time(dt):
-        unit = time_unit
+        unit = args.unit
 
         if unit is not None:
             scale = units[unit]
         else:
-            scales = [(scale, unit) for unit, scale in units.items()]
-            scales.sort(reverse=True)
             for scale, unit in scales:
                 if dt >= scale:
                     break
@@ -362,7 +365,6 @@ def format_time(dt):
           % (number, 's' if number != 1 else '',
              repeat, format_time(best)))
 
-    best = min(timings)
     worst = max(timings)
     if worst >= best * 4:
         import warnings
@@ -371,7 +373,7 @@ def format_time(dt):
                                "slower than the best time (%s)."
                                % (format_time(worst), format_time(best)),
                                UserWarning, '', 0)
-    return None
+    return 0
 
 
 if __name__ == "__main__":