[android] Use contextlib_ext.Optional for conditional context managers.

build/android/pylib/local/device/local_device_perf_test_run.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import collections
-import contextlib
 import io
 import json
 import logging
 import os
 import pickle
 import shutil
 import tempfile
 import threading
 import time
 import zipfile
 
 from devil.android import battery_utils
 from devil.android import device_errors
 from devil.android import device_list
 from devil.android import device_utils
 from devil.android import forwarder
 from devil.android.tools import device_recovery
 from devil.android.tools import device_status
 from devil.utils import cmd_helper
 from devil.utils import parallelizer
 from devil.utils import reraiser_thread
 from pylib import constants
 from pylib.base import base_test_result
 from pylib.constants import host_paths
 from pylib.local.device import local_device_environment
 from pylib.local.device import local_device_test_run
 from py_trace_event import trace_event
+from py_utils import contextlib_ext
 
 
 class HeartBeat(object):
 
   def __init__(self, shard, wait_time=60*10):
     """ HeartBeat Logger constructor.
 
     Args:
       shard: A perf test runner device shard.
       wait_time: time to wait between heartbeat messages.
@@ skipping 45 matching lines @@
 
   def _RunSingleTest(self, test):
     self._test_instance.WriteBuildBotJson(self._output_dir)
 
     timeout = self._tests[test].get('timeout', self._timeout)
     cmd = self._CreateCmd(test)
     cwd = os.path.abspath(host_paths.DIR_SOURCE_ROOT)
 
     self._LogTest(test, cmd, timeout)
 
-    @contextlib.contextmanager
-    def trace_if_enabled(test):
-      try:
-        if self._test_instance.trace_output:
-          trace_event.trace_begin(test)
-        yield
-      finally:
-        if self._test_instance.trace_output:
-          trace_event.trace_end(test)
-
     try:
       start_time = time.time()
 
-      with trace_if_enabled(test):
+      with contextlib_ext.Optional(
+          trace_event.trace(test),
+          self._test_instance.trace_output):
         exit_code, output = cmd_helper.GetCmdStatusAndOutputWithTimeout(
             cmd, timeout, cwd=cwd, shell=True)
       end_time = time.time()
       json_output = self._test_instance.ReadChartjsonOutput(self._output_dir)
       if exit_code == 0:
         result_type = base_test_result.ResultType.PASS
       else:
         result_type = base_test_result.ResultType.FAIL
     except cmd_helper.TimeoutError as e:
       end_time = time.time()
@@ skipping 411 matching lines @@
   # override
   def _RunTest(self, _device, _test):
     raise NotImplementedError
 
 
 class TestDictVersionError(Exception):
   pass
 
 class PerfTestRunGetStepsError(Exception):
   pass