Fixed bug with StringIO import, incompatible with python 3

This commit is contained in:
M1ha 2017-04-06 10:54:00 +05:00
parent 9bca42fe7e
commit 76718426bf

View File

@ -50,7 +50,7 @@ class Database(object):
def insert(self, model_instances, batch_size=1000): def insert(self, model_instances, batch_size=1000):
from six import next from six import next
from cStringIO import StringIO from io import BytesIO
i = iter(model_instances) i = iter(model_instances)
try: try:
first_instance = next(i) first_instance = next(i)
@ -62,7 +62,7 @@ class Database(object):
raise DatabaseException("You can't insert into read only table") raise DatabaseException("You can't insert into read only table")
def gen(): def gen():
buf = StringIO() buf = BytesIO()
buf.write(self._substitute('INSERT INTO $table FORMAT TabSeparated\n', model_class).encode('utf-8')) buf.write(self._substitute('INSERT INTO $table FORMAT TabSeparated\n', model_class).encode('utf-8'))
first_instance.set_database(self) first_instance.set_database(self)
buf.write(first_instance.to_tsv(include_readonly=False).encode('utf-8')) buf.write(first_instance.to_tsv(include_readonly=False).encode('utf-8'))
@ -78,7 +78,7 @@ class Database(object):
# Return the current batch of lines # Return the current batch of lines
yield buf.getvalue() yield buf.getvalue()
# Start a new batch # Start a new batch
buf = StringIO() buf = BytesIO()
lines = 0 lines = 0
# Return any remaining lines in partial batch # Return any remaining lines in partial batch
if lines: if lines: