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')