355 lines
9.3 KiB
Python
Executable File
355 lines
9.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""开发分支工作区初始化脚本
|
||
|
||
用于创建基于远端分支或新建的开发分支工作区。
|
||
"""
|
||
|
||
import argparse
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
|
||
def run_git(args: list[str]) -> subprocess.CompletedProcess:
|
||
"""执行 git 命令
|
||
|
||
Args:
|
||
args: git 命令参数列表
|
||
|
||
Returns:
|
||
subprocess.CompletedProcess 对象
|
||
"""
|
||
return subprocess.run(
|
||
["git"] + args,
|
||
capture_output=True,
|
||
text=True,
|
||
check=False
|
||
)
|
||
|
||
|
||
def get_root_dir() -> Path:
|
||
"""获取 git 仓库根目录
|
||
|
||
Raises:
|
||
RuntimeError: 不在 git 仓库中
|
||
"""
|
||
result = run_git(["rev-parse", "--show-toplevel"])
|
||
if result.returncode != 0:
|
||
raise RuntimeError("不在 git 仓库中")
|
||
return Path(result.stdout.strip())
|
||
|
||
|
||
def fetch_remote() -> bool:
|
||
"""从远端获取最新分支信息
|
||
|
||
Returns:
|
||
是否成功获取
|
||
"""
|
||
result = run_git(["fetch", "--quiet"])
|
||
if result.returncode != 0:
|
||
print("警告: 无法获取远端信息，继续使用本地数据")
|
||
return False
|
||
return True
|
||
|
||
|
||
def get_all_remote_branches() -> list[str]:
|
||
"""获取所有远端分支列表
|
||
|
||
Returns:
|
||
远端分支列表
|
||
"""
|
||
result = run_git(["branch", "-r"])
|
||
if result.returncode != 0:
|
||
return []
|
||
|
||
return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
|
||
|
||
|
||
def get_remote_branches(branch_name: str) -> list[str]:
|
||
"""获取远端匹配的分支列表
|
||
|
||
Args:
|
||
branch_name: 分支名称
|
||
|
||
Returns:
|
||
匹配的远端分支列表
|
||
"""
|
||
result = run_git(["branch", "-r"])
|
||
if result.returncode != 0:
|
||
return []
|
||
|
||
return [
|
||
line.strip()
|
||
for line in result.stdout.strip().split("\n")
|
||
if line.strip() and line.strip().endswith(f"/{branch_name}")
|
||
]
|
||
|
||
|
||
def branch_exists_local(branch_name: str) -> bool:
|
||
"""检查本地分支是否存在
|
||
|
||
Args:
|
||
branch_name: 分支名称
|
||
|
||
Returns:
|
||
是否存在
|
||
"""
|
||
result = run_git(["show-ref", "--verify", "--quiet", f"refs/heads/{branch_name}"])
|
||
return result.returncode == 0
|
||
|
||
|
||
def worktree_exists(worktree_path: Path) -> bool:
|
||
"""检查工作区是否存在
|
||
|
||
Args:
|
||
worktree_path: 工作区路径
|
||
|
||
Returns:
|
||
是否存在
|
||
"""
|
||
result = run_git(["worktree", "list"])
|
||
if result.returncode != 0:
|
||
return False
|
||
return str(worktree_path) in result.stdout
|
||
|
||
|
||
def extract_branch_name(remote_branch: str) -> str:
|
||
"""从远端分支名提取分支名
|
||
|
||
Args:
|
||
remote_branch: 远端分支名（如 origin/feature-123）
|
||
|
||
Returns:
|
||
分支名（如 feature-123）
|
||
"""
|
||
return remote_branch.split("/", 1)[1] if "/" in remote_branch else remote_branch
|
||
|
||
|
||
def validate_branch_creation(branch_name: str, worktree_path: Path) -> None:
|
||
"""验证是否可以创建分支和工作区
|
||
|
||
Args:
|
||
branch_name: 分支名称
|
||
worktree_path: 工作区路径
|
||
|
||
Raises:
|
||
RuntimeError: 验证失败
|
||
"""
|
||
if worktree_path.exists():
|
||
raise RuntimeError(f"工作区已存在于 {worktree_path}")
|
||
|
||
if worktree_exists(worktree_path):
|
||
raise RuntimeError(f"工作区 '{branch_name}' 已存在")
|
||
|
||
if branch_exists_local(branch_name):
|
||
raise RuntimeError(f"本地分支 '{branch_name}' 已存在")
|
||
|
||
|
||
def select_from_list(items: list[str], prompt: str, allow_create: bool = False) -> str | None:
|
||
"""让用户从列表中选择
|
||
|
||
Args:
|
||
items: 选项列表
|
||
prompt: 提示信息
|
||
allow_create: 是否允许创建新分支
|
||
|
||
Returns:
|
||
用户选择的项，或 None 表示创建新分支
|
||
"""
|
||
if not items:
|
||
return None
|
||
|
||
print(prompt)
|
||
for i, item in enumerate(items, 1):
|
||
print(f" {i}\t{item}")
|
||
|
||
if allow_create:
|
||
print(f" {len(items) + 1}\t创建新分支")
|
||
|
||
print()
|
||
|
||
max_choice = len(items) + 1 if allow_create else len(items)
|
||
|
||
while True:
|
||
try:
|
||
selection = int(input(f"请选择 (1-{max_choice}): "))
|
||
if 1 <= selection <= len(items):
|
||
selected = items[selection - 1]
|
||
print(f"已选择: {selected}")
|
||
return selected
|
||
if allow_create and selection == len(items) + 1:
|
||
return None
|
||
print(f"错误: 请输入 1-{max_choice} 之间的数字")
|
||
except ValueError:
|
||
print("错误: 请输入有效的数字")
|
||
except KeyboardInterrupt:
|
||
print("\n操作已取消")
|
||
sys.exit(1)
|
||
|
||
|
||
def input_branch_name() -> str:
|
||
"""让用户输入分支名称
|
||
|
||
Returns:
|
||
分支名称
|
||
"""
|
||
while True:
|
||
try:
|
||
name = input("请输入新分支名称: ").strip()
|
||
if name:
|
||
return name
|
||
print("错误: 分支名称不能为空")
|
||
except KeyboardInterrupt:
|
||
print("\n操作已取消")
|
||
sys.exit(1)
|
||
|
||
|
||
def create_worktree(branch_name: str, worktree_path: Path, base_branch: str | None = None) -> None:
|
||
"""创建工作区
|
||
|
||
Args:
|
||
branch_name: 新分支名称
|
||
worktree_path: 工作区路径
|
||
base_branch: 基础分支（可选）
|
||
|
||
Raises:
|
||
RuntimeError: 创建失败
|
||
"""
|
||
args = ["worktree", "add", "-b", branch_name, str(worktree_path)]
|
||
if base_branch:
|
||
args.append(base_branch)
|
||
|
||
result = run_git(args)
|
||
if result.returncode != 0:
|
||
raise RuntimeError(f"创建工作区失败: {result.stderr.strip()}")
|
||
|
||
|
||
def create_worktree_with_validation(
|
||
branch_name: str,
|
||
worktree_path: Path,
|
||
base_branch: str | None = None
|
||
) -> None:
|
||
"""验证并创建工作区
|
||
|
||
Args:
|
||
branch_name: 新分支名称
|
||
worktree_path: 工作区路径
|
||
base_branch: 基础分支（可选）
|
||
"""
|
||
validate_branch_creation(branch_name, worktree_path)
|
||
create_worktree(branch_name, worktree_path, base_branch)
|
||
print(f"工作区已创建于 {worktree_path}")
|
||
|
||
|
||
def handle_with_branch_name(branch_name: str, worktrees_dir: Path) -> None:
|
||
"""处理有分支名参数的情况
|
||
|
||
Args:
|
||
branch_name: 分支名称
|
||
worktrees_dir: 工作区目录
|
||
"""
|
||
worktree_path = worktrees_dir / branch_name
|
||
|
||
try:
|
||
validate_branch_creation(branch_name, worktree_path)
|
||
except RuntimeError as e:
|
||
print(f"错误: {e}")
|
||
sys.exit(1)
|
||
|
||
remote_branches = get_remote_branches(branch_name)
|
||
|
||
try:
|
||
if remote_branches:
|
||
selected_branch = select_from_list(
|
||
remote_branches,
|
||
"找到远端分支:",
|
||
allow_create=True
|
||
)
|
||
if selected_branch:
|
||
create_worktree(branch_name, worktree_path, selected_branch)
|
||
else:
|
||
create_worktree(branch_name, worktree_path)
|
||
else:
|
||
print("未找到远端分支，创建新分支")
|
||
create_worktree(branch_name, worktree_path)
|
||
|
||
print(f"工作区已创建于 {worktree_path}")
|
||
except RuntimeError as e:
|
||
print(f"错误: {e}")
|
||
sys.exit(1)
|
||
|
||
|
||
def handle_without_branch_name(worktrees_dir: Path) -> None:
|
||
"""处理无分支名参数的情况
|
||
|
||
Args:
|
||
worktrees_dir: 工作区目录
|
||
"""
|
||
remote_branches = get_all_remote_branches()
|
||
|
||
if not remote_branches:
|
||
print("未找到远端分支")
|
||
sys.exit(1)
|
||
|
||
selected_branch = select_from_list(
|
||
remote_branches,
|
||
"远端分支列表:",
|
||
allow_create=True
|
||
)
|
||
|
||
if selected_branch:
|
||
branch_name = extract_branch_name(selected_branch)
|
||
worktree_path = worktrees_dir / branch_name
|
||
|
||
try:
|
||
create_worktree_with_validation(branch_name, worktree_path, selected_branch)
|
||
except RuntimeError as e:
|
||
print(f"错误: {e}")
|
||
sys.exit(1)
|
||
else:
|
||
branch_name = input_branch_name()
|
||
worktree_path = worktrees_dir / branch_name
|
||
|
||
try:
|
||
create_worktree_with_validation(branch_name, worktree_path)
|
||
except RuntimeError as e:
|
||
print(f"错误: {e}")
|
||
sys.exit(1)
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(
|
||
description="创建开发分支工作区",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例:
|
||
%(prog)s # 列出远端分支供选择
|
||
%(prog)s feature-123 # 创建 feature-123 工作区
|
||
%(prog)s bugfix-456 # 创建 bugfix-456 工作区
|
||
"""
|
||
)
|
||
parser.add_argument("branch_name", nargs="?", help="分支名称（可选）")
|
||
|
||
args = parser.parse_args()
|
||
|
||
try:
|
||
root_dir = get_root_dir()
|
||
except RuntimeError as e:
|
||
print(f"错误: {e}")
|
||
sys.exit(1)
|
||
|
||
worktrees_dir = root_dir / ".worktrees"
|
||
worktrees_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
print("正在从远端获取最新分支信息...")
|
||
fetch_remote()
|
||
|
||
if args.branch_name:
|
||
handle_with_branch_name(args.branch_name, worktrees_dir)
|
||
else:
|
||
handle_without_branch_name(worktrees_dir)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|