Correctly start task in ketos compile progress bar

Fixes #504.
mittagessen · May 10, 2024 · 3899b49 · 3899b49
1 parent 2fb846a
commit 3899b49
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/kraken/ketos/dataset.py b/kraken/ketos/dataset.py
@@ -84,6 +84,11 @@ def compile(ctx, output, workers, format_type, files, random_split, force_type,
     with KrakenProgressBar() as progress:
         extract_task = progress.add_task('Extracting lines', total=0, start=False, visible=True if not ctx.meta['verbose'] else False)
 
+        def _update_bar(advance, total):
+            if not progress.tasks[0].started:
+                progress.start_task(extract_task)
+            progress.update(extract_task, total=total, advance=advance)
+
         arrow_dataset.build_binary_dataset(ground_truth,
                                            output,
                                            format_type,
@@ -93,7 +98,7 @@ def compile(ctx, output, workers, format_type, files, random_split, force_type,
                                            force_type,
                                            recordbatch_size,
                                            skip_empty_lines,
-                                           lambda advance, total: progress.update(extract_task, total=total, advance=advance),
+                                           _update_bar,
                                            legacy_polygons=legacy_polygons)
 
     message(f'Output file written to {output}')