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                              'tracing'))
  21 # pylint: disable=F0401
  22 from tracing_build import trace2html
  23
  24
  25 def _PackageTracesAsHtml(trace_files, html_file):
  26   with codecs.open(html_file, mode='w', encoding='utf-8') as f:
  27     trace2html.WriteHTMLForTracesToFile(trace_files, f)
  28   for trace_file in trace_files:
  29     os.unlink(trace_file)
  30
  31
  32 def _CompressFile(host_file, output):
  33   with gzip.open(output, 'wb') as out, \
  34       open(host_file, 'rb') as input_file:
  35     out.write(input_file.read())
  36   os.unlink(host_file)
  37
  38
  39 def _ArchiveFiles(host_files, output):
  40   with zipfile.ZipFile(output, 'w', zipfile.ZIP_DEFLATED) as z:
  41     for host_file in host_files:
  42       z.write(host_file)
  43       os.unlink(host_file)
  44
  45
  46 def _MergeTracesIfNeeded(trace_files):
  47   if len(trace_files) <= 1:
  48     return trace_files
  49   merge_candidates = []
  50   for trace_file in trace_files:
  51     with open(trace_file) as f:
  52       # Try to detect a JSON file cheaply since that's all we can merge.
  53       if f.read(1) != '{':
  54         continue
  55       f.seek(0)
  56       try:
  57         json_data = json.load(f)
  58       except ValueError:
  59         continue
  60       merge_candidates.append((trace_file, json_data))
  61   if len(merge_candidates) <= 1:
  62     return trace_files
  63
  64   other_files = [f for f in trace_files
  65                  if not f in [c[0] for c in merge_candidates]]
  66   merged_file, merged_data = merge_candidates[0]
  67   for trace_file, json_data in merge_candidates[1:]:
  68     for key, value in json_data.items():
  69       if not merged_data.get(key) or json_data[key]:
  70         merged_data[key] = value
  71     os.unlink(trace_file)
  72
  73   with open(merged_file, 'w') as f:
  74     json.dump(merged_data, f)
  75   return [merged_file] + other_files
  76
  77
  78 def PackageTraces(trace_files, output=None, compress=False, write_json=False):
  79   trace_files = _MergeTracesIfNeeded(trace_files)
  80   if not write_json:
  81     html_file = os.path.splitext(trace_files[0])[0] + '.html'
  82     _PackageTracesAsHtml(trace_files, html_file)
  83     trace_files = [html_file]
  84
  85   if compress and len(trace_files) == 1:
  86     result = output or trace_files[0] + '.gz'
  87     _CompressFile(trace_files[0], result)
  88   elif len(trace_files) > 1:
  89     result = output or 'chrome-combined-trace-%s.zip' % util.GetTraceTimestamp()
  90     _ArchiveFiles(trace_files, result)
  91   elif output:
  92     result = output
  93     shutil.move(trace_files[0], result)
  94   else:
  95     result = trace_files[0]
  96   return result