app.py

import os
import shutil
import time
import git
import subprocess

REPO_DIR = '/data/repo/'
USER = os.getenv('HF_USER')
REPO_URL = os.getenv('HF_REPO_URL')
ACCESS_TOKEN = os.getenv('HF_ACCESS_TOKEN')
HF_HOME = os.getenv('HF_HOME')
AUTH_REPO_URL = REPO_URL.replace('https://', f'https://{USER}:{ACCESS_TOKEN}@')


def timer(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f'建置時間：{elapsed_time:.2f} 秒', )
        return result
    return wrapper


def check_envs():
    envs = [USER, REPO_URL, ACCESS_TOKEN]
    missing_vars = [var for var in envs if var is None]

    if missing_vars:
        raise EnvironmentError(f"Missing environment variables: {', '.join(missing_vars)}")


def pull() -> int:
    try:
        repo = git.Repo(REPO_DIR)
        print('嘗試拉取代代碼...')
        repo.remotes.origin.pull()
        print('拉取代代碼成功！')
        return 200
    except git.exc.InvalidGitRepositoryError:
        # 404
        raise ValueError('Invalid Git repository.')

    except git.exc.NoSuchPathError:
        return 404
    
    except git.ecx.GitCommandError as e:
        print(f"拉取代碼時出錯：{e}")
        return 500


def clone(status_code):
    if status_code == 200:
        return
    if status_code == 500:
        if os.path.exists(REPO_DIR):
            shutil.rmtree(REPO_DIR)
            print('狀態碼500，刪除舊repo...')
        
    try:
        print('嘗試CLONE REPO...')
        git.Repo.clone_from(AUTH_REPO_URL, REPO_DIR, branch='main')
        print('嘗試CLONE 成功！')
    except git.exc.GitCommandError as e:
        print(f"克隆代碼時出錯：{e}")
        raise e


def run():
    os.chdir(os.path.join(REPO_DIR))
    print('嘗試運行主程式...')
    subprocess.run(["python", "app.py"], check=True)
    print('運行成功！')

def main():
    check_envs()
    status_code = pull()
    clone(status_code)
    run()


if __name__ == "__main__":
    main()