showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 1 | import csv |
| 2 | import django.http |
| 3 | import common |
| 4 | from autotest_lib.frontend.afe import rpc_utils |
| 5 | |
| 6 | class 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 NotImplemented |
| 28 | |
| 29 | |
| 30 | class UnhandledMethodEncoder(CsvEncoder): |
| 31 | def encode(self): |
| 32 | return rpc_utils.raw_http_response( |
showard | 37059df | 2009-05-08 02:30:20 +0000 | [diff] [blame] | 33 | 'Unhandled method %s (this indicates a bug)\r\n' % |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 34 | self._request['method']) |
| 35 | |
| 36 | |
| 37 | class 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'] |
showard | 194a59d | 2009-06-08 23:25:48 +0000 | [diff] [blame] | 47 | if 'extra_info' in group: |
| 48 | result = '\n'.join([result] + group['extra_info']) |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 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 | |
showard | 37059df | 2009-05-08 02:30:20 +0000 | [diff] [blame] | 91 | class 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 | |
| 114 | class GroupedTableCsvEncoder(TableCsvEncoder): |
| 115 | def encode(self): |
| 116 | return self._encode_table(self._response['groups']) |
| 117 | |
| 118 | |
| 119 | class 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 | |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 146 | _ENCODER_MAP = { |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 147 | 'get_latest_tests' : SpreadsheetCsvEncoder, |
showard | 37059df | 2009-05-08 02:30:20 +0000 | [diff] [blame] | 148 | 'get_test_views' : TableCsvEncoder, |
| 149 | 'get_group_counts' : GroupedTableCsvEncoder, |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 150 | } |
| 151 | |
| 152 | |
showard | 37059df | 2009-05-08 02:30:20 +0000 | [diff] [blame] | 153 | def _get_encoder_class(request): |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 154 | method = request['method'] |
showard | 37059df | 2009-05-08 02:30:20 +0000 | [diff] [blame] | 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 | |
| 164 | def encoder(request, response): |
| 165 | EncoderClass = _get_encoder_class(request) |
showard | 3d6ae11 | 2009-05-02 00:45:48 +0000 | [diff] [blame] | 166 | return EncoderClass(request, response) |