Files
TaskPool/custom/sync.py
T

332 lines
11 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import os
import subprocess
import sys
import urllib.error
import urllib.request
from typing import TYPE_CHECKING, Protocol, cast
# 仅在类型检查时导入,避免运行时依赖
if TYPE_CHECKING:
from http.client import HTTPResponse
class SyncArgs(Protocol):
"""用于类型检查的参数协议,映射 argparse 的解析结果"""
source_type: str
source_url: str
target_path: str
branch: str | None
path: str | None
single_file: bool
proxy: str | None
proxy_url: str | None
auth_token: str | None
http_proxy: str | None
def run(
cmd: list[str],
env: dict[str, str] | None = None,
cwd: str | None = None,
capture_output: bool = False
) -> str | None:
"""
执行系统命令。
Args:
cmd: 命令列表
env: 环境变量字典
cwd: 当前工作目录
capture_output: 是否捕获输出。
"""
if not capture_output:
print(">>", " ".join(cmd))
if capture_output:
# 捕获模式:不打印到屏幕,返回 stdout
result = subprocess.run(
cmd,
cwd=cwd,
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
timeout=30,
encoding="utf-8",
errors="ignore"
)
else:
# 直通模式:直接打印到屏幕
result = subprocess.run(
cmd,
cwd=cwd,
env=env,
stdout=sys.stdout,
stderr=sys.stderr,
)
if result.returncode != 0:
if capture_output:
return None
sys.exit(result.returncode)
if capture_output:
return str(result.stdout).strip()
return str(result)
def get_remote_default_branch(repo_url: str, env: dict[str, str]) -> str:
"""检测远程仓库的默认分支名称(通常是 main 或 master)。"""
print(f"正在检测远程仓库默认分支: {repo_url}")
cmd = ["git", "ls-remote", "--symref", repo_url, "HEAD"]
output = run(cmd, env=env, capture_output=True)
if isinstance(output, str):
for line in output.splitlines():
parts = line.split()
if len(parts) >= 2 and parts[0] == "ref:" and "refs/heads/" in parts[1]:
branch = parts[1].removeprefix("refs/heads/")
print(f"检测到默认分支: {branch}")
return branch
print("无法检测到默认分支,回退使用 'main'")
return "main"
def build_proxy_url(url: str, proxy_type: str | None, proxy_url: str | None) -> str:
"""根据配置构建带有代理前缀的 URL。"""
if not proxy_type or proxy_type == "none":
return url
proxy_base = ""
if proxy_type == "ghproxy":
proxy_base = "https://gh-proxy.com/"
elif proxy_type == "mirror":
proxy_base = "https://mirror.ghproxy.com/"
elif proxy_type == "custom" and proxy_url:
proxy_base = proxy_url.rstrip("/") + "/"
if proxy_base and url.startswith("http"):
return proxy_base + url
return url
def _download_file(url: str, dest: str, auth_token: str | None) -> None:
"""
内部通用下载函数,处理请求构建、Token 认证和文件写入。
"""
print(f"下载地址: {url}")
print(f"目标路径: {dest}")
parent_dir = os.path.dirname(dest)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
req = urllib.request.Request(url)
if auth_token:
req.add_header("Authorization", f"token {auth_token}")
req.add_header("User-Agent", "Mozilla/5.0 (compatible; sync.py)")
try:
with cast("HTTPResponse", urllib.request.urlopen(req, timeout=300)) as response:
content: bytes = response.read()
with open(dest, "wb") as f:
_ = f.write(content)
print(f"文件大小: {len(content)} 字节")
print("下载完成")
except urllib.error.HTTPError as e:
print(f"下载失败, HTTP 状态码: {e.code}")
sys.exit(1)
except urllib.error.URLError as e:
print(f"下载失败: {e.reason}")
sys.exit(1)
def sync_git_file(args: SyncArgs, repo_url: str, env: dict[str, str]) -> None:
"""
从 Git 仓库同步单个文件(通过构造 Raw URL 下载)。
"""
source_url = args.source_url
file_path = args.path or ""
dest = args.target_path
# 如果目标是目录,自动拼接文件名
if os.path.isdir(dest) or dest.endswith(os.sep) or (os.altsep and dest.endswith(os.altsep)):
filename = os.path.basename(file_path)
dest = os.path.join(dest, filename)
print(f"检测到目标路径为目录 '{args.target_path}',自动修正为: '{dest}'")
branch = args.branch or get_remote_default_branch(repo_url, env)
# 构建 raw 文件 URL
# GitHub: https://github.com/user/repo -> https://raw.githubusercontent.com/user/repo/branch/path
# GitLab: https://gitlab.com/user/repo -> https://gitlab.com/user/repo/-/raw/branch/path
# Gitee: https://gitee.com/user/repo -> https://gitee.com/user/repo/raw/branch/path
clean_url = args.source_url.rstrip(".git")
raw_url = ""
if "github.com" in source_url:
base = args.source_url.replace("github.com", "raw.githubusercontent.com").rstrip(".git")
raw_url = f"{base}/{branch}/{file_path}"
elif "gitlab.com" in source_url:
raw_url = f"{clean_url}/-/raw/{branch}/{file_path}"
elif "gitee.com" in source_url:
raw_url = f"{clean_url}/raw/{branch}/{file_path}"
else:
# 通用策略:尝试 GitHub 风格
raw_url = f"{clean_url}/raw/{branch}/{file_path}"
raw_url = build_proxy_url(raw_url, args.proxy, args.proxy_url)
# 调用通用下载函数
_download_file(raw_url, dest, args.auth_token)
def is_raw_file_url(url: str) -> bool:
"""判断 URL 是否已经是 Raw 文件链接。"""
raw_patterns = [
"raw.githubusercontent.com",
"/raw/",
"/-/raw/",
"/blob/",
]
return any(pattern in url for pattern in raw_patterns)
def get_repo_name(url: str) -> str:
"""从 Git URL 中提取仓库名称。"""
url_stripped = url.rstrip("/").rstrip(".git")
return os.path.basename(url_stripped)
def sync_git(args: SyncArgs) -> None:
"""处理 Git 类型的同步逻辑(Clone, Pull 或 Sparse Checkout)。"""
env = os.environ.copy()
if is_raw_file_url(args.source_url):
print("检测到 raw 文件 URL,自动切换到 URL 下载模式")
sync_url(args)
return
if args.http_proxy:
env["http_proxy"] = args.http_proxy
env["https_proxy"] = args.http_proxy
repo_url = build_proxy_url(args.source_url, args.proxy, args.proxy_url)
if args.auth_token and repo_url.startswith("https://"):
repo_url = repo_url.replace("https://", f"https://{args.auth_token}@")
dest = args.target_path
# 单文件模式:使用 Raw URL 下载
if args.path and args.single_file:
sync_git_file(args, repo_url, env)
return
# 自动追加仓库名逻辑
git_dir = os.path.join(dest, ".git")
if os.path.isdir(dest) and not os.path.exists(git_dir):
repo_name = get_repo_name(args.source_url)
dest = os.path.join(dest, repo_name)
print(f"目标路径自动追加仓库名: {dest}")
git_dir = os.path.join(dest, ".git")
if os.path.exists(git_dir):
print("检测到已存在仓库,执行 git pull")
if args.branch:
_ = run(["git", "checkout", args.branch], cwd=dest, env=env)
_ = run(["git", "pull"], cwd=dest, env=env)
else:
print("执行 git clone")
parent_dir = os.path.dirname(dest)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
if os.path.exists(dest) and os.listdir(dest):
print(f"错误: 目标目录 '{dest}' 已存在且不为空,无法执行 git clone")
print("提示: 请清空目标目录或指定一个新目录")
sys.exit(1)
clone_cmd = ["git", "clone", "--depth", "1"]
if args.branch:
clone_cmd.extend(["-b", args.branch])
# 稀疏检出 (Sparse Checkout)
if args.path:
clone_cmd.extend(["--filter=blob:none", "--no-checkout", repo_url, dest])
_ = run(clone_cmd, env=env)
_ = run(["git", "sparse-checkout", "init", "--cone"], cwd=dest, env=env)
_ = run(["git", "sparse-checkout", "set", args.path], cwd=dest, env=env)
_ = run(["git", "checkout"], cwd=dest, env=env)
else:
# 普通 Clone
clone_cmd.extend([repo_url, dest])
_ = run(clone_cmd, env=env)
print("同步完成")
def sync_url(args: SyncArgs) -> None:
"""处理普通 URL 文件下载逻辑。"""
download_url = build_proxy_url(args.source_url, args.proxy, args.proxy_url)
print(f"下载地址: {download_url}")
dest = args.target_path
if os.path.isdir(dest) or dest.endswith("/"):
url_path = args.source_url.split("?")[0]
filename = os.path.basename(url_path) or "downloaded_file"
dest = os.path.join(dest, filename)
print(f"目标文件: {dest}")
# 调用通用下载函数
_download_file(download_url, dest, args.auth_token)
def main() -> None:
parser = argparse.ArgumentParser(description="仓库/文件同步工具")
_ = parser.add_argument("--source-type", choices=["git", "url"], default="git",
help="源类型: git(Git仓库) 或 url(URL下载)")
_ = parser.add_argument("--source-url", required=True,
help="源地址(Git仓库URL或文件URL")
_ = parser.add_argument("--target-path", required=True,
help="目标路径")
_ = parser.add_argument("--branch",
help="Git 分支名(仅 git 类型有效)")
_ = parser.add_argument("--path",
help="仅拉取指定文件或目录(仅 git 类型有效)")
_ = parser.add_argument("--single-file", action="store_true",
help="单文件模式,直接下载指定文件而非 sparse-checkout(需配合 --path 使用)")
_ = parser.add_argument("--proxy", choices=["none", "ghproxy", "mirror", "custom"], default="none",
help="代理类型")
_ = parser.add_argument("--proxy-url",
help="自定义代理地址(仅 proxy=custom 时有效)")
_ = parser.add_argument("--auth-token",
help="认证 Token(用于私有仓库)")
_ = parser.add_argument("--http-proxy",
help="HTTP 代理(如 http://127.0.0.1:7890")
# 使用 cast 将 Namespace 转换为 SyncArgs 协议,满足静态类型检查
args = cast(SyncArgs, cast(object, parser.parse_args()))
print("参数:", " ".join(sys.argv[1:]))
if args.source_type == "git":
sync_git(args)
else:
sync_url(args)
if __name__ == "__main__":
main()