Chromy/chromy/handlers/import_data.py

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

from ..utilities import is_probably_text_file

SUCCESS_EXIT_CODE: Final = 0
FAILURE_EXIT_CODE: Final = 1


def _get_absolute_path(file: str) -> str:
    """
    A helper method that, given a valid relative path to a file, returns its
    absolute path.

    Args:
        file (str): The relative path to the file.

    Raises:
        FileNotFoundError(): If the file does not exist.
    """
    if not os.path.exists(file):
        raise FileNotFoundError()

    file_path = Path(file)
    return str(file_path.resolve())


def _import_one(collection: str, file: str) -> int:
    absolute_path = _get_absolute_path(file)

    if not Path(absolute_path).is_file():
        raise FileNotFoundError()

    if not is_probably_text_file(absolute_path):
        raise UnsupportedTextFileError()

    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)
        except FileNotFoundError:
            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."
    )

    if failed_imports:
        return FAILURE_EXIT_CODE

    return SUCCESS_EXIT_CODE