From 3c6af9a6105bc51a2d06f4a978e40d0e5f5a0dee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Herv=C3=A9=20=20MENAGER?= <herve.menager@pasteur.fr>
Date: Mon, 4 Mar 2019 18:09:09 +0100
Subject: [PATCH] fix progress bar in dataframe processing
---
django_diu/import_command.py | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/django_diu/import_command.py b/django_diu/import_command.py
index 515b81a..83b62a9 100644
--- a/django_diu/import_command.py
+++ b/django_diu/import_command.py
@@ -73,8 +73,8 @@ class ImportTask(object):
def post_process(self):
pass
- def _process_rows(self, rows):
- for idx, row in enumerate(rows):
+ def _process_rows(self, rows, progress=None):
+ for row in rows:
try:
new_object = None
with transaction.atomic():
@@ -93,8 +93,8 @@ class ImportTask(object):
if self.progress_bar == False:
self.out_stream.write(self.style.SUCCESS(
'Successfully inserted {}'.format(new_object)))
- if self.progress_bar:
- self.rows.update(idx)
+ if self.progress_bar == True:
+ progress.update(1)
def count_source(self):
self.source_count = len(self.rows)
@@ -104,8 +104,8 @@ class ImportTask(object):
self.count_source()
if self.progress_bar is True:
with click.progressbar(length=self.source_count,
- label='Importing ' + self.description + ' (' + str(self.source_count) + ' rows to process).') as rows_list:
- self._process_rows(rows_list)
+ label='Importing ' + self.description + ' (' + str(self.source_count) + ' rows to process).') as progress:
+ self._process_rows(self.rows, progress)
else:
self._process_rows(self.rows)
self.post_process()
@@ -152,7 +152,7 @@ class DataFrameImportTask(ImportTask):
dataframe = None
def open_data_source(self):
- self.rows = self.dataframe.iterrows()
+ self.rows = (row for idx, row in self.dataframe.iterrows())
def count_source(self):
self.source_count = self.dataframe.shape[0]
--
GitLab