IOHelperUtilities.py

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/helper_utilities.ipynb.

# %% auto 0
__all__ = ['check_is_colab', 'MultiFileChooser', 'setup_drives']

# %% ../nbs/helper_utilities.ipynb 3
import ipywidgets as widgets
from IPython.display import display, clear_output
from functools import partial
from ipyfilechooser import FileChooser
import os

# %% ../nbs/helper_utilities.ipynb 4
def check_is_colab():
    """
    Check if the current environment is Google Colab.
    """
    try:
        import google.colab
        return True
    except:
        return False

# %% ../nbs/helper_utilities.ipynb 7
class MultiFileChooser:
    def __init__(self):
        self.fc = FileChooser('.')
        self.fc.title = "Use the following file chooser to add each file individually.\n You can remove files by clicking the remove button."
        self.fc.use_dir_icons = True
        self.fc.show_only_dirs = False
        self.selected_files = []
        
        self.fc.register_callback(self.file_selected)
        
        self.output = widgets.Output()
        
    def file_selected(self, chooser):
        if self.fc.selected is not None and self.fc.selected not in self.selected_files:
            self.selected_files.append(self.fc.selected)
            self.update_display()
    
    def update_display(self):
        with self.output:
            clear_output()
            for this_file in self.selected_files:
                remove_button = widgets.Button(description="Remove", tooltip="Remove this file")
                remove_button.on_click(partial(self.remove_file, file=this_file))
                display(widgets.HBox([widgets.Label(value=this_file), remove_button]))
    
    def remove_file(self, button, this_file):
        if this_file in self.selected_files:
            self.selected_files.remove(this_file)
            self.update_display()
    
    def display(self):
        display(self.fc, self.output)
    
    def get_selected_files(self):
        return self.selected_files

# %% ../nbs/helper_utilities.ipynb 12
def setup_drives(upload_set):

    upload_set = upload_set.lower()
    uploaded = None

    # allow them to mount the drive if they chose Google Colab.
    if upload_set == 'google drive':
        if check_is_colab():
            from google.colab import drive
            drive.mount('/content/drive')
        else:
            raise ValueError("It looks like you're not on Google Colab. Google Drive mounting is currently only implemented for Google Colab.")

    # Everything else means that they'll need to use a file chooser (including Google Drive)
    if check_is_colab():
        from google.colab import files
        uploaded = files.upload()
    else:
        # Create file chooser and interact
        mfc = MultiFileChooser()
        mfc.display()
        uploaded = mfc.get_selected_files()
    
    return uploaded