Home | History | Annotate | Download | only in tko 
      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 NotImplementedError
     28 
     29 
     30 class 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 
     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']
     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 
     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 
    146 _ENCODER_MAP = {
    147     'get_latest_tests' : SpreadsheetCsvEncoder,
    148     'get_test_views' : TableCsvEncoder,
    149     'get_group_counts' : GroupedTableCsvEncoder,
    150 }
    151 
    152 
    153 def _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 
    164 def encoder(request, response):
    165     EncoderClass = _get_encoder_class(request)
    166     return EncoderClass(request, response)
    167