TrixyZ/trainer/components/callbacks/data_source_callbacks.py

import base64
import os
from dash.dependencies import Input, Output, State, ALL
from dash import html, callback_context
from utils.file_utils import list_files_in_directory, save_file, delete_file, get_file_metadata
from components.data_source import generate_data_source_table, DATA_SOURCE_DIR

def register_data_source_callbacks(app):

    @app.callback(
        Output('data-source-upload-status', 'children'),
        [
            Input('upload-data', 'contents'),
            Input({'type': 'delete-data-source-button', 'index': ALL}, 'n_clicks')
        ],
        [
            State('upload-data', 'filename'),
            State('upload-data', 'last_modified'),
            State({'type': 'delete-data-source-button', 'index': ALL}, 'id')
        ]
    )
    def manage_data_source_files(contents, filenames, last_modified, delete_ids):
        ctx = callback_context

        if not ctx.triggered:
            return html.Div()

        triggered_id = ctx.triggered[0]['prop_id'].split('.')[0]

        if 'upload-data' in triggered_id and contents:
            for content, filename in zip(contents, filenames):
                data = content.split(",")[1]
                save_file(DATA_SOURCE_DIR, filename, base64.b64decode(data))
            files = list_files_in_directory(DATA_SOURCE_DIR)
            return generate_data_source_table(files)

        elif 'delete-data-source-button' in triggered_id:
            for item in delete_ids:
                if item['id'] == triggered_id:
                    file_path = os.path.join(DATA_SOURCE_DIR, item['index'])
                    delete_file(file_path)
            files = list_files_in_directory(DATA_SOURCE_DIR)
            return generate_data_source_table(files)

        return html.Div()