Django REST framework的各种技巧——7.导入导出

其实这个东西跟restframework没有卵关系，顺便写在这里

Django REST framework的各种技巧【目录索引】

导入导出在cms中是一个非常常用的功能，思考做成通用的东西，最终选择django-import-export，虽然这玩意儿最初是想放在admin中用的，虽然用起来很繁琐，但是可以做成通用的东西，而且用起来比较像rest的serializer。

django-import-export==0.4.2
文档

需要看的源码
cd 你的virtualenv/local/lib/python2.7/site-packages/import_export

resources.py
instance_loaders.py

先看用法

通过view可以看到，代码在这里是很干净的，跟一个正常的restframework的api没有什么区别。

class SchoolExportView(ExportMixin, GenericAPIView):

    serializer_class = SchoolSerializer
    permission_classes = (IsAuthenticated, ModulePermission)
    queryset = School.objects.filter(is_active=True).order_by('-id')
    resource_class = SchoolResource
    filter_backends = (filters.DjangoFilterBackend, filters.SearchFilter)
    filter_class = SchoolFilter
    search_fields = ('name', 'contact')
    module_perms = ['school.school']


class SchoolImportView(ImportMixin, GenericAPIView):

    serializer_class = SchoolSerializer
    permission_classes = (IsAuthenticated, ModulePermission)
    queryset = School.objects.filter(is_active=True).order_by('-id')
    resource_class = SchoolResource
    module_perms = ['school.school']

Mixin

class ExportMixin(object):

    @GET('filename', type='string', default='download.xls')
    @GET('format', type='string', default='xls', validators='in: xls,xlsx')
    @GET('empty', type='bool', default=False)
    def get(self, request, format, filename, empty):
        queryset = None
        if not empty:
            queryset = self.filter_queryset(self.get_queryset())
        resourse = self.resource_class()
        export_data = resourse.export(queryset, empty)
        return attachment_response(getattr(export_data, format), filename=filename)


class ImportMixin(object):

    @POST('file', validators='required')
    def post(self, request, file):
        import_file = request.FILES['file']
        resource = self.resource_class()
        extra_data = {} if not hasattr(self, 'get_resoucre_extra_data') else self.get_resoucre_extra_data()
        resource.set_extra_data(extra_data)
        dataset = resource.get_dataset(import_file)
        result = resource.import_data(dataset, use_transactions=True, raise_errors=True)
        return Response()

重点是实现Resource, 先说export

export非常简单，因此先说export,先看demo(仅仅写export)

# -*- coding: utf-8 -*-
from __future__ import absolute_import
from import_export import resources
from school.models import School
class SchoolResource(resources.ModelResource):
    def dehydrate_category(self, school):
        if school.category == School.MIDDLE_SCHOOL:
            return u'中学'
        elif school.category == School.COLLEGE:
            return u'高校'
        return ''
    def get_export_headers(self):
        return [u'分类', u'省份', u'城市', u'学校', u'地址', u'联系人',
                u'职务', u'联系电话', u'邮箱']
    class Meta:
        model = School
        fields = ('category', 'city__province__name', 'city__name',
                'name', 'address', 'contact', 'position', 'phone',
                'email')
        export_order = ('category', 'city__province__name', 'city__name',
                'name', 'address', 'contact', 'position', 'phone',
                'email')

resource的写法如下

• Meta中的fields指导出那些列，可以用外键的__方法

• Meta中的export_order指导出列的顺序

• get_export_headers是指excel的表头

• dehydrate_%filed%是指你可以对某一列做一些定制，同类似serializer里面的SerializerMethodField，但是只能是model上存在的%filed%才可以

再说import

由于import的复杂性，导致import的resource写起来非常复杂，因为import的时候有各种需求，例如导入了某些列但是只更新某些列，导入了很多列只更新不新建，导入列的各种数据校检...

首先基础的import_export中的InstanceLoader不能满足一个非常重要的查询需求，例如我们的model上面有is_active字段，然而又不能把这个东西导出区，导入的时候is_active又是get_instance的一个查询条件；以及ModelResource上面有些东西支持也非常不够，例如我输入一个文件就可以拿到dataset数据，例如我export的时候想传一个可以迭代的东西而不是queryset，还有给出更人性化的错误提示等等。

class ModelExtraParamInstanceLoader(BaseInstanceLoader):
    """ get_instance时支持额外的附加参数, 对说的就是is_active=True"""

    def get_queryset(self):
        return self.resource._meta.model.objects.all()

    def get_instance(self, row):
        try:
            params = self.resource._meta.import_instanceloader_extra_params
            for key in self.resource.get_import_id_fields():
                field = self.resource.fields[key]
                params[field.attribute] = field.clean(row)
            return self.get_queryset().get(**params)
        except self.resource._meta.model.DoesNotExist:
            return None

class ModelResource(resources.ModelResource):

    def set_extra_data(self, extra_data):
        self.extra_data = extra_data

    def get_clean_row(self, row):
        _row = []
        for each in row:
if isinstance(each, float):
                each = int(each)
            each = unicode(each).strip()
            _row.append(each)
        return _row

    def get_dataset_data(self, file_obj):
        '''从前端传来的excel获得原始数据'''
        headers = self.get_export_headers()
        try:
            self._dataset_data = get_data_from_excel(file_obj=file_obj, header=headers)
        except Exception as ex:
            logger.warn(ex)
            raise Error(
                errors.ExcelFormatError,
                err_message=unicode(ex),
                message=unicode(ex)
            )
        return self._dataset_data   
    def get_printable_row(self, row):
        _row = [unicode(each) for each in row]
        return u'({})'.format(u', '.join(_row))

    def get_printable_error_message(self, error_type, index, row):
        return u'excel格式错误:[{}]\n错误的行:{}行\n内容:{}'.format(
            error_type,
            index, self.get_printable_row(row)
        )

    def get_error(self, error_type, index, row):
        return Error(
            errors.ExcelFormatError,
            err_message='excel格式错误',
            message=self.get_printable_error_message(error_type, index, row)
        )    
    def clean_dataset_data(self, data):
        '''洗清原始数据，将data洗成跟diff_header一样的对应格式
            diff_header即model上面对应的列, 因为import_data是直接在model上做的
            有逻辑写在这里，比如参数检查，错误直接在这里raise Error，提示用户
        '''
        headers = self.get_export_headers()
        header_length = len(headers)
        for index, row in enumerate(data):
            if len(row) != header_length:
                raise self.get_error(u'列数错误', index+2, row)
        return data        
    def get_dataset(self, file_obj=None):
        assert hasattr(self, '_dataset_data') or file_obj, 'You need call get_dataset_data first or pass file_obj'
        if file_obj:
            data = self.get_dataset_data(file_obj)
        else:
            data = self._dataset_data
        data = self.clean_dataset_data(data)
        headers = self.get_diff_headers()
        dataset = get_dataset(data, headers)
        return dataset      
    def export(self, queryset=None, empty=False):
        """
        Exports a resource.
        """
        if queryset is None:
            if empty:
                if hasattr(self._meta, 'empty_export_data'):
                    queryset = self._meta.empty_export_data
                else:
                    queryset = []
            else:
                queryset = self.get_queryset()
        headers = self.get_export_headers()
        data = tablib.Dataset(headers=headers)

        if isinstance(queryset, QuerySet):
            # Iterate without the queryset cache, to avoid wasting memory when
            # exporting large datasets.
            iterable = queryset.iterator()
        else:
            iterable = queryset
        for obj in iterable:
            if empty and isinstance(obj, Iterable):    
                data.append(obj)
            else:
                data.append(self.export_resource(obj))
        return data

    def init_instance(self, row=None):
        if not row:
            row = {}
        instance = self._meta.model()
        for attr, value in row.iteritems():
            setattr(instance, attr, value)
        return instance                           

先给出一个没有复杂外键的model的导入Resource

class SchoolResource(ModelResource):

    def dehydrate_category(self, school):
        if school.category == School.MIDDLE_SCHOOL:
            return u'中学'
        elif school.category == School.COLLEGE:
            return u'高校'
        return ''

    def get_export_headers(self):
        return [u'分类', u'省份', u'城市', u'学校', u'地址', u'联系人',
                u'职务', u'联系电话', u'邮箱']

    def get_diff_headers(self):
        return ['category', 'city', 'name', 'address', 'contact', 'position', 'phone', 'email']

    def clean_dataset_data(self, data):
        data = super(SchoolResource, self).clean_dataset_data(data)
        clean_data = []
        for index, row in enumerate(data):
            _index = index + 2
            _row = self.get_clean_row(row)
            category = self.clean_dataset_category(_row[0], _index, row)
            city = self.clean_dataset_city((_row[1], _row[2]), _index, row)
            clean_data.append([category, city]+ _row[3:])
        return clean_data

    def clean_dataset_category(self, category, index, row):
        if category not in (u'中学', u'高校'):
            raise self.get_error(u'分类错误', index, row)
        if category == u'中学':
            return 1
        else:
            return 2

    class Meta:
        model = School
        import_id_fields = ['name',]
        import_instanceloader_extra_params = {'is_active': True}
        instance_loader_class = ModelExtraParamInstanceLoader
        empty_export_data = [...]          
        fields = ('category', 'city__province__name', 'city__name',
                'name', 'address', 'contact', 'position', 'phone',
                'email')
        export_order = ('category', 'city__province__name', 'city__name',
                'name', 'address', 'contact', 'position', 'phone',
                'email')          

resource的写法如下

• get_export_header是指导出时excel的表头

• get_diff_headers是指import时用的header，可以叫做任何东西（貌似是model上要能找到的，可以通过 外键__属性的方式）

• init_instance是指如果通过instanceloader没有get到数据时需要新建一条记录，可以根据传过来的row做一些事情，有时候你要做一些奇怪的事情，例如diff_header为city__name,但是我想把city设置为id，可以先通过clean_data拿到city然后在赋值，请见courseresource

• clean_dataset_data做数据清洗，对每行的数据都要做数据校检，可以做一些厉害的事情，比如city有关外键检查的事情，由于导出的时候city是用的city__name,city__province__name,导入就需要用这两列来确定一个city对象，看下下面clean_dataset_city的写法

• raise_error直接用self.get_error(u'分类错误', index, row)，第一个是说大的错误是毛，index是实际excel中的行，由于我们skip了header，如果你又用的enumerate做计数的话index应该+2

• class meta的import_id_fields，通过那几列来唯一确定数据，取值从git_diff_header中取，如果excel的信息不够(例如我们很多地方要写is_active=True)则填写import_instanceloader_extra_params

• skip_unchanged 是指如果excel的数据跟数据库的相同是否跟新数据，如果设为true有的时候会出问题

稍微复杂的demo

class CourseResource(ModelResource):

    def dehydrate_is_authentication(self, course):
        if course.is_authentication:
            return u'已认证'
        else:
            return u'未认证'

    def get_export_headers(self):
        return [
            u'年份', u'开课科目', u'学校', u'教师姓名',
            u'身份证号', u'联系方式', u'邮箱', u'教师认证状态',
            u'班级规模'
        ]

    def get_diff_headers(self):
        return ['term__name', 'name', 'school__name', 'teacher', 'ID_number', 'phone',
                'email', 'is_authentication', 'enrollment']

    def init_instance(self, row=None):
        if not row:
            row = {}
        instance = self._meta.model()
        for attr, value in row.iteritems():
            setattr(instance, attr, value)
        instance.term = row['term__name']
        instance.school = row['school__name']
        return instance

    def clean_dataset_data(self, data):
        data = super(CourseResource, self).clean_dataset_data(data)
        clean_data = []
        for index, row in enumerate(data):
            _index = index + 2
            _row = self.get_clean_row(row)
            term = self.clean_dataset_term(_row[0], _index, row)
            school = self.clean_dataset_school(_row[2], _index, row)
            is_authentication = self.clean_dataset_is_authentication(_row[7], _index, row)
            enrollment = self.clean_dataset_enrollment(_row[8], _index, row)
            clean_data.append([term, _row[1], school, _row[3], _row[4],
                _row[5], _row[6], is_authentication, enrollment])
        return clean_data

    def clean_dataset_term(self, term, index, row):
        try:
            return Term.objects.get(name=term, is_active=True)
        except Term.DoesNotExist:
            raise self.get_error(u'年份错误', index, row)     

    def clean_dataset_school(self, school, index, row):
        try:
            school = School.objects.get(name=school, is_active=True)
            user = self.extra_data['user']
            if not SchoolPermissionFilterBackend().has_school_permission(user,
                    school):
                raise self.get_error(u'没有对应的学校权限', index, row)
            return school
        except School.DoesNotExist:
            raise self.get_error(u'学校错误', index, row)

    def clean_dataset_is_authentication(self, is_authentication, index, row):
        if is_authentication == u'已认证':
            return True
        if is_authentication == u'未认证':
            return False
        raise self.get_error(u'教师认证状态错误', index, row)
    def clean_dataset_enrollment(self, enrollment, index, row):
        try:
            if not enrollment:
                enrollment = 0
            return int(float(enrollment))
        except:
            raise self.get_error(u'班级规模错误', index, row)

    class Meta:
        model = Course
        import_id_fields = ['term__name', 'name', 'school__name']
        import_instanceloader_extra_params = {
                'is_active': True, 'term__is_active': True, 'school__is_active': True}
        instance_loader_class = ModelExtraParamInstanceLoader
        fields = ('term__name', 'name', 'school__name',
                'teacher', 'ID_number', 'phone', 'email', 'is_authentication',
                'enrollment')
        export_order = ('term__name', 'name', 'school__name',
                'teacher', 'ID_number', 'phone', 'email', 'is_authentication',
                'enrollment')
        empty_export_data = [...]               

几个方法

def extract_data(sheet, header, skip_header=True, row_type='list'):
    assert header
    data = []
    for row_index in xrange(1 if skip_header else 0, sheet.nrows):
        row = sheet.row_values(row_index)
        assert len(header) == len(row), u'excel 第{}行,列数对应数据不对'.format(row_index)
        if row_type == 'list':
            data.append(row)
        else:
            each_data = {}
            for col_index in xrange(len(header)):
                each_data[header[col_index]] = row[col_index]
            data.append(each_data)
    return data
    
def get_data_from_excel(file_path=None, file_obj=None, header=None,
        sheet_index=0, skip_header=True):
    assert header
    assert file_path or file_obj
    if file_path:
        with open_workbook(file_path) as wb: 
            data = extract_data(wb.sheet_by_index(sheet_index), header, skip_header)
    else:
        with tempinput(file_obj) as tempfilename:
            with open_workbook(tempfilename) as wb:
                data = extract_data(wb.sheet_by_index(sheet_index), header, skip_header)
    return data


def get_dataset(data, header):
    return tablib.Dataset(*data, headers=header)    

def attachment_response(export_data, filename='download.xls', content_type='application/vnd.ms-excel'):
    # Django 1.7 uses the content_type kwarg instead of mimetype
    try:
        response = HttpResponse(export_data, content_type=content_type)
    except TypeError:
        response = HttpResponse(export_data, mimetype=content_type)
    response['Content-Disposition'] = 'attachment; filename={}'.format(filename)
    return response