utils/hub_storage.py

from pathlib import Path
from huggingface_hub import HfApi
from .token_handler import TokenHandler
import time
import tempfile
import os

class HubStorage:
    def __init__(self, repo_id):
        self.repo_id = repo_id
        self.api = HfApi()

    def get_file_content(self, path_in_repo):
        """Get content of a file from the hub"""
        try:
            # Download the file content as bytes
            path = self.api.hf_hub_download(
                repo_id=self.repo_id,
                filename=path_in_repo,
                repo_type="space"
            )
            # Read file content
            with open(path, 'r') as f:
                return f.read()
        except Exception as e:
            print(f"Error reading file {path_in_repo}: {str(e)}")
            return None

    def save_to_hub(self, file_content, path_in_repo, commit_message="Update file"):
        """Save content to hub"""
        try:
            # If file_content is a string (like JSON), write it to a temporary file first
            if isinstance(file_content, str):
                with tempfile.NamedTemporaryFile(mode='w', delete=False) as temp_file:
                    temp_file.write(file_content)
                    temp_path = temp_file.name
            else:
                # If it's already a file path, use it directly
                temp_path = file_content

            # Upload to hub
            self.api.upload_file(
                path_or_fileobj=temp_path,
                path_in_repo=path_in_repo,
                repo_id=self.repo_id,
                repo_type="space",
                commit_message=commit_message
            )
            os.unlink(temp_path)

            return True
        except Exception as e:
            print(f"Error saving to hub: {str(e)}")
            return False