You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
jumpserver/apps/common/parsers/csv.py

102 lines
3.0 KiB

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

# ~*~ coding: utf-8 ~*~
#
import json
import unicodecsv
from rest_framework.parsers import BaseParser
from rest_framework.exceptions import ParseError
from ..utils import get_logger
logger = get_logger(__file__)
class JMSCSVParser(BaseParser):
"""
Parses CSV file to serializer data
"""
media_type = 'text/csv'
@staticmethod
def _universal_newlines(stream):
"""
保证在`通用换行模式`下打开文件
"""
for line in stream.splitlines():
yield line
@staticmethod
def _gen_rows(csv_data, charset='utf-8', **kwargs):
csv_reader = unicodecsv.reader(csv_data, encoding=charset, **kwargs)
for row in csv_reader:
if not any(row): # 空行
continue
yield row
@staticmethod
def _get_fields_map(serializer):
fields_map = {}
fields = serializer.get_fields()
fields_map.update({v.label: k for k, v in fields.items()})
fields_map.update({k: k for k, _ in fields.items()})
return fields_map
@staticmethod
def _process_row(row):
"""
构建json数据前的行处理
"""
_row = []
for col in row:
# 列表转换
if isinstance(col, str) and col.find("[") != -1 and col.find("]") != -1:
# 替换中文格式引号
col = col.replace("", '"').replace("", '"').\
replace("", '"').replace('', '"').replace("'", '"')
col = json.loads(col)
_row.append(col)
return _row
@staticmethod
def _process_row_data(row_data):
"""
构建json数据后的行数据处理
"""
_row_data = {}
for k, v in row_data.items():
if isinstance(v, list) \
or isinstance(v, str) and k.strip() and v.strip():
_row_data[k] = v
return _row_data
def parse(self, stream, media_type=None, parser_context=None):
parser_context = parser_context or {}
encoding = parser_context.get('encoding', 'utf-8')
try:
serializer = parser_context["view"].get_serializer()
except Exception as e:
logger.debug(e, exc_info=True)
raise ParseError('The resource does not support imports!')
try:
stream_data = stream.read()
binary = self._universal_newlines(stream_data)
rows = self._gen_rows(binary, charset=encoding)
header = next(rows)
fields_map = self._get_fields_map(serializer)
header = [fields_map.get(name, '') for name in header]
data = []
for row in rows:
row = self._process_row(row)
row_data = dict(zip(header, row))
row_data = self._process_row_data(row_data)
data.append(row_data)
return data
except Exception as e:
logger.debug(e, exc_info=True)
raise ParseError('CSV parse error!')