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