fix progress bar in dataframe processing

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]