corrlib/corrlib/tracker.py

import os
from configparser import ConfigParser
import datalad.api as dl
from typing import Optional
import shutil
from .tools import get_db_file


def get_tracker(path: str) -> 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.
    """
    config_path = os.path.join(path, '.corrlib')
    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: str, file: str) -> 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.
    """
    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: str, message: str, files: Optional[list[str]]=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.
    """
    tracker = get_tracker(path)
    if tracker == 'datalad':
        if files is not None:
            files = [os.path.join(path, f) for f in files]
        dl.save(files, message=message, dataset=path)
    elif tracker == 'None':
        Warning("Tracker 'None' does not implement save.")
        pass
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")


def init(path: str, 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.
    """
    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: str, file: str) -> 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.
    """
    tracker = get_tracker(path)
    if tracker == 'datalad':
        dl.unlock(file, dataset=path)
    elif tracker == 'None':
        Warning("Tracker 'None' does not implement unlock.")
        pass
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return


def clone(path: str, 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.
    """
    tracker = get_tracker(path)
    if tracker == 'datalad':
        dl.clone(target=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: str, reckless: Optional[str]=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.
    """
    tracker = get_tracker(path)
    if tracker == 'datalad':
        dl.drop(path, reckless=reckless)
    elif tracker == 'None':
        Warning("Tracker 'None' does not implement drop.")
        pass
    else:
        raise ValueError(f"Tracker {tracker} is not supported.")
    return