xiaoshiai-hub 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

xiaoshiai_hub/upload.py

@@ -12,7 +12,7 @@ from typing import List, Optional, Union, Dict
 import requests
 
 from xiaoshiai_hub.client import DEFAULT_BASE_URL, HubClient
-from xiaoshiai_hub.envelope_crypto import DataKey, envelope_enc_file
+from xiaoshiai_hub.envelope_crypto import Algorithm, DataKey, envelope_enc_file
 from .exceptions import HubException, AuthenticationError, RepositoryNotFoundError
 
 
@@ -70,61 +70,98 @@ def _calculate_file_sha256(file_path: Path) -> str:
     return sha256_hash.hexdigest()
 
 
-class _TqdmUploadWrapper:
-    """Wrapper to add tqdm progress bar to file upload."""
+class _ProgressFileReader:
+    """
+    文件读取器，跟踪实际网络上传进度。
+
+    通过分块读取文件并在每次 read() 调用后更新进度条，
+    配合 requests 的流式上传，确保进度条反映实际网络传输进度。
+    """
 
-    def __init__(self, file_obj, total_size, desc=None):
-        self.file_obj = file_obj
-        self.total_size = total_size
+    def __init__(self, file_path: Path, chunk_size: int = 8192, desc: Optional[str] = None):
+        self.file_path = file_path
+        self.file_size = file_path.stat().st_size
+        self.chunk_size = chunk_size
+        self.file_obj = open(file_path, 'rb')
+        self.bytes_read = 0
         self.pbar = None
         if tqdm:
             self.pbar = tqdm(
-                total=total_size,
+                total=self.file_size,
                 unit='B',
                 unit_scale=True,
                 unit_divisor=1024,
                 desc=desc,
             )
 
-    def read(self, size=-1):
-        """Read data and update progress bar."""
+    def read(self, size: int = -1) -> bytes:
+        """读取数据并更新进度条。"""
+        if size == -1:
+            size = self.chunk_size
         data = self.file_obj.read(size)
-        if self.pbar is not None and data:
-            self.pbar.update(len(data))
+        if data:
+            self.bytes_read += len(data)
+            if self.pbar is not None:
+                self.pbar.update(len(data))
         return data
 
-    def __enter__(self):
+    def __len__(self) -> int:
+        """返回文件大小，供 requests 设置 Content-Length。"""
+        return self.file_size
+
+    def __iter__(self):
+        """迭代器，用于流式上传。"""
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore[no-untyped-def]
-        del exc_type, exc_val, exc_tb  # Unused but required by context manager protocol
+    def __next__(self) -> bytes:
+        data = self.read(self.chunk_size)
+        if data:
+            return data
+        raise StopIteration
+
+    def close(self):
         if self.pbar is not None:
             self.pbar.close()
         self.file_obj.close()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        del exc_type, exc_val, exc_tb
+        self.close()
+
 
 def _upload_file_with_progress(
     upload_url: str,
     file_path: Path,
     desc: Optional[str] = None,
+    chunk_size: int = 1024 * 1024,  # 1MB chunks
 ) -> None:
-    """Upload file to URL with progress bar."""
+    """
+    上传文件到 URL 并显示真实的网络传输进度。
+
+    使用分块流式上传，进度条反映实际发送到服务器的字节数。
+    """
     file_size = file_path.stat().st_size
 
-    with open(file_path, 'rb') as f:
-        with _TqdmUploadWrapper(f, file_size, desc=desc) as wrapped_file:
-            upload_response = requests.put(
-                upload_url,
-                data=wrapped_file,
-                headers={'Content-Type': 'application/octet-stream'}
-            )
-            upload_response.raise_for_status()
+    with _ProgressFileReader(file_path, chunk_size=chunk_size, desc=desc) as reader:
+        upload_response = requests.put(
+            upload_url,
+            data=reader,
+            headers={
+                'Content-Type': 'application/octet-stream',
+                'Content-Length': str(file_size),
+            }
+        )
+        upload_response.raise_for_status()
 
 
 def _encrypt_file_if_needed(
     file_path: Path,
     encryption_password: Optional[str] = None,
     data_key: Optional[DataKey] = None,
+    algorithm: Optional[Algorithm] = None,
 ) -> tuple[Optional[Path], Optional[Path]]:
     """
     Encrypt file if encryption_password is provided, file is large enough, and file extension is encryptable.
@@ -133,6 +170,7 @@ def _encrypt_file_if_needed(
         file_path: Path to the file to encrypt
         encryption_password: Password for encryption
         data_key: DataKey for encryption (required if encryption_password is provided)
+        algorithm: Encryption algorithm (default: AES-256-CTR)
 
     Returns:
         Tuple of (encrypted_file_path, temp_dir_path)
@@ -154,13 +192,15 @@ def _encrypt_file_if_needed(
     encrypted_file = temp_dir / file_path.name
 
     # Encrypt the file
-    envelope_enc_file(
-        source=file_path,
-        password=encryption_password,
-        data_key=data_key,
-        dest=encrypted_file,
-        chunked=True,
-    )
+    enc_kwargs = {
+        "source": file_path,
+        "data_key": data_key,
+        "dest": encrypted_file,
+        "chunked": True,
+    }
+    if algorithm:
+        enc_kwargs["algorithm"] = algorithm
+    envelope_enc_file(**enc_kwargs)
     print(f"Encrypted file path: {encrypted_file}")
 
     return encrypted_file, temp_dir
@@ -295,6 +335,7 @@ def upload_folder(
     encryption_password: Optional[str] = None,
     ignore_patterns: Optional[List[str]] = None,
     temp_dir: Optional[Union[str, Path]] = None,
+    algorithm: Optional[str] = None,
 ) -> Dict:
     """
     Upload a folder to a repository using HTTP API.
@@ -313,6 +354,7 @@ def upload_folder(
         encryption_password: Password for file encryption (optional)
         ignore_patterns: List of patterns to ignore
         temp_dir: Temporary directory for encrypted files (optional, auto-created if not specified)
+        algorithm: Encryption algorithm ("AES-256-CTR" or "SM4-CTR", default: AES-256-CTR)
 
     Returns:
         Upload response
@@ -361,7 +403,7 @@ def upload_folder(
             # Auto-create temp directory
             temp_dir_path = Path(tempfile.mkdtemp())
         # Generate data key for encryption
-        data_key = client.generate_data_key(encryption_password)
+        data_key = client.generate_data_key(algorithm,encryption_password)
 
     try:
         # Create session and API URL
@@ -409,13 +451,16 @@ def upload_folder(
                     # Preserve directory structure in temp dir
                     encrypted_file = temp_dir_path / rel_file_path
                     encrypted_file.parent.mkdir(parents=True, exist_ok=True)
-                    envelope_enc_file(
-                        source=local_file,
-                        password=encryption_password,
-                        data_key=data_key,
-                        dest=encrypted_file,
-                        chunked=True,
-                    )
+                    # Build encryption kwargs
+                    enc_kwargs = {
+                        "source": local_file,
+                        "data_key": data_key,
+                        "dest": encrypted_file,
+                        "chunked": True,
+                    }
+                    if algorithm:
+                        enc_kwargs["algorithm"] = Algorithm(algorithm)
+                    envelope_enc_file(**enc_kwargs)
                     print(f"Encrypted file path: {encrypted_file}")
                     # 加密后重新计算大小，避免被huggingface下载的时候大小一致性检查不通过
                     file_size = encrypted_file.stat().st_size
@@ -512,6 +557,7 @@ def upload_file(
     password: Optional[str] = None,
     token: Optional[str] = None,
     encryption_password: Optional[str] = None,
+    algorithm: Optional[str] = None,
 ) -> Dict:
     """
     Upload a single file to a repository using HTTP API.
@@ -529,6 +575,7 @@ def upload_file(
         password: Password for authentication
         token: Token for authentication (preferred)
         encryption_password: Password for file encryption (optional)
+        algorithm: Encryption algorithm ("AES" or "SM4", default: AES)
 
     Returns:
         Upload response
@@ -561,10 +608,11 @@ def upload_file(
     # Generate data key if encryption is needed
     data_key: Optional[DataKey] = None
     if encryption_password:
-        data_key = client.generate_data_key(encryption_password)
+        data_key = client.generate_data_key(algorithm,encryption_password)
 
     # Encrypt file if needed
-    encrypted_file, temp_dir = _encrypt_file_if_needed(path_file, encryption_password, data_key)
+    algo = Algorithm(algorithm) if algorithm else None
+    encrypted_file, temp_dir = _encrypt_file_if_needed(path_file, encryption_password, data_key, algo)
     actual_file = encrypted_file if encrypted_file else path_file
     file_was_encrypted = encrypted_file is not None
 

{xiaoshiai_hub-1.1.1.dist-info → xiaoshiai_hub-1.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xiaoshiai-hub
-Version: 1.1.1
+Version: 1.1.3
 Summary: Python SDK for XiaoShi AI Hub - Upload, download, and manage AI models and datasets with xpai-enc encryption support
 Home-page: https://github.com/poxiaoyun/moha-sdk
 Author: XiaoShi AI
@@ -244,25 +244,65 @@ client = HubClient(
     password="your-password",
 )
 
+# 创建仓库
+repo = client.create_repository(
+    organization="demo",
+    repo_type="models",
+    repo_name="my-model",
+    description="我的模型",
+    visibility="internal",
+    metadata={
+        "license": ["apache-2.0"],
+        "frameworks": ["transformers"],
+    },
+)
+print(f"仓库已创建: {repo.name}")
+
 # 获取仓库信息
 repo_info = client.get_repository_info("demo", "models", "my-model")
 print(f"仓库名称: {repo_info.name}")
 print(f"组织: {repo_info.organization}")
+print(f"所有者: {repo_info.owner}")
+print(f"可见性: {repo_info.visibility}")
+
+# 更新仓库
+client.update_repository(
+    organization="demo",
+    repo_type="models",
+    repo_name="my-model",
+    description="更新后的描述",
+)
 
 # 列出分支
-branches = client.list_branches("demo", "models", "my-model")
-for branch in branches:
-    print(f"分支: {branch.name} (commit: {branch.commit_sha})")
+refs = client.get_repository_refs("demo", "models", "my-model")
+for ref in refs:
+    print(f"分支: {ref.name} (commit: {ref.hash[:8]})")
+
+# 创建分支（幂等操作，已存在则直接返回）
+client.create_branch("demo", "models", "my-model", "dev", "main")
+
+# 删除分支（幂等操作，不存在则直接返回）
+client.delete_branch("demo", "models", "my-model", "dev")
 
 # 浏览仓库内容
 content = client.get_repository_content("demo", "models", "my-model", "main")
 for entry in content.entries:
     print(f"{entry.type}: {entry.name}")
+
+# 删除仓库
+client.delete_repository("demo", "models", "my-model")
 ```
 
 ## 🔐 加密功能
 
-SDK 提供了智能加密功能，使用 AES-256-CTR 算法对大型模型文件进行加密。
+SDK 提供了智能加密功能，支持 **AES** 和 **SM4** 两种加密算法对大型模型文件进行加密。
+
+### 支持的加密算法
+
+| 算法 | 说明 |
+|------|------|
+| `AES` | AES-256-CTR 模式，国际通用标准（默认） |
+| `SM4` | SM4-CTR 模式，国密标准 |
 
 ### 自动加密规则
 
@@ -278,11 +318,21 @@ SDK 提供了智能加密功能，使用 AES-256-CTR 算法对大型模型文件
 ```python
 from xiaoshiai_hub import upload_folder
 
-# 上传文件夹，自动加密大型模型文件
+# 上传文件夹，使用 AES 加密（默认）
 result = upload_folder(
     folder_path="./llama-7b",
     repo_id="demo/llama-7b",
-    encryption_password="my-secure-password-123",  # 设置加密密码
+    encryption_password="my-secure-password-123",
+    username="your-username",
+    password="your-password",
+)
+
+# 使用 SM4 国密算法加密
+result = upload_folder(
+    folder_path="./llama-7b",
+    repo_id="demo/llama-7b",
+    encryption_password="my-secure-password-123",
+    algorithm="SM4",  # 使用 SM4 加密
     username="your-username",
     password="your-password",
 )
@@ -326,7 +376,7 @@ export MOHA_ENCRYPTION_PASSWORD="your-encryption-password"
 
 ## 🖥️ 命令行工具 (CLI)
 
-SDK 提供了 `moha` 命令行工具，支持常见的上传下载操作。
+SDK 提供了 `moha` 命令行工具，支持登录认证、仓库管理、分支管理、上传下载等操作。
 
 ### 基本用法
 
@@ -334,6 +384,75 @@ SDK 提供了 `moha` 命令行工具，支持常见的上传下载操作。
 moha --help
 ```
 
+### 登录认证
+
+```bash
+# 登录（交互式输入用户名和密码）
+moha login
+
+# 直接指定用户名和密码
+moha login --username your-username --password your-password
+
+# 查看当前登录状态
+moha whoami
+
+# 退出登录
+moha logout
+```
+
+登录后，Token 会保存到 `~/.moha/token.json`，后续命令无需重复输入认证信息。
+
+### 仓库管理
+
+```bash
+# 创建仓库
+moha repo-create org/my-model \
+    --description "我的模型" \
+    --visibility internal \
+    --license apache-2.0 \
+    --tasks text-generation \
+    --frameworks transformers
+
+# 创建数据集仓库
+moha repo-create org/my-dataset \
+    --repo-type datasets \
+    --description "我的数据集" \
+    --visibility private
+
+# 查看仓库信息
+moha repo-info org/my-model
+
+# 更新仓库信息
+moha repo-update org/my-model \
+    --description "更新后的描述" \
+    --tags production
+
+# 删除仓库（需要确认）
+moha repo-delete org/my-model
+
+# 跳过确认直接删除
+moha repo-delete org/my-model -y
+```
+
+### 分支管理
+
+```bash
+# 列出仓库的所有分支
+moha branch-list org/my-model
+
+# 创建分支（基于 main 分支）
+moha branch-create org/my-model dev
+
+# 创建分支（基于指定分支）
+moha branch-create org/my-model feature --from dev
+
+# 删除分支
+moha branch-delete org/my-model dev
+
+# 跳过确认直接删除
+moha branch-delete org/my-model dev -y
+```
+
 ### 上传文件夹
 
 ```bash
@@ -353,10 +472,18 @@ moha upload ./my_model org/my-model \
     --username your-username \
     --password your-password
 
-# 启用加密
+# 启用加密（默认使用 AES）
+moha upload ./my_model org/my-model \
+    --encrypt \
+    --encryption-password "your-secret" \
+    --username your-username \
+    --password your-password
+
+# 使用 SM4 国密算法加密
 moha upload ./my_model org/my-model \
     --encrypt \
     --encryption-password "your-secret" \
+    --algorithm SM4 \
     --username your-username \
     --password your-password
 ```
@@ -419,26 +546,75 @@ moha download-file org/my-model model.safetensors \
     --password your-password
 ```
 
+### CLI 命令列表
+
+| 命令 | 说明 |
+|------|------|
+| `moha login` | 登录并保存 Token |
+| `moha logout` | 退出登录并删除 Token |
+| `moha whoami` | 查看当前登录状态 |
+| `moha repo-create` | 创建仓库 |
+| `moha repo-update` | 更新仓库 |
+| `moha repo-delete` | 删除仓库 |
+| `moha repo-info` | 查看仓库信息 |
+| `moha branch-create` | 创建分支 |
+| `moha branch-delete` | 删除分支 |
+| `moha branch-list` | 列出仓库的所有分支 |
+| `moha upload` | 上传文件夹到仓库 |
+| `moha upload-file` | 上传单个文件到仓库 |
+| `moha download` | 下载整个仓库 |
+| `moha download-file` | 从仓库下载单个文件 |
+
 ### CLI 参数说明
 
+#### 通用参数
+
 | 参数 | 说明 | 适用命令 |
 |------|------|----------|
-| `--repo-type, -t` | 仓库类型：`models` 或 `datasets`（默认：models） | 所有 |
-| `--revision, -r` | 分支/标签/提交（默认：main） | 所有 |
+| `--repo-type, -t` | 仓库类型：`models` 或 `datasets`（默认：models） | 大部分命令 |
 | `--base-url` | API 基础 URL（默认：环境变量 MOHA_ENDPOINT） | 所有 |
 | `--token` | 认证令牌 | 所有 |
 | `--username` | 用户名 | 所有 |
 | `--password` | 密码 | 所有 |
+
+#### 上传/下载参数
+
+| 参数 | 说明 | 适用命令 |
+|------|------|----------|
+| `--revision, -r` | 分支/标签/提交（默认：main） | upload, download |
 | `--message, -m` | 提交消息 | upload, upload-file |
 | `--ignore, -i` | 忽略模式（可多次使用） | upload, download |
 | `--include` | 包含模式（可多次使用） | download |
 | `--encrypt, -e` | 启用加密 | upload, upload-file |
 | `--encryption-password` | 加密密码 | upload, upload-file |
+| `--algorithm, -a` | 加密算法：`AES` 或 `SM4`（默认：AES） | upload, upload-file |
 | `--path-in-repo, -p` | 仓库中的文件路径 | upload-file |
 | `--temp-dir` | 加密临时目录 | upload |
 | `--local-dir, -o` | 本地保存目录 | download, download-file |
 | `--quiet, -q` | 禁用进度条 | download, download-file |
 
+#### 仓库管理参数
+
+| 参数 | 说明 | 适用命令 |
+|------|------|----------|
+| `--description, -d` | 仓库描述 | repo-create, repo-update |
+| `--visibility, -v` | 可见性：`public`、`internal`、`private` | repo-create, repo-update |
+| `--license` | 许可证（可多次使用） | repo-create, repo-update |
+| `--tasks` | 任务类型（可多次使用） | repo-create, repo-update |
+| `--languages` | 语言（可多次使用） | repo-create, repo-update |
+| `--tags` | 标签（可多次使用） | repo-create, repo-update |
+| `--frameworks` | 框架（可多次使用） | repo-create, repo-update |
+| `--base-model` | 基础模型（可多次使用） | repo-create, repo-update |
+| `--relationship` | 与基础模型的关系 | repo-create, repo-update |
+| `--yes, -y` | 跳过确认提示 | repo-delete, branch-delete |
+
+#### 分支管理参数
+
+| 参数 | 说明 | 适用命令 |
+|------|------|----------|
+| `--from, -f` | 基于哪个分支创建（默认：main） | branch-create |
+| `--yes, -y` | 跳过确认提示 | branch-delete |
+
 ### 使用环境变量
 
 可以通过环境变量设置认证信息，避免每次输入：

xiaoshiai_hub-1.1.3.dist-info/RECORD

@@ -0,0 +1,15 @@
+xiaoshiai_hub/__init__.py,sha256=2hnhTGad1GPTzRstfwZhAF2204Pl2VF8YHD8HQkAYGY,1469
+xiaoshiai_hub/auth.py,sha256=Pv0P6f76flOefYmALyehuMu2Klyc-u_gtXlJYJR4eMY,4105
+xiaoshiai_hub/cli.py,sha256=E2xWvybOz6Tr8bSk8pp-QZVHoCyQfnmCYg7442aq4-Y,29790
+xiaoshiai_hub/client.py,sha256=3ASE_KuV-S0CmDvutvmMpY3GU2Hdrui2WAfwHs3iEkw,18775
+xiaoshiai_hub/download.py,sha256=9Uido7cJjGVd6ERDKu_xNMPliarPdZVVb8hkLgYfFcU,13613
+xiaoshiai_hub/envelope_crypto.py,sha256=zjrt5fc3ya86Y4N7_4OI5_NnmFpQ5HiHgoP31ioJRmw,8235
+xiaoshiai_hub/exceptions.py,sha256=24QzgHWq_4bes07UkC3vGi2oT8SMH6Xu4FNlKt52QHo,672
+xiaoshiai_hub/types.py,sha256=-h8EctZo_bRJALufjSTaxIWf2bqHz2ZDCy3SCPfucHs,2004
+xiaoshiai_hub/upload.py,sha256=GKlnsxsIVuLhCWpymupya74TNgTP0s8T-gKTOR-RMgw,22371
+xiaoshiai_hub-1.1.3.dist-info/licenses/LICENSE,sha256=tS28u6VpvqNisRWGeufp-XYQc6p194vOGARl3OIjidA,9110
+xiaoshiai_hub-1.1.3.dist-info/METADATA,sha256=RCvKnYGmiU5upd5HXYkK0RzitGZT-Wsjvw_6CDoBa8M,21630
+xiaoshiai_hub-1.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xiaoshiai_hub-1.1.3.dist-info/entry_points.txt,sha256=fVh_IA1mbRWl7LEd4-RhENMaspSBa4Hxbsg8_HDWc6Y,48
+xiaoshiai_hub-1.1.3.dist-info/top_level.txt,sha256=9AQDFb5Xn7RLQPdbk1aA0QpntbKhlhlT6Z_g-zUBtlM,14
+xiaoshiai_hub-1.1.3.dist-info/RECORD,,

xiaoshiai_hub-1.1.1.dist-info/RECORD

@@ -1,15 +0,0 @@
-xiaoshiai_hub/__init__.py,sha256=ZDZEGKBbv8zEDO9GkKPy0uiTva9gPPEMIU9_A_MzCZA,1339
-xiaoshiai_hub/auth.py,sha256=Pv0P6f76flOefYmALyehuMu2Klyc-u_gtXlJYJR4eMY,4105
-xiaoshiai_hub/cli.py,sha256=IKuHYGyk-Kg57BUquWyZFMK6rPfLMmwEOWGZq8E7LJs,14885
-xiaoshiai_hub/client.py,sha256=OQyyKD5kyDaPlI81-js1Uj9gYjVPkS60UicZAtokDVw,11914
-xiaoshiai_hub/download.py,sha256=9Uido7cJjGVd6ERDKu_xNMPliarPdZVVb8hkLgYfFcU,13613
-xiaoshiai_hub/envelope_crypto.py,sha256=PWZrVCkgQckFQO6putRDlZXkQoEYOkq5O0PPltdBDE4,7161
-xiaoshiai_hub/exceptions.py,sha256=24QzgHWq_4bes07UkC3vGi2oT8SMH6Xu4FNlKt52QHo,672
-xiaoshiai_hub/types.py,sha256=ZC5elxqles8_ODl-fCSssOzm9q8_KjA9mGpiGgObgls,1915
-xiaoshiai_hub/upload.py,sha256=cYKCvaxfO5wOmjCmWp_rgbTOdEr73MN2EoBUCwyqqvE,20713
-xiaoshiai_hub-1.1.1.dist-info/licenses/LICENSE,sha256=tS28u6VpvqNisRWGeufp-XYQc6p194vOGARl3OIjidA,9110
-xiaoshiai_hub-1.1.1.dist-info/METADATA,sha256=0zbLRQQGJbyp5nAcmloTSvoRe9LEXuB8pKGBWNgUudo,16689
-xiaoshiai_hub-1.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xiaoshiai_hub-1.1.1.dist-info/entry_points.txt,sha256=fVh_IA1mbRWl7LEd4-RhENMaspSBa4Hxbsg8_HDWc6Y,48
-xiaoshiai_hub-1.1.1.dist-info/top_level.txt,sha256=9AQDFb5Xn7RLQPdbk1aA0QpntbKhlhlT6Z_g-zUBtlM,14
-xiaoshiai_hub-1.1.1.dist-info/RECORD,,