utils/hub_storage.py

from pathlib import Path
from huggingface_hub import HfApi
from .token_handler import TokenHandler

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

    def get_file_content(self, file_path):
        """
        Get content of a file from the repository
        """
        try:
            content = self.api.hf_hub_download(
                repo_id=self.repo_id,
                repo_type="space",
                filename=file_path,
                text=True
            )
            return content
        except Exception as e:
            print(f"Error reading file {file_path}: {str(e)}")
            return None

    def save_to_hub(self, file_content, path_in_repo, commit_message):
        """
        Save a file to the hub
        """
        try:
            self.api.upload_file(
                path_or_fileobj=file_content,
                path_in_repo=path_in_repo,
                repo_id=self.repo_id,
                repo_type="space",
                commit_message=commit_message
            )
            return True
        except Exception as e:
            print(f"Error saving file to hub: {str(e)}")
            return False