Skip the decorator if an environment variable is set.

This should be a minor speedup and gives cleaner back traces. This can't be done with a flag since imports happen before flags.

PiperOrigin-RevId: 213631051

Author: tewalds

pysc2/lib/stopwatch.py

@@ -20,6 +20,7 @@
 from collections import defaultdict
 import functools
 import math
+import os
 import sys
 import threading
 import time
@@ -190,6 +191,9 @@ def random_func_name():
       If a name is passed, returns this as a decorator, otherwise returns the
       decorated function.
     """
+    if os.environ.get("SC2_NO_STOPWATCH"):
+      return name_or_func if callable(name_or_func) else lambda func: func
+
     def decorator(name, func):
       @functools.wraps(func)
       def _stopwatch(*args, **kwargs):
@@ -236,6 +240,7 @@ def merge(self, other):
 
   @staticmethod
   def parse(s):
+    """Parse the output below to create a new StopWatch."""
     stopwatch = StopWatch()
     for line in s.splitlines():
       if line.strip():

pysc2/lib/stopwatch_test.py

@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 from absl.testing import absltest
 from future.builtins import range  # pylint: disable=redefined-builtin
 
@@ -103,6 +105,18 @@ def testDivideZero(self):
     # Just make sure this doesn't have a divide by 0 for when the total is 0.
     self.assertIn("zero", str(sw))
 
+  @mock.patch.dict(os.environ, {"SC2_NO_STOPWATCH": "1"})
+  def testDecoratorDisabled(self):
+    sw = stopwatch.StopWatch()
+    self.assertEqual(round, sw.decorate(round))
+    self.assertEqual(round, sw.decorate("name")(round))
+
+  @mock.patch.dict(os.environ, {"SC2_NO_STOPWATCH": ""})
+  def testDecoratorEnabled(self):
+    sw = stopwatch.StopWatch()
+    self.assertNotEqual(round, sw.decorate(round))
+    self.assertNotEqual(round, sw.decorate("name")(round))
+
   def testSpeed(self):
     count = 1000