feat: add export
This commit is contained in:
101
backend/utils/export_mixin.py
Normal file
101
backend/utils/export_mixin.py
Normal file
@@ -0,0 +1,101 @@
|
||||
import pandas as pd
|
||||
|
||||
from rest_framework.decorators import action
|
||||
from django.http import HttpResponse
|
||||
|
||||
class ExportMixin:
|
||||
"""
|
||||
导出功能 Mixin,提供数据导出为 Excel 或 CSV 的功能
|
||||
"""
|
||||
# 导出配置
|
||||
export_fields = {} # 字段映射配置
|
||||
export_filename = None # 导出文件名
|
||||
|
||||
@action(detail=False, methods=['get'], url_path='export')
|
||||
def export_data(self, request):
|
||||
"""
|
||||
导出数据功能
|
||||
支持通过参数控制导出字段:
|
||||
- fields: 指定要导出的字段,多个字段用逗号分隔
|
||||
- format: 导出格式 (excel/csv),默认excel
|
||||
"""
|
||||
# 获取查询集
|
||||
queryset = self.filter_queryset(self.get_queryset())
|
||||
|
||||
# 获取导出字段配置
|
||||
export_config = self.get_export_fields(request)
|
||||
|
||||
# 序列化数据
|
||||
serializer = self.get_serializer(queryset, many=True)
|
||||
data = serializer.data
|
||||
|
||||
# 处理数据
|
||||
processed_data = self.process_export_data(data, export_config)
|
||||
|
||||
# 生成文件名
|
||||
filename = self.export_filename or f"{self.__class__.__name__}_export"
|
||||
|
||||
# 根据格式返回响应
|
||||
export_format = request.query_params.get('format', 'excel').lower()
|
||||
|
||||
if export_format == 'csv':
|
||||
return self.generate_csv_response(processed_data, filename)
|
||||
else:
|
||||
return self.generate_excel_response(processed_data, filename)
|
||||
|
||||
def get_export_fields(self, request):
|
||||
"""
|
||||
获取导出字段配置
|
||||
支持通过URL参数指定字段: ?fields=id,name,email
|
||||
"""
|
||||
# 优先使用请求参数中的字段
|
||||
fields_param = request.query_params.get('fields')
|
||||
if fields_param:
|
||||
field_names = [f.strip() for f in fields_param.split(',')]
|
||||
# 只返回指定字段的配置
|
||||
return {field: self.export_fields.get(field, field) for field in field_names}
|
||||
# 默认返回所有配置字段
|
||||
return self.export_fields or {}
|
||||
|
||||
def process_export_data(self, data, export_config):
|
||||
"""
|
||||
处理导出数据,根据配置映射字段名和处理数据
|
||||
"""
|
||||
if not export_config:
|
||||
return data
|
||||
|
||||
processed_data = []
|
||||
for item in data:
|
||||
processed_item = {}
|
||||
for field_key, field_config in export_config.items():
|
||||
# field_config 可以是字符串(列名)或字典(包含列名和处理函数)
|
||||
if isinstance(field_config, dict):
|
||||
column_name = field_config.get('name', field_key)
|
||||
processor = field_config.get('processor')
|
||||
value = item.get(field_key)
|
||||
if processor and callable(processor):
|
||||
value = processor(value, item)
|
||||
processed_item[column_name] = value
|
||||
else:
|
||||
# 简单映射
|
||||
processed_item[field_config] = item.get(field_key, '')
|
||||
processed_data.append(processed_item)
|
||||
return processed_data
|
||||
|
||||
def generate_excel_response(self, data, filename):
|
||||
"""生成Excel格式响应"""
|
||||
df = pd.DataFrame(data)
|
||||
response = HttpResponse(
|
||||
content_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
|
||||
)
|
||||
response['Content-Disposition'] = f'attachment; filename="{filename}.xlsx"'
|
||||
df.to_excel(response, index=False)
|
||||
return response
|
||||
|
||||
def generate_csv_response(self, data, filename):
|
||||
"""生成CSV格式响应"""
|
||||
df = pd.DataFrame(data)
|
||||
response = HttpResponse(content_type='text/csv')
|
||||
response['Content-Disposition'] = f'attachment; filename="{filename}.csv"'
|
||||
df.to_csv(response, index=False)
|
||||
return response
|
||||
Reference in New Issue
Block a user