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