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