CSVField for forms

import csv

from django import forms
from django.core.exceptions import ValidationError


class CSVField(forms.FileField):
    """
    FileField that checks that the file is a valid CSV and if specified
    in `expected_fieldnames` checks that the fields match exactly.

    The widget's `accept` parameter is set to accept csv, text and excel files.

    TODO: Validate the entirety of the CSV file, not just the headers.
          But this should be enough for most use cases, as checking the 
          whole file could be computationally expensive for huge files.

    Example usage:
        people = CSVField(expected_fieldnames=['First Name', 'Last Name'])
    """

    def __init__(self, *args, **kwargs):
        self.expected_fieldnames = kwargs.pop('expected_fieldnames', None)
        super(CSVField, self).__init__(*args, **kwargs)
        self.error_messages['required'] = 'You must select a file'
        self.widget.attrs.update(
            {'accept': '.csv,'
                       'text/plain,'
                       'application/vnd.ms-excel,'
                       'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'})

    def clean(self, data, initial=None):
        value = super(CSVField, self).clean(data)
        reader = csv.reader(data)
        # Check it's a valid CSV file
        try:
            fieldnames = reader.next()
        except csv.Error:
            raise ValidationError('You must upload a valid CSV file')

        # Check the fieldnames are as specified, if requested
        if self.expected_fieldnames and fieldnames != self.expected_fieldnames:
            raise ValidationError(
                u'The CSV fields are expected to be "{0}"'.format(
                    u','.join(self.expected_fieldnames)))

        return value