diff --git a/apps/common/drf/const.py b/apps/common/drf/const.py new file mode 100644 index 000000000..50a415fea --- /dev/null +++ b/apps/common/drf/const.py @@ -0,0 +1 @@ +CSV_FILE_ESCAPE_CHARS = ['=', '@', '0'] diff --git a/apps/common/drf/parsers/csv.py b/apps/common/drf/parsers/csv.py index 0dd11aa4b..adfc56e0e 100644 --- a/apps/common/drf/parsers/csv.py +++ b/apps/common/drf/parsers/csv.py @@ -4,13 +4,25 @@ import chardet import unicodecsv +from common.utils import lazyproperty from .base import BaseFileParser +from ..const import CSV_FILE_ESCAPE_CHARS class CSVFileParser(BaseFileParser): - media_type = 'text/csv' + @lazyproperty + def match_escape_chars(self): + chars = [] + for c in CSV_FILE_ESCAPE_CHARS: + dq_char = '"{}'.format(c) + sg_char = "'{}".format(c) + chars.append(dq_char) + chars.append(sg_char) + return tuple(chars) + + @staticmethod def _universal_newlines(stream): """ @@ -18,6 +30,14 @@ class CSVFileParser(BaseFileParser): """ for line in stream.splitlines(): yield line + + def __parse_row(self, row): + row_escape = [] + for d in row: + if isinstance(d, str) and d.strip().startswith(self.match_escape_chars): + d = d.lstrip("'").lstrip('"') + row_escape.append(d) + return row_escape def generate_rows(self, stream_data): detect_result = chardet.detect(stream_data) @@ -25,4 +45,5 @@ class CSVFileParser(BaseFileParser): lines = self._universal_newlines(stream_data) csv_reader = unicodecsv.reader(lines, encoding=encoding) for row in csv_reader: + row = self.__parse_row(row) yield row diff --git a/apps/common/drf/renders/csv.py b/apps/common/drf/renders/csv.py index 88466f533..a8729294b 100644 --- a/apps/common/drf/renders/csv.py +++ b/apps/common/drf/renders/csv.py @@ -6,31 +6,36 @@ import unicodecsv from six import BytesIO from .base import BaseFileRenderer - +from ..const import CSV_FILE_ESCAPE_CHARS class CSVFileRenderer(BaseFileRenderer): + media_type = 'text/csv' format = 'csv' writer = None buffer = None + escape_chars = tuple(CSV_FILE_ESCAPE_CHARS) + def initial_writer(self): csv_buffer = BytesIO() csv_buffer.write(codecs.BOM_UTF8) csv_writer = unicodecsv.writer(csv_buffer, encoding='utf-8') self.buffer = csv_buffer self.writer = csv_writer - - def write_row(self, row): + + def __render_row(self, row): row_escape = [] for d in row: - if isinstance(d, str) and d.strip().startswith(('=', '@')): + if isinstance(d, str) and d.strip().startswith(self.escape_chars): d = "'{}".format(d) - row_escape.append(d) - else: - row_escape.append(d) - self.writer.writerow(row_escape) + row_escape.append(d) + return row_escape + + def write_row(self, row): + row = self.__render_row(row) + self.writer.writerow(row) def get_rendered_value(self): value = self.buffer.getvalue()