toolbox/hf_netdisk/netdisk/repo_client.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse

from toolbox.hf_netdisk.hf_netdisk_client import HfNetdiskClient


class RepoClient(HfNetdiskClient):
    def __init__(self, hf_token: str):
        super(RepoClient, self).__init__(hf_token=hf_token)

    def create_repo(self, repo_id: str) -> str:
        self.hf_api.create_repo(
            repo_id=repo_id,
            repo_type="dataset",
            private=False,
            exist_ok=True
        )
        return repo_id

    def delete_repo(self, repo_id: str) -> str:
        self.hf_api.delete_repo(
            repo_id=repo_id,
            repo_type="dataset",
            missing_ok=True
        )
        return repo_id

    def check_repo(self, repo_id: str) -> bool:
        exists: bool = self.hf_api.repo_exists(
            repo_id=repo_id,
            repo_type="dataset",
        )
        return exists

    def make_sure_repo(self, repo_id: str) -> bool:
        exists: bool = self.hf_api.repo_exists(
            repo_id=repo_id,
            repo_type="dataset",
        )
        if exists is not True:
            self.hf_api.create_repo(
                repo_id=repo_id,
                repo_type="dataset",
                private=False,
                exist_ok=True
            )
        exists: bool = self.hf_api.repo_exists(
            repo_id=repo_id,
            repo_type="dataset",
        )
        return exists


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--key_of_credentials",
        default="hf_miyuki_write_token",
        type=str,
    )
    args = parser.parse_args()
    return args


def main():
    args = get_args()

    import log
    from project_settings import environment, project_path, log_directory, time_zone_info

    log.setup_size_rotating(log_directory=log_directory, tz_info=time_zone_info)

    hf_token = environment.get(key=args.key_of_credentials)

    client = RepoClient(
        hf_token=hf_token,
    )

    flag = client.check_repo("tianxing1994/hf_netdisk")
    print(f"flag: {flag}")
    repo_id = client.create_repo("tianxing1994/hf_netdisk")
    flag = client.check_repo("tianxing1994/hf_netdisk")
    print(f"flag: {flag}")
    repo_id = client.delete_repo("tianxing1994/hf_netdisk")
    flag = client.check_repo("tianxing1994/hf_netdisk")
    print(f"flag: {flag}")

    return


if __name__ == "__main__":
    main()