<reponame>nvonahsen/toolbelt<filename>colab/functions.py
from typing import Callable
from ..config import get_dir
import os
import shutil
from google.colab import files

def improved_file_upload(dir: str=None, clear_each_time: bool=False, no_loop: bool=False):
    """    
    Use colab file uploader to upload files and run function for each one


    Parameters
    ----------
    dir : str, optional
        Directory used for storing images, by default aquired from config/temp_dir
    clear_each_time : bool, optional
        Clear uploaded images each time, if false existing will be used, by default False

    Returns
    -------
    list
        List of paths to all files uploaded with this tool
    """

    if dir is None:
        dir = get_dir("uploaded_temp")
    else:
        dir = get_dir(dir, external=True)

    if clear_each_time:
        shutil.rmtree(dir)
    else:
        print(f"Found {len(os.listdir(dir))} existing images, upload more?")

    uploaded = []
    while True:
        new = files.upload()
        if len(new) == 0:
            break
        uploaded += new
        if no_loop:
            break
        print()
        print(f"Uploaded files (total {len(uploaded)}), upload more?")

    for img_name in uploaded:
        shutil.move(f"{img_name}", os.path.join(dir, img_name))

    paths = [os.path.join(dir, d) for d in os.listdir(dir)]
    return [p for p in paths if os.path.isfile(p)]