add progress bar when importing multiple files
This commit is contained in:
@@ -1,10 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from rich import print
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
MofNCompleteColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TextColumn,
|
||||
)
|
||||
|
||||
from chromy.errors import UnsupportedTextFileError
|
||||
from chromy.utilities import ingest_file
|
||||
@@ -42,34 +50,76 @@ def _import_one(collection: str, file: str) -> int:
|
||||
if not is_probably_text_file(absolute_path):
|
||||
raise UnsupportedTextFileError()
|
||||
|
||||
records_added = ingest_file(collection, absolute_path)
|
||||
print(
|
||||
"[bold green]Added[/] "
|
||||
f"{records_added} records from '{file}' to collection '{collection}'."
|
||||
)
|
||||
return SUCCESS_EXIT_CODE
|
||||
return ingest_file(collection, absolute_path)
|
||||
|
||||
|
||||
def _should_show_progress(file_count: int) -> bool:
|
||||
return file_count > 1 and sys.stdout.isatty()
|
||||
|
||||
|
||||
def _truncate_file_name(file_name: str, max_length: int = 20) -> str:
|
||||
if len(file_name) <= max_length:
|
||||
return file_name
|
||||
|
||||
return f"{file_name[: max_length - 3]}"
|
||||
|
||||
|
||||
def handle_import(collection: str, files: list[str]) -> int:
|
||||
successful_imports = 0
|
||||
failed_imports = 0
|
||||
seen_paths: set[str] = set()
|
||||
unique_files: list[str] = []
|
||||
|
||||
for file in files:
|
||||
try:
|
||||
absolute_path = _get_absolute_path(file)
|
||||
if absolute_path in seen_paths:
|
||||
continue
|
||||
|
||||
seen_paths.add(absolute_path)
|
||||
_import_one(collection, file)
|
||||
successful_imports += 1
|
||||
except FileNotFoundError:
|
||||
failed_imports += 1
|
||||
print(f"[bold red]Error[/]: The file '{file}' was not found.")
|
||||
except UnsupportedTextFileError:
|
||||
failed_imports += 1
|
||||
print(f"[bold red]Error[/]: The file '{file}' is not a text file.")
|
||||
unique_files.append(file)
|
||||
continue
|
||||
|
||||
if absolute_path in seen_paths:
|
||||
continue
|
||||
|
||||
seen_paths.add(absolute_path)
|
||||
unique_files.append(file)
|
||||
|
||||
show_progress = _should_show_progress(len(unique_files))
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
MofNCompleteColumn(),
|
||||
transient=True,
|
||||
disable=not show_progress,
|
||||
) as progress:
|
||||
task_id = progress.add_task("Importing files...", total=len(unique_files))
|
||||
|
||||
for file in unique_files:
|
||||
file_name = _truncate_file_name(Path(file).name)
|
||||
description = f"Importing [bold]{file_name}[/]..."
|
||||
progress.update(task_id, description=description)
|
||||
try:
|
||||
records_added = _import_one(collection, file)
|
||||
successful_imports += 1
|
||||
if not show_progress:
|
||||
progress.console.print(
|
||||
"[bold green]Added[/] "
|
||||
f"{records_added} records from '{file}' to "
|
||||
f"collection '{collection}'."
|
||||
)
|
||||
except FileNotFoundError:
|
||||
failed_imports += 1
|
||||
progress.console.print(
|
||||
f"[bold red]Error[/]: The file '{file}' was not found."
|
||||
)
|
||||
except UnsupportedTextFileError:
|
||||
failed_imports += 1
|
||||
progress.console.print(
|
||||
f"[bold red]Error[/]: The file '{file}' is not a text file."
|
||||
)
|
||||
finally:
|
||||
progress.advance(task_id)
|
||||
|
||||
print(
|
||||
f"Imported {successful_imports} file(s) successfully; {failed_imports} failed."
|
||||
|
||||
Reference in New Issue
Block a user