Исходный код EJIO.utils.dataset.operations

## /utils/dataset/operations.py

from pathlib import Path
from typing import Any
from ..system.file_utils import save_atomic
from .dataset_registry import DatasetRegistry

__all__ = [
    'save',
    'load',
    'register_storage_for_extension',
]

mapping = {
    'txt': 'csv',
    'csv': 'csv',
}




def register_storage_for_extension(extension: str, storage_name: str) -> None:
    """
    Выставляет соответствие между расширением файла и хранилищем, ответственным за файлы с этим расширением.

    :param extension: Расширение файла.
    :param storage_name: Имя хранилища, под которым оно зарегестрировано в реестре.
    :raises KeyError: Если соответствие для расширения уже задано.
    """
    if extension in mapping:
        raise KeyError(f"Serializer for extension '{extension}' already registered.")

    mapping[extension] = storage_name



def _detect_dataset_key(filename: str, forced_format: str | None) -> str:
    """
    Внутренний хелпер: автоматически определяет ключ типа хранилища.
    на основе расширения файла, если forced_format равен None или дефолтному 'csv'.

    :param filename: Имя файла.
    :param forced_format: Ключ типа хранилища.
    :return: Ключ типа хранилища для подходящего расширения файла, если соответствие задано,
        или forced_format если он был явно передан в функцию или csv для всех остальных случаев.
    """
    fmt = forced_format or Path(filename).suffix.lstrip('.').lower()

    detected_key = mapping.get(fmt)
    if detected_key and DatasetRegistry.has_key(detected_key):  # type: ignore[attr-defined]
        return detected_key

    return forced_format or 'csv'




def save[T](data: list[dict[str, T]], filename: str, *, format: str | None = None, **kwargs: Any) -> None:
    """
    Сохранить табличный набор данных (датасет) на диск через скрытый файл.

    :param data: Список словарей для сериализации.
    :param filename: Имя или путь к файлу.
    :param format: Строковый ключ типа хранилища. в глобальном регистре.
    :param kwargs: Дополнительные параметры.
    """
    actual_key = _detect_dataset_key(filename, format)
    dataset_cls = DatasetRegistry.get(actual_key)

    def execute_write(temp_path: Path) -> None:
        dataset_cls.save(data, str(temp_path), **kwargs)

    save_atomic(filename, execute_write)





def load[T](filename: str, *, format: str | None = None, **kwargs: Any) -> list[dict[str, T]]:
    """
    Загрузить табличный набор данных (датасет) из файла.

    :param filename: Имя или путь к файлу.
    :param format: Строковый ключ типа хранилища. в глобальном регистре.
    :param kwargs: Дополнительные параметры.
    :return: Список словарей.
    """
    actual_key = _detect_dataset_key(filename, format)
    dataset_cls = DatasetRegistry.get(actual_key)

    return dataset_cls.load(filename, **kwargs)