waifuc_gui/file_handler.py

import os
import zipfile
import shutil

class FileHandler:
    def extract_zip(self, uploaded_file):
        extract_dir = "/tmp/input"
        shutil.rmtree(extract_dir, ignore_errors=True)
        os.makedirs(extract_dir)
        with zipfile.ZipFile(uploaded_file, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
        return extract_dir

    def create_zip(self, dataset_name):
        output_dir = f"/tmp/{dataset_name}"
        zip_path = f"/tmp/{dataset_name}.zip"
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, dirs, files in os.walk(dataset_name):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, '.')
                    zipf.write(file_path, arcname)
        return zip_path

    def save_log(self, log_content):
        log_path = "/tmp/waifuc_log.txt"
        with open(log_path, 'w') as f:
            f.write(log_content)
        return log_path