• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import csv
2import django.http
3import common
4from autotest_lib.frontend.afe import rpc_utils
5
6class CsvEncoder(object):
7    def __init__(self, request, response):
8        self._request = request
9        self._response = response
10        self._output_rows = []
11
12
13    def _append_output_row(self, row):
14        self._output_rows.append(row)
15
16
17    def _build_response(self):
18        response = django.http.HttpResponse(mimetype='text/csv')
19        response['Content-Disposition'] = (
20            'attachment; filename=tko_query.csv')
21        writer = csv.writer(response)
22        writer.writerows(self._output_rows)
23        return response
24
25
26    def encode(self):
27        raise NotImplementedError
28
29
30class UnhandledMethodEncoder(CsvEncoder):
31    def encode(self):
32        return rpc_utils.raw_http_response(
33            'Unhandled method %s (this indicates a bug)\r\n' %
34            self._request['method'])
35
36
37class SpreadsheetCsvEncoder(CsvEncoder):
38    def _total_index(self, group, num_columns):
39        row_index, column_index = group['header_indices']
40        return row_index * num_columns + column_index
41
42
43    def _group_string(self, group):
44        result = '%s / %s' % (group['pass_count'], group['complete_count'])
45        if group['incomplete_count'] > 0:
46            result +=  ' (%s incomplete)' % group['incomplete_count']
47        if 'extra_info' in group:
48            result = '\n'.join([result] + group['extra_info'])
49        return result
50
51
52    def _build_value_table(self):
53        value_table = [''] * self._num_rows * self._num_columns
54        for group in self._response['groups']:
55            total_index = self._total_index(group, self._num_columns)
56            value_table[total_index] = self._group_string(group)
57        return value_table
58
59
60    def _header_string(self, header_value):
61        return '/'.join(header_value)
62
63
64    def _process_value_table(self, value_table, row_headers):
65        total_index = 0
66        for row_index in xrange(self._num_rows):
67            row_header = self._header_string(row_headers[row_index])
68            row_end_index = total_index + self._num_columns
69            row_values = value_table[total_index:row_end_index]
70            self._append_output_row([row_header] + row_values)
71            total_index += self._num_columns
72
73
74    def encode(self):
75        header_values = self._response['header_values']
76        assert len(header_values) == 2
77        row_headers, column_headers = header_values
78        self._num_rows, self._num_columns = (len(row_headers),
79                                             len(column_headers))
80
81        value_table = self._build_value_table()
82
83        first_line = [''] + [self._header_string(header_value)
84                            for header_value in column_headers]
85        self._append_output_row(first_line)
86        self._process_value_table(value_table, row_headers)
87
88        return self._build_response()
89
90
91class TableCsvEncoder(CsvEncoder):
92    def __init__(self, request, response):
93        super(TableCsvEncoder, self).__init__(request, response)
94        self._column_specs = request['columns']
95
96
97    def _format_row(self, row_object):
98        """Extract data from a row object into a list of strings"""
99        return [row_object.get(field) for field, name in self._column_specs]
100
101
102    def _encode_table(self, row_objects):
103        self._append_output_row([column_spec[1] # header row
104                                 for column_spec in self._column_specs])
105        for row_object in row_objects:
106            self._append_output_row(self._format_row(row_object))
107        return self._build_response()
108
109
110    def encode(self):
111        return self._encode_table(self._response)
112
113
114class GroupedTableCsvEncoder(TableCsvEncoder):
115    def encode(self):
116        return self._encode_table(self._response['groups'])
117
118
119class StatusCountTableCsvEncoder(GroupedTableCsvEncoder):
120    _PASS_RATE_FIELD = '_test_pass_rate'
121
122    def __init__(self, request, response):
123        super(StatusCountTableCsvEncoder, self).__init__(request, response)
124        # inject a more sensible field name for test pass rate
125        for column_spec in self._column_specs:
126            field, name = column_spec
127            if name == 'Test pass rate':
128                column_spec[0] = self._PASS_RATE_FIELD
129                break
130
131
132    def _format_pass_rate(self, row_object):
133        result = '%s / %s' % (row_object['pass_count'],
134                              row_object['complete_count'])
135        incomplete_count = row_object['incomplete_count']
136        if incomplete_count:
137            result += ' (%s incomplete)' % incomplete_count
138        return result
139
140
141    def _format_row(self, row_object):
142        row_object[self._PASS_RATE_FIELD] = self._format_pass_rate(row_object)
143        return super(StatusCountTableCsvEncoder, self)._format_row(row_object)
144
145
146_ENCODER_MAP = {
147    'get_latest_tests' : SpreadsheetCsvEncoder,
148    'get_test_views' : TableCsvEncoder,
149    'get_group_counts' : GroupedTableCsvEncoder,
150}
151
152
153def _get_encoder_class(request):
154    method = request['method']
155    if method in _ENCODER_MAP:
156        return _ENCODER_MAP[method]
157    if method == 'get_status_counts':
158        if 'columns' in request:
159            return StatusCountTableCsvEncoder
160        return SpreadsheetCsvEncoder
161    return UnhandledMethodEncoder
162
163
164def encoder(request, response):
165    EncoderClass = _get_encoder_class(request)
166    return EncoderClass(request, response)
167