tools/profile_chrome/trace_packager.py

   1 # Copyright 2014 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 import codecs
   6 import gzip
   7 import json
   8 import os
   9 import shutil
  10 import sys
  11 import zipfile
  12
  13 from profile_chrome import util
  14
  15 from pylib import constants
  16
  17 sys.path.append(os.path.join(constants.DIR_SOURCE_ROOT,
  18                              'third_party',
  19                              'catapult'))
  20 # pylint: disable=F0401
  21 from tracing.build import trace2html
  22
  23
  24 def _PackageTracesAsHtml(trace_files, html_file):
  25   with codecs.open(html_file, mode='w', encoding='utf-8') as f:
  26     trace2html.WriteHTMLForTracesToFile(trace_files, f)
  27   for trace_file in trace_files:
  28     os.unlink(trace_file)
  29
  30
  31 def _CompressFile(host_file, output):
  32   with gzip.open(output, 'wb') as out, \
  33       open(host_file, 'rb') as input_file:
  34     out.write(input_file.read())
  35   os.unlink(host_file)
  36
  37
  38 def _ArchiveFiles(host_files, output):
  39   with zipfile.ZipFile(output, 'w', zipfile.ZIP_DEFLATED) as z:
  40     for host_file in host_files:
  41       z.write(host_file)
  42       os.unlink(host_file)
  43
  44
  45 def _MergeTracesIfNeeded(trace_files):
  46   if len(trace_files) <= 1:
  47     return trace_files
  48   merge_candidates = []
  49   for trace_file in trace_files:
  50     with open(trace_file) as f:
  51       # Try to detect a JSON file cheaply since that's all we can merge.
  52       if f.read(1) != '{':
  53         continue
  54       f.seek(0)
  55       try:
  56         json_data = json.load(f)
  57       except ValueError:
  58         continue
  59       merge_candidates.append((trace_file, json_data))
  60   if len(merge_candidates) <= 1:
  61     return trace_files
  62
  63   other_files = [f for f in trace_files
  64                  if not f in [c[0] for c in merge_candidates]]
  65   merged_file, merged_data = merge_candidates[0]
  66   for trace_file, json_data in merge_candidates[1:]:
  67     for key, value in json_data.items():
  68       if not merged_data.get(key) or json_data[key]:
  69         merged_data[key] = value
  70     os.unlink(trace_file)
  71
  72   with open(merged_file, 'w') as f:
  73     json.dump(merged_data, f)
  74   return [merged_file] + other_files
  75
  76
  77 def PackageTraces(trace_files, output=None, compress=False, write_json=False):
  78   trace_files = _MergeTracesIfNeeded(trace_files)
  79   if not write_json:
  80     html_file = os.path.splitext(trace_files[0])[0] + '.html'
  81     _PackageTracesAsHtml(trace_files, html_file)
  82     trace_files = [html_file]
  83
  84   if compress and len(trace_files) == 1:
  85     result = output or trace_files[0] + '.gz'
  86     _CompressFile(trace_files[0], result)
  87   elif len(trace_files) > 1:
  88     result = output or 'chrome-combined-trace-%s.zip' % util.GetTraceTimestamp()
  89     _ArchiveFiles(trace_files, result)
  90   elif output:
  91     result = output
  92     shutil.move(trace_files[0], result)
  93   else:
  94     result = trace_files[0]
  95   return result