corrlib/corrlib/tracker.py

import os
import shutil
import warnings
from configparser import ConfigParser
from pathlib import Path

import datalad.api as dl

from .tools import CONFIG_FILENAME, get_db_file


def get_tracker(path: Path) -> str:
    """
    Get the tracker used in the dataset located at path.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.

    Returns
    -------
    tracker: str
        The tracker used in the dataset.
    """
    path = Path(path)
    config_path = path / CONFIG_FILENAME
    config = ConfigParser()
    if os.path.exists(config_path):
        config.read(config_path)
    else:
        raise FileNotFoundError(f"No config file found in {path}.")
    tracker = config.get('core', 'tracker', fallback='datalad')
    return tracker


def get(path: Path, file: Path) -> None:
    """
    Wrapper function to get a file from the dataset located at path with the specified tracker.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    file: str
        The file to get.
    """
    path = Path(path)
    tracker = get_tracker(path)
    if tracker == 'datalad':
        if file == get_db_file(path):
            print("Downloading database...")
        else:
            print("Downloading data...")
        dl.get(os.path.join(path, file), dataset=path)
        print("> downloaded file")
    elif tracker == 'None':
        pass
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return


def save(path: Path, message: str, files: list[Path] | None=None) -> None:
    """
    Wrapper function to save a file to the dataset located at path with the specified tracker.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    message: str
        The commit message.
    files: list[str], optional
        The files to save. If None, all changes are saved.
    """
    path = Path(path)
    tracker = get_tracker(path)
    if tracker == 'datalad':
        if files is not None:
            files = [path / f for f in files]
        dl.save(files, message=message, dataset=path)
    elif tracker == 'None':
        warnings.warn("Tracker 'None' does not implement save.", Warning, 1)
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")


def init(path: Path, tracker: str='datalad') -> None:
    """
    Initialize a dataset at the specified path with the specified tracker.

    Parameters
    ----------
    path: str
        The path to initialize the dataset.
    tracker: str
        The tracker to use. Currently only 'datalad' and 'None' are supported.
    """
    path = Path(path)
    if tracker == 'datalad':
        dl.create(path)
    elif tracker == 'None':
        os.makedirs(path, exist_ok=True)
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return


def unlock(path: Path, file: Path) -> None:
    """
    Wrapper function to unlock a file in the dataset located at path with the specified tracker.

    Parameters
    ----------
    path    : str
        The path to the backlogger folder.
    file    : str
        The file to unlock.
    """
    path = Path(path)
    tracker = get_tracker(path)
    if tracker == 'datalad':
        dl.unlock(os.path.join(path, file), dataset=path)
    elif tracker == 'None':
        warnings.warn("Tracker 'None' does not implement unlock.", Warning, 1)
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return


def clone(path: Path, source: str, target: str) -> None:
    """
    Wrapper function to clone a dataset from source to target with the specified tracker.
    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    source: str
        The source dataset to clone.
    target: str
        The target path to clone the dataset to.
    """
    path = Path(path)
    tracker = get_tracker(path)
    if tracker == 'datalad':
        dl.clone(path=target, source=source, dataset=path)
    elif tracker == 'None':
        os.makedirs(path, exist_ok=True)
        # Implement a simple clone by copying files
        shutil.copytree(source, target, dirs_exist_ok=False)
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return


def drop(path: Path, reckless: str | None=None) -> None:
    """
    Wrapper function to drop data from a dataset located at path with the specified tracker.

    Parameters
    ----------
    path: str
        The path to the backlogger folder.
    reckless: Optional[str]
        The datalad's reckless option for dropping data.
    """
    path = Path(path)
    tracker = get_tracker(path)
    if tracker == 'datalad':
        dl.drop(path, reckless=reckless)
    elif tracker == 'None':
        warnings.warn("Tracker 'None' does not implement drop.", Warning, 1)
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return