Browse Source

update writer.py

fix csv writer: first close filename, then write zip file, or it will zip an empty file
fix json and jsonL writer: add ensure_ascii=False, now it can write chinese in file, but not \uxxxx
pull/1340/head
pq4362050 3 years ago
committed by GitHub
parent
commit
a8efafc286
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 6 deletions
  1. 13
      backend/api/views/download/writer.py

13
backend/api/views/download/writer.py

@ -64,7 +64,7 @@ class CsvWriter(BaseWriter):
for record in records: for record in records:
filename = os.path.join(self.tmpdir, f'{record.user}.{self.extension}') filename = os.path.join(self.tmpdir, f'{record.user}.{self.extension}')
if filename not in writers: if filename not in writers:
f = open(filename, mode='a')
f = open(filename, mode='a', encoding='utf-8')
writer = csv.DictWriter(f, header) writer = csv.DictWriter(f, header)
writer.writeheader() writer.writeheader()
writers[filename] = writer writers[filename] = writer
@ -72,11 +72,12 @@ class CsvWriter(BaseWriter):
writer = writers[filename] writer = writers[filename]
line = self.create_line(record) line = self.create_line(record)
writer.writerow(line) writer.writerow(line)
for f in file_handlers:
f.close()
save_file = self.write_zip(writers) save_file = self.write_zip(writers)
for file in writers: for file in writers:
os.remove(file) os.remove(file)
for f in file_handlers:
f.close()
return save_file return save_file
def create_line(self, record) -> Dict: def create_line(self, record) -> Dict:
@ -102,14 +103,14 @@ class JSONWriter(BaseWriter):
for record in records: for record in records:
filename = os.path.join(self.tmpdir, f'{record.user}.{self.extension}') filename = os.path.join(self.tmpdir, f'{record.user}.{self.extension}')
if filename not in writers: if filename not in writers:
f = open(filename, mode='a')
f = open(filename, mode='a', encoding='utf-8')
writers[filename] = f writers[filename] = f
line = self.create_line(record) line = self.create_line(record)
contents[filename].append(line) contents[filename].append(line)
for filename, f in writers.items(): for filename, f in writers.items():
content = contents[filename] content = contents[filename]
json.dump(content, f)
json.dump(content, f, ensure_ascii=False)
f.close() f.close()
save_file = self.write_zip(writers) save_file = self.write_zip(writers)
@ -135,7 +136,7 @@ class JSONLWriter(LineWriter):
'data': record.data, 'data': record.data,
'label': record.label, 'label': record.label,
**record.metadata **record.metadata
})
},ensure_ascii=False)
class FastTextWriter(LineWriter): class FastTextWriter(LineWriter):

Loading…
Cancel
Save