feat: 新增函数注册机制与 CLI 工具函数模块
CI / Lint, Typecheck & Test (push) Has been cancelled

- 新增 registry.py 提供 register_fn/get_fn/has_fn 函数注册机制, 支持 @register_fn 和 @register_fn("name") 两种用法
- 新增 cli/_ops 包 (files/dev/media/system 四个子模块), 聚合 59 个可复用函数供 YAML fn 字段引用
- 扩展 yaml_loader 支持 fn 字段、args/kwargs 传参、${VAR} 变量占位符
- 新增 test_registry.py (20 个测试) 和扩展 test_yaml_loader.py
- 更新自驱动规则: 自动 commit+push, 删除需要用户明确指示的步骤
This commit is contained in:
2026-07-04 18:24:52 +08:00
parent db02443463
commit 6931f36fd1
13 changed files with 2730 additions and 52 deletions
+1 -1
View File
@@ -150,7 +150,7 @@ uvx --from pyflowx pymake cov
## Git 与提交
- **自动提交/push**:除非用户明确要求
- **自动提交**:任务完成后自动 `git add`(按文件名)+ `git commit` + `git push`(仅当分支已跟踪远程时执行 push;新分支跳过 push 并在总结中说明)
- **不修改 git config**。
- **不运行破坏性命令**`push --force`/`reset --hard`/`clean -f`)除非用户明确要求。
- **staging**:按文件名添加,不用 `git add -A`/`git add .`,避免误加敏感文件。
+6 -5
View File
@@ -21,8 +21,9 @@ alwaysApply: true
1. **目标与范围**:要解决什么问题?交付物是什么?显式列出不在范围内的内容。
2. **验收标准**:怎样算"完成"?可观测的判定条件(功能、性能、覆盖率阈值)。
3. **特殊约束**:除 `python-standards.md` 之外的约束(兼容性、依赖限制、API 兼容策略等)。
4. **提交策略**是否允许自动 `git commit`?是否允许 push?(默认:均不允许,仅在用户明确要求时执行。)
5. **测试要求**:覆盖率门槛(项目默认 ≥95%,branch);是否需要新增 `slow` 标记。
4. **测试要求**覆盖率门槛(项目默认 ≥95%,branch);是否需要新增 `slow` 标记。
**git commit/push 不在确认范围内**:任务完成后自动 commit + push(仅当分支已跟踪远程时执行 push;新分支跳过 push 并在总结中说明),遵循 `.trae/rules/git-commit-message.md` 风格。仅 force-push、reset --hard、clean -f、修改 git config 等真正破坏性操作才需暂停确认。
确认后,将目标与验收标准固化进 `TaskCreate` 任务列表,后续不再就同一信息反复询问。
@@ -75,7 +76,7 @@ alwaysApply: true
## 暂停条件(仅在以下情况中断自驱动找用户)
1. **歧义无法自决**:需求存在多种合理解读且无既有约定可循。
2. **高风险/不可逆操作**:删除文件、`git push --force`、`reset --hard`、删表、修改 CI 配置、修改 git config、卸载依赖等。
2. **高风险/不可逆操作**:删除非临时文件、`git push --force`、`reset --hard`、删表、修改 CI 配置、修改 git config、卸载依赖等。**普通 `git commit`/`push` 不属于此类**(任务完成后自动执行)。
3. **不可恢复的失败**:根因不在本仓库、需外部环境/权限配合、或经两轮尝试仍无法定位。
4. **超出初始确认范围**:用户目标在执行中发现需要显著扩大范围或改变方向。
5. **用户主动询问**:用户在对话中提出新问题或要求澄清。
@@ -105,7 +106,7 @@ alwaysApply: true
**必须暂停询问的典型情况**
- 删除非临时文件、重命名公共模块/包。
- `git commit`/`push`/`force-push`/`reset --hard`/`clean -f`(除非用户已在初始确认中授权)。
- `git push --force``reset --hard``clean -f`、修改 git config(普通 commit/push 自动执行,无需询问)。
- 引入新的运行时依赖(违反项目零依赖原则)。
- 修改 CI 配置、pre-commit 钩子、pyproject.toml 的工具链配置。
- 卸载或降级既有依赖。
@@ -128,6 +129,6 @@ alwaysApply: true
## 收尾
- 验收标准全部满足后,**直接输出最终总结并结束任务**:交付物、关键决策、遗留事项。
- **不询问**"是否需要提交"或"是否扩展范围";用户明确要求提交时再执行 `git commit`/`push`,并按 `.trae/rules/git-commit-message.md` 风格写提交信息
- **自动提交**:收尾时自动 `git add`(按文件名)+ `git commit`(遵循 `.trae/rules/git-commit-message.md` 风格)+ `git push`(仅当分支已跟踪远程时执行;新分支跳过 push 并在总结中说明);**不询问**"是否需要提交"或"是否扩展范围"。
- 若验收标准未全部满足,回到「计划」继续下一轮,不停下询问。
- 将本次会话的关键产出与决策更新到 memory,便于后续会话续接。
+5
View File
@@ -83,6 +83,7 @@ from .errors import (
from .executors import Strategy, run
from .graph import Graph, GraphDefaults
from .profiling import ProfileReport, TaskProfile
from .registry import FnRegistry, get_fn, has_fn, register_fn
from .report import RunReport
from .runner import CliExitCode, CliRunner
from .storage import JSONBackend, MemoryBackend, StateBackend
@@ -117,6 +118,7 @@ __all__ = [
"Context",
"CycleError",
"DuplicateTaskError",
"FnRegistry",
"Graph",
"GraphComposer",
"GraphDefaults",
@@ -145,8 +147,11 @@ __all__ = [
"cmd",
"compose",
"describe_injection",
"get_fn",
"has_fn",
"load_yaml",
"parse_yaml_string",
"register_fn",
"run",
"run_command",
"task",
+18
View File
@@ -0,0 +1,18 @@
"""CLI 工具通用函数模块.
按类别组织 CLI 工具中可复用的函数, 每个子模块使用 ``@px.register_fn`` 注册函数,
供 YAML 任务编排通过 ``fn`` 字段引用.
子模块
------
- :mod:`files` —— 文件日期/等级/备份/压缩相关函数
- :mod:`dev` —— 开发工具 (ruff/版本号/pip/git) 相关函数
- :mod:`media` —— PDF/截图相关函数
- :mod:`system` —— LS-DYNA/SSH/打包相关函数
"""
from __future__ import annotations
from . import dev, files, media, system
__all__ = ["dev", "files", "media", "system"]
+514
View File
@@ -0,0 +1,514 @@
"""开发工具类函数模块.
聚合自动格式化 (autofmt)、版本号管理 (bumpversion)、pip 包管理 (piptool)、
git 工具 (gittool) 的可复用函数. 所有公共函数通过 ``@px.register_fn`` 注册,
供 YAML 任务编排引用.
"""
from __future__ import annotations
import ast
import fnmatch
import re
import subprocess
from pathlib import Path
from typing import Literal
import pyflowx as px
__all__ = [
"IGNORE_PATTERNS",
"PACKAGE_DIR",
"REQUIREMENTS_FILE",
"_PROTECTED_PACKAGES",
"BumpVersionType",
"add_docstring",
"auto_add_docstrings",
"bump_file_version",
"format_all",
"format_with_ruff",
"generate_module_docstring",
"has_files",
"init_sub_dirs",
"lint_with_ruff",
"not_has_git_repo",
"pip_download",
"pip_freeze",
"pip_reinstall",
"pip_uninstall",
"sync_pyproject_config",
]
# ============================================================================
# autofmt 配置
# ============================================================================
IGNORE_PATTERNS = [
"__pycache__",
"*.pyc",
"*.pyo",
".git",
".venv",
".idea",
".vscode",
"*.egg-info",
"dist",
"build",
".pytest_cache",
".tox",
".mypy_cache",
]
# ============================================================================
# bumpversion 配置
# ============================================================================
BumpVersionType = Literal["patch", "minor", "major"]
_PYPROJECT_VERSION_PATTERN = re.compile(
r'(?:^|\n)\s*version\s*=\s*["\']'
r"(?P<major>0|[1-9]\d*)\.(?P<minor>0|[1-9]\d*)\.(?P<patch>0|[1-9]\d*)"
r"(?:-(?P<prerelease>(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?"
r"(?:\+(?P<buildmetadata>[0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?"
r'["\']',
re.MULTILINE,
)
_INIT_VERSION_PATTERN = re.compile(
r'(?:^|\n)\s*__version__\s*=\s*["\']'
r"(?P<major>0|[1-9]\d*)\.(?P<minor>0|[1-9]\d*)\.(?P<patch>0|[1-9]\d*)"
r"(?:-(?P<prerelease>(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?"
r"(?:\+(?P<buildmetadata>[0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?"
r'["\']',
re.MULTILINE,
)
# ============================================================================
# piptool 配置
# ============================================================================
PACKAGE_DIR = "packages"
REQUIREMENTS_FILE = "requirements.txt"
_PROTECTED_PACKAGES: frozenset[str] = frozenset(
{
"pyflowx",
"bitool",
}
)
# ============================================================================
# autofmt 私有辅助函数
# ============================================================================
# ============================================================================
# autofmt 函数
# ============================================================================
@px.register_fn
def format_with_ruff(target: Path, fix: bool = True) -> None:
"""使用 ruff 格式化代码.
Parameters
----------
target : Path
目标路径
fix : bool
是否自动修复
"""
cmd = ["ruff", "format", str(target)]
if fix:
cmd.append("--fix")
subprocess.run(cmd, check=True)
print(f"ruff format 完成: {target}")
@px.register_fn
def lint_with_ruff(target: Path, fix: bool = True) -> None:
"""使用 ruff 检查代码.
Parameters
----------
target : Path
目标路径
fix : bool
是否自动修复
"""
cmd = ["ruff", "check", str(target)]
if fix:
cmd.extend(["--fix", "--unsafe-fixes"])
subprocess.run(cmd, check=True)
print(f"ruff check 完成: {target}")
@px.register_fn
def add_docstring(file_path: Path, docstring: str) -> bool:
"""为文件添加 docstring.
Parameters
----------
file_path : Path
文件路径
docstring : str
docstring 内容
Returns
-------
bool
是否成功添加
"""
try:
content = file_path.read_text(encoding="utf-8")
tree = ast.parse(content)
first_node = tree.body[0] if tree.body else None
if first_node and isinstance(first_node, ast.Expr) and isinstance(first_node.value, ast.Constant):
return False
lines = content.splitlines()
doc_lines = docstring.splitlines()
doc_lines.append("")
new_content = "\n".join(doc_lines + lines)
file_path.write_text(new_content, encoding="utf-8")
print(f"添加 docstring: {file_path}")
return True
except (OSError, UnicodeDecodeError, SyntaxError) as e:
print(f"处理失败: {file_path} - {e}")
return False
@px.register_fn
def generate_module_docstring(file_path: Path) -> str:
"""生成模块 docstring.
Parameters
----------
file_path : Path
文件路径
Returns
-------
str
生成的 docstring
"""
stem = file_path.stem
parent = file_path.parent.name
keywords = {
"cli": f"Command-line interface for {parent}",
"gui": f"Graphical user interface for {parent}",
"core": f"Core functionality for {parent}",
"util": f"Utility functions for {parent}",
"model": f"Data models for {parent}",
"test": f"Tests for {parent}",
}
for key, desc in keywords.items():
if key in stem.lower():
return f'"""{desc}."""'
return f'"""{stem.replace("_", " ").title()} module."""'
@px.register_fn
def auto_add_docstrings(root_dir: Path) -> int:
"""自动为所有 Python 文件添加 docstring.
Parameters
----------
root_dir : Path
根目录
Returns
-------
int
添加的 docstring 数量
"""
count = 0
for py_file in root_dir.rglob("*.py"):
if any(pattern in str(py_file) for pattern in IGNORE_PATTERNS):
continue
docstring = generate_module_docstring(py_file)
if add_docstring(py_file, docstring):
count += 1
print(f"共添加 {count} 个 docstring")
return count
@px.register_fn
def sync_pyproject_config(root_dir: Path) -> None:
"""同步 pyproject.toml 配置到子项目.
Parameters
----------
root_dir : Path
根目录
"""
main_toml = root_dir / "pyproject.toml"
if not main_toml.exists():
print(f"主项目配置文件不存在: {main_toml}")
return
sub_tomls = [p for p in root_dir.rglob("pyproject.toml") if p != main_toml and ".venv" not in str(p)]
if not sub_tomls:
print("没有找到子项目的 pyproject.toml")
return
print(f"找到 {len(sub_tomls)} 个子项目配置文件")
for sub_toml in sub_tomls:
subprocess.run(["ruff", "format", str(sub_toml)], check=False)
print("配置同步完成")
@px.register_fn
def format_all(root_dir: Path) -> None:
"""格式化所有 Python 文件.
Parameters
----------
root_dir : Path
根目录
"""
subprocess.run(["ruff", "format", str(root_dir)], check=True)
subprocess.run(["ruff", "check", "--fix", "--unsafe-fixes", str(root_dir)], check=True)
print(f"格式化完成: {root_dir}")
# ============================================================================
# bumpversion 私有辅助函数
# ============================================================================
def _get_pattern_for_file(file_name: str) -> re.Pattern[str] | None:
"""根据文件类型获取对应的正则表达式."""
if file_name == "pyproject.toml":
return _PYPROJECT_VERSION_PATTERN
if file_name == "__init__.py":
return _INIT_VERSION_PATTERN
return None
def _calculate_new_version(major: int, minor: int, patch: int, part: BumpVersionType) -> str:
"""计算新版本号."""
if part == "major":
return f"{major + 1}.0.0"
if part == "minor":
return f"{major}.{minor + 1}.0"
return f"{major}.{minor}.{patch + 1}"
def _build_replacement_string(original_match: str, new_version: str, file_name: str) -> str:
"""构建替换字符串, 保留原始格式."""
quote_char = '"' if '"' in original_match else "'"
if file_name == "pyproject.toml":
prefix_match = re.match(r'(\s*version\s*=\s*)["\']', original_match)
prefix = prefix_match.group(1) if prefix_match else "version = "
return f"{prefix}{quote_char}{new_version}{quote_char}"
if file_name == "__init__.py":
prefix_match = re.match(r'(\s*__version__\s*=\s*)["\']', original_match)
prefix = prefix_match.group(1) if prefix_match else "__version__ = "
return f"{prefix}{quote_char}{new_version}{quote_char}"
return new_version
# ============================================================================
# bumpversion 函数
# ============================================================================
@px.register_fn
def bump_file_version(file_path: Path, part: BumpVersionType = "patch") -> str | None:
"""更新文件中的版本号.
Parameters
----------
file_path : Path
要更新的文件路径
part : BumpVersionType
版本部分: patch, minor, major
Returns
-------
str | None
更新后的新版本号, 如果文件中未找到版本号则返回 None
"""
try:
content = file_path.read_text(encoding="utf-8")
except Exception as e:
print(f"读取文件 {file_path} 时出错: {e}")
raise
pattern = _get_pattern_for_file(file_path.name)
if pattern:
match = pattern.search(content)
else:
match = _PYPROJECT_VERSION_PATTERN.search(content) or _INIT_VERSION_PATTERN.search(content)
if not match:
print(f"文件 {file_path} 中未找到版本号模式")
return None
major = int(match.group("major"))
minor = int(match.group("minor"))
patch = int(match.group("patch"))
new_version = _calculate_new_version(major, minor, patch, part)
original_match = match.group(0)
replacement = _build_replacement_string(original_match, new_version, file_path.name)
content = content.replace(original_match, replacement)
try:
file_path.write_text(content, encoding="utf-8")
except Exception as e:
print(f"更新文件 {file_path} 版本号时出错: {e}")
raise
return new_version
# ============================================================================
# piptool 私有辅助函数
# ============================================================================
def _get_installed_packages() -> list[str]:
"""获取当前环境中所有已安装的包名."""
try:
result = subprocess.run(
["pip", "list", "--format=freeze"],
capture_output=True,
text=True,
check=True,
)
packages: list[str] = []
for line in result.stdout.strip().split("\n"):
if line and "==" in line:
pkg_name = line.split("==")[0].strip()
packages.append(pkg_name)
except (subprocess.SubprocessError, OSError):
return []
return packages
def _expand_wildcard_packages(pattern: str) -> list[str]:
"""展开通配符模式为实际的包名列表."""
if not any(char in pattern for char in ["*", "?", "[", "]"]):
return [pattern]
installed_packages = _get_installed_packages()
matched = [pkg for pkg in installed_packages if fnmatch.fnmatchcase(pkg.lower(), pattern.lower())]
return matched
def _filter_protected_packages(packages: list[str]) -> list[str]:
"""过滤掉受保护的包名."""
safe = [p for p in packages if p.lower() not in {p.lower() for p in _PROTECTED_PACKAGES}]
filtered = [p for p in packages if p.lower() in {p.lower() for p in _PROTECTED_PACKAGES}]
if filtered:
print(f"跳过受保护的包: {', '.join(filtered)}")
return safe
# ============================================================================
# piptool 函数
# ============================================================================
@px.register_fn
def pip_uninstall(pkg_names: list[str]) -> None:
"""卸载包."""
packages_to_uninstall: list[str] = []
for pattern in pkg_names:
packages_to_uninstall.extend(_expand_wildcard_packages(pattern))
packages_to_uninstall = _filter_protected_packages(packages_to_uninstall)
if not packages_to_uninstall:
return
subprocess.run(["pip", "uninstall", "-y", *packages_to_uninstall], check=True)
@px.register_fn
def pip_reinstall(pkg_names: list[str], offline: bool = False) -> None:
"""重新安装包."""
safe_ps = _filter_protected_packages(pkg_names)
if not safe_ps:
print("所有指定的包均为受保护包, 跳过重装")
return
subprocess.run(["pip", "uninstall", "-y", *safe_ps], check=True)
options = ["--no-index", "--find-links", "."] if offline else []
subprocess.run(["pip", "install", *options, *safe_ps], check=True)
@px.register_fn
def pip_download(pkg_names: list[str], offline: bool = False) -> None:
"""下载包."""
options = ["--no-index", "--find-links", "."] if offline else []
subprocess.run(
["pip", "download", *pkg_names, *options, "-d", PACKAGE_DIR],
check=True,
)
@px.register_fn
def pip_freeze() -> None:
"""冻结依赖."""
result = subprocess.run(
["pip", "freeze", "--exclude-editable"],
capture_output=True,
text=True,
check=True,
)
Path(REQUIREMENTS_FILE).write_text(result.stdout)
# ============================================================================
# gittool 函数
# ============================================================================
@px.register_fn
def init_sub_dirs() -> None:
"""初始化子目录的 Git 仓库."""
sub_dirs = [subdir for subdir in Path.cwd().iterdir() if subdir.is_dir()]
for subdir in sub_dirs:
px.run(
px.Graph().chain(
px.cmd(["git", "init"], conditions=(lambda _: not_has_git_repo(),), cwd=subdir),
px.cmd(["git", "add", "."]),
px.cmd(["git", "commit", "-m", "init commit"]),
),
)
@px.register_fn
def not_has_git_repo() -> bool:
"""检查当前目录没有 Git 仓库."""
return not Path.cwd().exists() or not (Path.cwd() / ".git").is_dir()
@px.register_fn
def has_files() -> bool:
"""检查当前目录是否有文件."""
return bool(list(Path.cwd().glob("*")))
+327
View File
@@ -0,0 +1,327 @@
"""文件类函数模块.
聚合文件日期处理、文件等级重命名、文件夹备份、文件夹压缩工具的可复用函数.
所有公共函数通过 ``@px.register_fn`` 注册, 供 YAML 任务编排引用.
"""
from __future__ import annotations
import re
import shutil
import time
import zipfile
from pathlib import Path
import pyflowx as px
__all__ = [
"BRACKETS",
"DATE_PATTERN",
"IGNORE_DIRS",
"IGNORE_EXT",
"IGNORE_FILES",
"LEVELS",
"SEP",
"add_date_prefix",
"archive_folder",
"backup_folder",
"folderback_default",
"folderzip_default",
"get_file_timestamp",
"process_file_date",
"process_file_level",
"process_files_date",
"process_files_level",
"remove_date_prefix",
"remove_dump",
"remove_marks",
"zip_folders",
"zip_target",
]
# ============================================================================
# filedate 配置
# ============================================================================
DATE_PATTERN = re.compile(r"(20|19)\d{2}[-_#.~]?((0[1-9])|(1[012]))[-_#.~]?((0[1-9])|([12]\d)|(3[01]))[-_#.~]?")
SEP = "_"
# ============================================================================
# filelevel 配置
# ============================================================================
LEVELS: dict[str, str] = {
"0": "",
"1": "PUB,NOR",
"2": "INT",
"3": "CON",
"4": "CLA",
}
BRACKETS: tuple[str, str] = (" ([_(【-", " )]_)】")
# ============================================================================
# folderzip 配置
# ============================================================================
IGNORE_DIRS: list[str] = [".git", ".idea", ".vscode", "__pycache__"]
IGNORE_FILES: list[str] = [".gitignore"]
IGNORE: list[str] = [*IGNORE_DIRS, *IGNORE_FILES]
IGNORE_EXT: list[str] = [".zip", ".rar", ".7z", ".tar", ".gz"]
# ============================================================================
# filedate 函数
# ============================================================================
@px.register_fn
def get_file_timestamp(filepath: Path) -> str:
"""获取文件时间戳."""
modified_time = filepath.stat().st_mtime
created_time = filepath.stat().st_ctime
return time.strftime("%Y%m%d", time.localtime(max((modified_time, created_time))))
@px.register_fn
def remove_date_prefix(filepath: Path) -> Path:
"""移除文件日期前缀."""
stem = filepath.stem
new_stem = DATE_PATTERN.sub("", stem)
if new_stem != stem:
new_path = filepath.with_name(new_stem + filepath.suffix)
filepath.rename(new_path)
return new_path
return filepath
@px.register_fn
def add_date_prefix(filepath: Path) -> Path:
"""添加文件日期前缀."""
timestamp = get_file_timestamp(filepath)
stem = filepath.stem
new_stem = f"{timestamp}{SEP}{stem}"
new_path = filepath.with_name(new_stem + filepath.suffix)
if new_path != filepath:
filepath.rename(new_path)
return new_path
return filepath
@px.register_fn
def process_file_date(filepath: Path, clear: bool = False) -> None:
"""处理单个文件的日期前缀.
Parameters
----------
filepath : Path
文件路径
clear : bool
是否清除日期前缀
"""
if clear:
remove_date_prefix(filepath)
else:
new_path = remove_date_prefix(filepath)
add_date_prefix(new_path)
@px.register_fn
def process_files_date(targets: list[Path], clear: bool = False) -> None:
"""批量处理文件日期前缀.
Parameters
----------
targets : list[Path]
文件路径列表
clear : bool
是否清除日期前缀
"""
for target in targets:
if target.exists() and not target.name.startswith("."):
process_file_date(target, clear)
# ============================================================================
# filelevel 函数
# ============================================================================
@px.register_fn
def remove_marks(stem: str, marks: list[str]) -> str:
"""从文件名主干中移除所有标记."""
left_brackets, right_brackets = BRACKETS
for mark in marks:
pos = 0
while True:
pos = stem.find(mark, pos)
if pos == -1:
break
b, e = pos - 1, pos + len(mark)
if b >= 0 and e < len(stem) and stem[b] in left_brackets and stem[e] in right_brackets:
stem = stem[:b] + stem[e + 1 :]
else:
pos = e
return stem
@px.register_fn
def process_file_level(filepath: Path, level: int = 0) -> None:
"""处理单个文件的等级标记.
Parameters
----------
filepath : Path
文件路径
level : int
文件等级 (0-4), 0 用于清除等级
"""
if not (0 <= level < len(LEVELS)):
print(f"无效的等级 {level}, 必须在 0 和 {len(LEVELS) - 1} 之间")
return
if not filepath.exists():
print(f"文件不存在: {filepath}")
return
filestem = filepath.stem
original_stem = filestem
for level_names in LEVELS.values():
if level_names:
filestem = remove_marks(filestem, level_names.split(","))
for digit in map(str, range(1, 10)):
filestem = remove_marks(filestem, [digit])
if level > 0:
levelstr = LEVELS.get(str(level), "").split(",")[0]
if levelstr:
filestem = f"{filestem}({levelstr})"
if filestem != original_stem:
new_path = filepath.with_name(filestem + filepath.suffix)
filepath.rename(new_path)
print(f"重命名: {filepath} -> {new_path}")
@px.register_fn
def process_files_level(targets: list[Path], level: int = 0) -> None:
"""批量处理文件等级标记.
Parameters
----------
targets : list[Path]
文件路径列表
level : int
文件等级 (0-4)
"""
for target in targets:
process_file_level(target, level)
# ============================================================================
# folderback 函数
# ============================================================================
@px.register_fn
def remove_dump(src: Path, dst: Path, max_zip: int) -> None:
"""递归删除旧的备份 zip 文件."""
zip_paths = [filepath for filepath in dst.rglob("*.zip") if src.stem in str(filepath)]
zip_files = sorted(zip_paths, key=lambda fn: str(fn)[-19:-4])
if len(zip_files) > max_zip:
zip_files[0].unlink()
remove_dump(src, dst, max_zip)
@px.register_fn
def zip_target(src: Path, dst: Path, max_zip: int) -> None:
"""将单个文件或文件夹压缩为 zip 文件."""
files = [str(_) for _ in src.rglob("*")]
timestamp = time.strftime("_%Y%m%d_%H%M%S")
target_path = dst / (src.stem + timestamp + ".zip")
with zipfile.ZipFile(target_path, "w") as zip_file:
for file in files:
zip_file.write(file, arcname=file.replace(str(src.parent), ""))
remove_dump(src, dst, max_zip)
print(f"备份完成: {target_path}")
@px.register_fn
def backup_folder(src: str, dst: str, max_zip: int = 5) -> None:
"""备份文件夹.
Parameters
----------
src : str
源文件夹路径
dst : str
目标文件夹路径
max_zip : int
最大备份数量
"""
src_path = Path(src)
dst_path = Path(dst)
if not src_path.exists():
print(f"源文件夹不存在: {src_path}")
return
if not dst_path.exists():
dst_path.mkdir(parents=True, exist_ok=True)
print(f"创建目标文件夹: {dst_path}")
zip_target(src_path, dst_path, max_zip)
@px.register_fn("folderback_default")
def folderback_default() -> None:
"""备份当前目录到 ./backup."""
backup_folder(".", "./backup", 5)
# ============================================================================
# folderzip 函数
# ============================================================================
@px.register_fn
def archive_folder(folder: Path) -> None:
"""压缩单个文件夹."""
shutil.make_archive(
str(folder.with_name(folder.name)),
format="zip",
base_dir=folder,
)
print(f"压缩完成: {folder.name}.zip")
@px.register_fn
def zip_folders(cwd: str = ".") -> None:
"""压缩目录下的所有文件夹.
Parameters
----------
cwd : str
工作目录
"""
cwd_path = Path(cwd)
if not cwd_path.exists():
print(f"目录不存在: {cwd_path}")
return
dirs: list[Path] = [
e for e in cwd_path.iterdir() if e.is_dir() and e.name not in IGNORE_DIRS and e.suffix not in IGNORE_EXT
]
for dir_path in dirs:
archive_folder(dir_path)
@px.register_fn("folderzip_default")
def folderzip_default() -> None:
"""压缩当前目录下的所有文件夹."""
zip_folders(".")
+498
View File
@@ -0,0 +1,498 @@
"""媒体类函数模块.
聚合 PDF 工具 (pdftool) 和截图工具 (screenshot) 的可复用函数.
所有公共函数通过 ``@px.register_fn`` 注册, 供 YAML 任务编排引用.
"""
from __future__ import annotations
import subprocess
from datetime import datetime
from pathlib import Path
import pyflowx as px
from pyflowx.conditions import Constants
__all__ = [
"DEFAULT_PASSWORD",
"DEFAULT_QUALITY",
"PDF_SUFFIX",
"get_screenshot_path",
"pdf_add_watermark",
"pdf_compress",
"pdf_crop",
"pdf_decrypt",
"pdf_encrypt",
"pdf_extract_images",
"pdf_extract_text",
"pdf_info",
"pdf_merge",
"pdf_ocr",
"pdf_reorder",
"pdf_repair",
"pdf_rotate",
"pdf_split",
"pdf_to_images",
"take_screenshot_area",
"take_screenshot_full",
]
try:
import fitz # PyMuPDF
HAS_PYMUPDF = True
except ImportError:
HAS_PYMUPDF = False
try:
import pypdf
HAS_PYPDF = True
except ImportError:
HAS_PYPDF = False
# ============================================================================
# 配置
# ============================================================================
PDF_SUFFIX = ".pdf"
DEFAULT_QUALITY = 75
DEFAULT_PASSWORD = ""
# ============================================================================
# PDF 函数
# ============================================================================
@px.register_fn
def pdf_merge(input_paths: list[Path], output_path: Path) -> None:
"""合并多个 PDF 文件."""
if not HAS_PYPDF:
print("未安装 pypdf 库, 请安装: pip install pypdf")
return
writer = pypdf.PdfWriter()
for input_path in input_paths:
if input_path.exists():
reader = pypdf.PdfReader(str(input_path))
for page in reader.pages:
writer.add_page(page)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "wb") as f:
writer.write(f)
print(f"合并完成: {output_path}")
@px.register_fn
def pdf_split(input_path: Path, output_dir: Path) -> None:
"""拆分 PDF 文件为单页."""
if not HAS_PYPDF:
print("未安装 pypdf 库, 请安装: pip install pypdf")
return
reader = pypdf.PdfReader(str(input_path))
output_dir.mkdir(parents=True, exist_ok=True)
for i, page in enumerate(reader.pages):
writer = pypdf.PdfWriter()
writer.add_page(page)
output_file = output_dir / f"{input_path.stem}_page_{i + 1}.pdf"
with open(output_file, "wb") as f:
writer.write(f)
print(f"拆分完成: {output_dir}")
@px.register_fn
def pdf_compress(input_path: Path, output_path: Path) -> None:
"""压缩 PDF 文件."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
output_path.parent.mkdir(parents=True, exist_ok=True)
doc.save(str(output_path), garbage=4, deflate=True, clean=True)
doc.close()
original_size = input_path.stat().st_size
new_size = output_path.stat().st_size
ratio = (1 - new_size / original_size) * 100
print(f"压缩完成: {output_path} (缩小 {ratio:.1f}%)")
@px.register_fn
def pdf_encrypt(input_path: Path, output_path: Path, password: str) -> None:
"""加密 PDF 文件."""
if not HAS_PYPDF:
print("未安装 pypdf 库, 请安装: pip install pypdf")
return
reader = pypdf.PdfReader(str(input_path))
writer = pypdf.PdfWriter()
for page in reader.pages:
writer.add_page(page)
writer.encrypt(user_password=password, owner_password=password, use_128bit=True)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "wb") as f:
writer.write(f)
print(f"加密完成: {output_path}")
@px.register_fn
def pdf_decrypt(input_path: Path, output_path: Path, password: str) -> None:
"""解密 PDF 文件."""
if not HAS_PYPDF:
print("未安装 pypdf 库, 请安装: pip install pypdf")
return
reader = pypdf.PdfReader(str(input_path))
if reader.is_encrypted:
reader.decrypt(password)
writer = pypdf.PdfWriter()
for page in reader.pages:
writer.add_page(page)
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "wb") as f:
writer.write(f)
print(f"解密完成: {output_path}")
@px.register_fn
def pdf_extract_text(input_path: Path, output_path: Path) -> None:
"""提取 PDF 文本."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
text = ""
for page in doc:
text += str(page.get_text()) + "\n\n"
doc.close()
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(text, encoding="utf-8")
print(f"文本提取完成: {output_path}")
@px.register_fn
def pdf_extract_images(input_path: Path, output_dir: Path) -> None:
"""提取 PDF 图片."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
output_dir.mkdir(parents=True, exist_ok=True)
image_count = 0
# pyrefly: ignore [bad-argument-type]
for page_num, page in enumerate(doc):
images = page.get_images(full=True)
for img_idx, img in enumerate(images):
xref = img[0]
base_image = doc.extract_image(xref)
image_data = base_image["image"]
image_ext = base_image["ext"]
image_path = output_dir / f"page_{page_num + 1}_img_{img_idx + 1}.{image_ext}"
image_path.write_bytes(image_data)
image_count += 1
doc.close()
print(f"图片提取完成: {output_dir} (共 {image_count} 张)")
@px.register_fn
def pdf_add_watermark(input_path: Path, output_path: Path, text: str = "CONFIDENTIAL") -> None:
"""添加 PDF 水印."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
for page in doc:
rect = page.rect
text_width = fitz.get_text_length(text, fontsize=48)
x = (rect.width - text_width) / 2
y = rect.height / 2
page.insert_text((x, y), text, fontsize=48, rotate=45, color=(0, 0, 0))
output_path.parent.mkdir(parents=True, exist_ok=True)
doc.save(str(output_path))
doc.close()
print(f"水印添加完成: {output_path}")
@px.register_fn
def pdf_rotate(input_path: Path, output_path: Path, rotation: int = 90) -> None:
"""旋转 PDF 页面."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
for page in doc:
page.set_rotation(rotation)
output_path.parent.mkdir(parents=True, exist_ok=True)
doc.save(str(output_path))
doc.close()
print(f"旋转完成: {output_path}")
@px.register_fn
def pdf_crop(input_path: Path, output_path: Path, margins: tuple[int, int, int, int]) -> None:
"""裁剪 PDF 页面."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
left, top, right, bottom = margins
for page in doc:
rect = page.rect
new_rect = fitz.Rect(
rect.x0 + left,
rect.y0 + top,
rect.x1 - right,
rect.y1 - bottom,
)
page.set_cropbox(new_rect)
output_path.parent.mkdir(parents=True, exist_ok=True)
doc.save(str(output_path))
doc.close()
print(f"裁剪完成: {output_path}")
@px.register_fn
def pdf_info(input_path: Path) -> None:
"""显示 PDF 信息."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
print(f"文件: {input_path}")
print(f"页数: {doc.page_count}")
# pyrefly: ignore [missing-attribute]
print(f"标题: {doc.metadata.get('title', 'N/A')}")
# pyrefly: ignore [missing-attribute]
print(f"作者: {doc.metadata.get('author', 'N/A')}")
# pyrefly: ignore [missing-attribute]
print(f"创建日期: {doc.metadata.get('creationDate', 'N/A')}")
# pyrefly: ignore [missing-attribute]
print(f"修改日期: {doc.metadata.get('modDate', 'N/A')}")
print(f"文件大小: {input_path.stat().st_size / 1024:.1f} KB")
doc.close()
@px.register_fn
def pdf_ocr(input_path: Path, output_path: Path, lang: str = "chi_sim+eng") -> None:
"""PDF OCR 识别."""
try:
import pytesseract
from PIL import Image
except ImportError:
print("未安装 OCR 相关库, 请安装: pip install pytesseract pillow")
return
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
new_doc = fitz.open()
for page in doc:
pix = page.get_pixmap()
img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
ocr_text = pytesseract.image_to_string(img, lang=lang)
new_page = new_doc.new_page(width=page.rect.width, height=page.rect.height)
new_page.insert_image(new_page.rect, pixmap=pix)
text_rect = fitz.Rect(0, 0, page.rect.width, page.rect.height)
# pyrefly: ignore [bad-argument-type]
new_page.insert_textbox(text_rect, ocr_text, fontname="china-ss", fontsize=11)
output_path.parent.mkdir(parents=True, exist_ok=True)
new_doc.save(str(output_path))
new_doc.close()
doc.close()
print(f"OCR 识别完成: {output_path}")
@px.register_fn
def pdf_reorder(input_path: Path, output_path: Path, order: list[int]) -> None:
"""重排 PDF 页面顺序."""
if not HAS_PYPDF:
print("未安装 pypdf 库, 请安装: pip install pypdf")
return
reader = pypdf.PdfReader(str(input_path))
writer = pypdf.PdfWriter()
for page_num in order:
if 0 <= page_num < len(reader.pages):
writer.add_page(reader.pages[page_num])
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "wb") as f:
writer.write(f)
print(f"重排完成: {output_path}")
@px.register_fn
def pdf_to_images(input_path: Path, output_dir: Path, dpi: int = 300) -> None:
"""PDF 转图片."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
output_dir.mkdir(parents=True, exist_ok=True)
# pyrefly: ignore [bad-argument-type]
for page_num, page in enumerate(doc):
pix = page.get_pixmap(dpi=dpi)
image_path = output_dir / f"{input_path.stem}_page_{page_num + 1}.png"
pix.save(str(image_path))
doc.close()
print(f"转换完成: {output_dir}")
@px.register_fn
def pdf_repair(input_path: Path, output_path: Path) -> None:
"""修复 PDF 文件."""
if not HAS_PYMUPDF:
print("未安装 PyMuPDF 库, 请安装: pip install PyMuPDF")
return
doc = fitz.open(str(input_path))
output_path.parent.mkdir(parents=True, exist_ok=True)
doc.save(str(output_path), garbage=4, deflate=True, clean=True)
doc.close()
print(f"修复完成: {output_path}")
# ============================================================================
# screenshot 函数
# ============================================================================
@px.register_fn
def get_screenshot_path(filename: str | None = None) -> Path:
"""获取截图保存路径.
Parameters
----------
filename : str | None
文件名, 如果为 None 则自动生成
Returns
-------
Path
截图保存路径
"""
if filename is None:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"screenshot_{timestamp}.png"
screenshots_dir = Path.home() / "Pictures" / "screenshots"
screenshots_dir.mkdir(parents=True, exist_ok=True)
return screenshots_dir / filename
@px.register_fn
def take_screenshot_full(filename: str | None = None) -> None:
"""全屏截图.
Parameters
----------
filename : str | None
文件名
"""
output_path = get_screenshot_path(filename)
if Constants.IS_WINDOWS:
ps_script = f"""
Add-Type -AssemblyName System.Windows.Forms
Add-Type -AssemblyName System.Drawing
$screen = [System.Windows.Forms.Screen]::PrimaryScreen
$bounds = $screen.Bounds
$bitmap = New-Object System.Drawing.Bitmap $bounds.Width, $bounds.Height
$graphics = [System.Drawing.Graphics]::FromImage($bitmap)
$graphics.CopyFromScreen($bounds.Location, [System.Drawing.Point]::Empty, $bounds.Size)
$bitmap.Save('{output_path.as_posix()}')
$graphics.Dispose()
$bitmap.Dispose()
"""
subprocess.run(["powershell", "-Command", ps_script], check=True)
elif Constants.IS_MACOS:
subprocess.run(["screencapture", "-x", str(output_path)], check=True)
else:
try:
subprocess.run(["gnome-screenshot", "-f", str(output_path)], check=True)
except FileNotFoundError:
subprocess.run(["scrot", str(output_path)], check=True)
print(f"截图已保存: {output_path}")
@px.register_fn
def take_screenshot_area(filename: str | None = None) -> None:
"""区域截图.
Parameters
----------
filename : str | None
文件名
"""
output_path = get_screenshot_path(filename)
if Constants.IS_WINDOWS:
ps_script = f"""
Add-Type -AssemblyName System.Windows.Forms
Add-Type -AssemblyName System.Drawing
$form = New-Object System.Windows.Forms.Form
$form.WindowState = 'Maximized'
$form.FormBorderStyle = 'None'
$form.BackColor = [System.Drawing.Color]::FromArgb(1, 0, 0)
$form.Opacity = 0.5
$form.TopMost = $true
$form.Show()
Start-Sleep -Milliseconds 100
$screen = [System.Windows.Forms.Screen]::PrimaryScreen
$bounds = $screen.Bounds
$bitmap = New-Object System.Drawing.Bitmap $bounds.Width, $bounds.Height
$graphics = [System.Drawing.Graphics]::FromImage($bitmap)
$graphics.CopyFromScreen($bounds.Location, [System.Drawing.Point]::Empty, $bounds.Size)
$form.Close()
$bitmap.Save('{output_path.as_posix()}')
$graphics.Dispose()
$bitmap.Dispose()
"""
subprocess.run(["powershell", "-Command", ps_script], check=True)
elif Constants.IS_MACOS:
subprocess.run(["screencapture", "-i", str(output_path)], check=True)
else:
try:
subprocess.run(["gnome-screenshot", "-a", "-f", str(output_path)], check=True)
except FileNotFoundError:
subprocess.run(["scrot", "-s", str(output_path)], check=True)
print(f"截图已保存: {output_path}")
+458
View File
@@ -0,0 +1,458 @@
"""系统类函数模块.
聚合 LS-DYNA 计算 (lscalc)、SSH 密钥部署 (sshcopyid)、Python 打包 (packtool)
的可复用函数. 所有公共函数通过 ``@px.register_fn`` 注册, 供 YAML 任务编排引用.
"""
from __future__ import annotations
import platform
import shutil
import subprocess
import sys
import urllib.request
import zipfile
from pathlib import Path
import pyflowx as px
from pyflowx.conditions import Constants
__all__ = [
"DEFAULT_BUILD_DIR",
"DEFAULT_CACHE_DIR",
"DEFAULT_DIST_DIR",
"DEFAULT_INPUT_FILE",
"DEFAULT_LIB_DIR",
"DEFAULT_NCPU",
"IGNORE_PATTERNS",
"LS_DYNA_COMMANDS",
"check_ls_dyna_status",
"clean_build_dir",
"create_zip_package",
"get_ls_dyna_command",
"install_embed_python",
"pack_dependencies",
"pack_source",
"pack_wheel",
"run_ls_dyna",
"run_ls_dyna_mpi",
"ssh_copy_id",
]
# ============================================================================
# lscalc 配置
# ============================================================================
LS_DYNA_COMMANDS: dict[str, list[str]] = {
"windows": ["ls-dyna_mpp", "i=input.k", "ncpu=4"],
"linux": ["ls-dyna_mpp", "i=input.k", "ncpu=8"],
"macos": ["ls-dyna_mpp", "i=input.k", "ncpu=4"],
}
DEFAULT_INPUT_FILE: str = "input.k"
DEFAULT_NCPU: int = 4
# ============================================================================
# packtool 配置
# ============================================================================
DEFAULT_BUILD_DIR = ".pypack"
DEFAULT_DIST_DIR = "dist"
DEFAULT_LIB_DIR = "libs"
DEFAULT_CACHE_DIR = ".cache/pypack"
IGNORE_PATTERNS = [
"__pycache__",
"*.pyc",
"*.pyo",
".git",
".venv",
".idea",
".vscode",
"*.egg-info",
"dist",
"build",
".pytest_cache",
".tox",
".mypy_cache",
]
# ============================================================================
# lscalc 函数
# ============================================================================
@px.register_fn
def get_ls_dyna_command(input_file: str, ncpu: int) -> list[str]:
"""获取 LS-DYNA 命令.
Parameters
----------
input_file : str
输入文件路径
ncpu : int
CPU 核心数
Returns
-------
list[str]
LS-DYNA 命令列表
"""
if Constants.IS_WINDOWS or Constants.IS_MACOS:
return ["ls-dyna_mpp", f"i={input_file}", f"ncpu={ncpu}"]
else:
return ["ls-dyna_mpp", f"i={input_file}", f"ncpu={ncpu}"]
@px.register_fn
def run_ls_dyna(input_file: str, ncpu: int = DEFAULT_NCPU) -> None:
"""运行 LS-DYNA 计算.
Parameters
----------
input_file : str
输入文件路径
ncpu : int
CPU 核心数
"""
input_path = Path(input_file)
if not input_path.exists():
print(f"输入文件不存在: {input_path}")
return
cmd = get_ls_dyna_command(input_file, ncpu)
try:
subprocess.run(cmd, check=True)
print(f"LS-DYNA 计算完成: {input_file}")
except FileNotFoundError:
print("未找到 ls-dyna_mpp 命令")
except subprocess.CalledProcessError as e:
print(f"LS-DYNA 计算失败: {e}")
@px.register_fn
def run_ls_dyna_mpi(input_file: str, ncpu: int = DEFAULT_NCPU) -> None:
"""运行 LS-DYNA MPI 计算.
Parameters
----------
input_file : str
输入文件路径
ncpu : int
CPU 核心数
"""
input_path = Path(input_file)
if not input_path.exists():
print(f"输入文件不存在: {input_path}")
return
cmd = ["mpirun", "-np", str(ncpu), "ls-dyna_mpp", f"i={input_file}"]
try:
subprocess.run(cmd, check=True)
print(f"LS-DYNA MPI 计算完成: {input_file}")
except FileNotFoundError:
print("未找到 mpirun 或 ls-dyna_mpp 命令")
except subprocess.CalledProcessError as e:
print(f"LS-DYNA MPI 计算失败: {e}")
@px.register_fn
def check_ls_dyna_status() -> None:
"""检查 LS-DYNA 进程状态."""
try:
if Constants.IS_WINDOWS:
result = subprocess.run(
["tasklist", "/fi", "imagename eq ls-dyna_mpp.exe"],
capture_output=True,
text=True,
check=True,
)
print(result.stdout)
else:
result = subprocess.run(
["pgrep", "-f", "ls-dyna"],
capture_output=True,
text=True,
check=False,
)
if result.stdout.strip():
print(f"运行中的 LS-DYNA 进程 PID: {result.stdout.strip()}")
else:
print("没有运行中的 LS-DYNA 进程")
except subprocess.CalledProcessError as e:
print(f"检查进程状态失败: {e}")
# ============================================================================
# sshcopyid 函数
# ============================================================================
@px.register_fn
def ssh_copy_id(
hostname: str,
username: str,
password: str,
port: int = 22,
keypath: str = "~/.ssh/id_rsa.pub",
timeout: int = 30,
) -> None:
"""将 SSH 公钥部署到远程服务器.
Parameters
----------
hostname : str
远程服务器主机名或 IP 地址
username : str
远程服务器用户名
password : str
远程服务器密码
port : int
SSH 端口, 默认 22
keypath : str
公钥文件路径, 默认 ~/.ssh/id_rsa.pub
timeout : int
SSH 操作超时秒数, 默认 30
"""
pub_key_path = Path(keypath).expanduser()
if not pub_key_path.exists():
print(f"公钥文件不存在: {pub_key_path}")
sys.exit(1)
pub_key = pub_key_path.read_text().strip()
script = f"""mkdir -p ~/.ssh && chmod 700 ~/.ssh
cd ~/.ssh && touch authorized_keys && chmod 600 authorized_keys
grep -qF '{pub_key.split()[1]}' authorized_keys 2>/dev/null || echo '{pub_key}' >> authorized_keys"""
try:
subprocess.run(
[
"sshpass",
"-p",
password,
"ssh",
"-p",
str(port),
"-o",
"StrictHostKeyChecking=no",
"-o",
"UserKnownHostsFile=/dev/null",
"-o",
f"ConnectTimeout={timeout}",
f"{username}@{hostname}",
script,
],
check=True,
timeout=timeout,
)
print(f"SSH 密钥已部署到 {username}@{hostname}:{port}")
except FileNotFoundError:
print(f"未找到 sshpass 工具, 请手动执行: ssh-copy-id -p {port} {username}@{hostname}")
sys.exit(1)
except subprocess.TimeoutExpired:
print("SSH 连接超时")
sys.exit(1)
except subprocess.CalledProcessError as e:
print(f"SSH 执行失败: {e}")
sys.exit(1)
# ============================================================================
# packtool 函数
# ============================================================================
@px.register_fn
def pack_source(project_dir: Path, output_dir: Path) -> None:
"""打包项目源码.
Parameters
----------
project_dir : Path
项目目录
output_dir : Path
输出目录
"""
output_dir.mkdir(parents=True, exist_ok=True)
pyproject_file = project_dir / "pyproject.toml"
project_name = project_dir.name
if pyproject_file.exists():
try:
import tomllib
content = pyproject_file.read_text(encoding="utf-8")
data = tomllib.loads(content)
project_name = data.get("project", {}).get("name", project_name)
except ImportError:
pass
source_dir = output_dir / "src" / project_name
source_dir.mkdir(parents=True, exist_ok=True)
src_subdir = project_dir / "src"
if src_subdir.exists():
shutil.copytree(
src_subdir,
source_dir / "src",
ignore=shutil.ignore_patterns(*IGNORE_PATTERNS),
dirs_exist_ok=True,
)
else:
for item in project_dir.iterdir():
if item.name in IGNORE_PATTERNS or item.name.startswith("."):
continue
dst_item = source_dir / item.name
if item.is_dir():
shutil.copytree(
item,
dst_item,
ignore=shutil.ignore_patterns(*IGNORE_PATTERNS),
dirs_exist_ok=True,
)
else:
shutil.copy2(item, dst_item)
print(f"源码打包完成: {source_dir}")
@px.register_fn
def pack_dependencies(lib_dir: Path, dependencies: list[str]) -> None:
"""打包项目依赖.
Parameters
----------
lib_dir : Path
依赖库目录
dependencies : list[str]
依赖列表
"""
lib_dir.mkdir(parents=True, exist_ok=True)
if not dependencies:
print("没有依赖需要打包")
return
cmd = [
"pip",
"install",
"--target",
str(lib_dir),
"--no-compile",
"--no-warn-script-location",
]
cmd.extend(dependencies)
subprocess.run(cmd, check=True)
print(f"依赖打包完成: {lib_dir}")
@px.register_fn
def pack_wheel(project_dir: Path, output_dir: Path) -> None:
"""打包项目为 wheel 文件.
Parameters
----------
project_dir : Path
项目目录
output_dir : Path
输出目录
"""
output_dir.mkdir(parents=True, exist_ok=True)
cmd = [
"pip",
"wheel",
"--no-deps",
"--wheel-dir",
str(output_dir),
str(project_dir),
]
subprocess.run(cmd, check=True)
print(f"Wheel 打包完成: {output_dir}")
@px.register_fn
def install_embed_python(version: str, output_dir: Path) -> None:
"""安装嵌入式 Python.
Parameters
----------
version : str
Python 版本 (如: 3.10, 3.11)
output_dir : Path
输出目录
"""
output_dir.mkdir(parents=True, exist_ok=True)
arch = platform.machine().lower()
if arch in ["x86_64", "amd64"]:
arch = "amd64"
elif arch in ["arm64", "aarch64"]:
arch = "arm64"
version_map = {
"3.8": "3.8.10",
"3.9": "3.9.13",
"3.10": "3.10.11",
"3.11": "3.11.9",
"3.12": "3.12.4",
}
full_version = version_map.get(version, f"{version}.0")
url = f"https://www.python.org/ftp/python/{full_version}/python-{full_version}-embed-{arch}.zip"
cache_file = Path(DEFAULT_CACHE_DIR) / f"python-{full_version}-embed-{arch}.zip"
cache_file.parent.mkdir(parents=True, exist_ok=True)
if not cache_file.exists():
print(f"正在下载嵌入式 Python {full_version}...")
urllib.request.urlretrieve(url, cache_file)
print(f"下载完成: {cache_file}")
with zipfile.ZipFile(cache_file, "r") as zf:
zf.extractall(output_dir)
print(f"嵌入式 Python 安装完成: {output_dir}")
@px.register_fn
def create_zip_package(source_dir: Path, output_file: Path) -> None:
"""创建 ZIP 打包文件.
Parameters
----------
source_dir : Path
源目录
output_file : Path
输出文件
"""
output_file.parent.mkdir(parents=True, exist_ok=True)
with zipfile.ZipFile(output_file, "w", zipfile.ZIP_DEFLATED) as zf:
for file in source_dir.rglob("*"):
if file.is_file():
arcname = file.relative_to(source_dir)
zf.write(file, arcname)
print(f"ZIP 打包完成: {output_file}")
@px.register_fn
def clean_build_dir(build_dir: Path) -> None:
"""清理构建目录.
Parameters
----------
build_dir : Path
构建目录
"""
if build_dir.exists():
shutil.rmtree(build_dir)
print(f"清理完成: {build_dir}")
else:
print(f"目录不存在: {build_dir}")
+8 -2
View File
@@ -236,7 +236,11 @@ class Graph:
return graph
@classmethod
def from_yaml(cls, path: str | Path) -> Graph:
def from_yaml(
cls,
path: str | Path,
variables: Mapping[str, Any] | None = None,
) -> Graph:
"""从 YAML 文件构建任务图。
参考 GitHub Actions 风格 schema, 支持 jobs/needs/strategy.matrix/if
@@ -246,6 +250,8 @@ class Graph:
----------
path : str | Path
YAML 文件路径
variables : Mapping[str, Any] | None
运行时变量, 用于替换 ``${VAR}`` 占位符
Returns
-------
@@ -259,7 +265,7 @@ class Graph:
"""
from .yaml_loader import load_yaml
return load_yaml(path)
return load_yaml(path, variables=variables)
def add_subgraph(self, sub: Graph, *, namespace: str | None = None) -> Graph:
"""将子图合并到当前图,任务名加命名空间前缀避免冲突。
+159
View File
@@ -0,0 +1,159 @@
"""函数注册表.
提供全局函数注册机制, 供 YAML 任务编排通过 ``fn`` 字段引用 Python 函数.
使用方式
--------
import pyflowx as px
@px.register_fn("pack_source")
def pack_source(project_dir, output_dir):
...
# YAML 中引用:
# jobs:
# pack:
# fn: pack_source
# args: ["./project", "./dist"]
"""
from __future__ import annotations
import sys
from typing import Any, Callable, TypeVar, overload
if sys.version_info >= (3, 10):
from typing import ParamSpec
else:
from typing_extensions import ParamSpec # pragma: no cover
__all__ = ["FnRegistry", "get_fn", "has_fn", "register_fn"]
P = ParamSpec("P")
T = TypeVar("T")
_REGISTRY: dict[str, Callable[..., Any]] = {}
@overload
def register_fn(name: Callable[P, T]) -> Callable[P, T]: ...
@overload
def register_fn(name: str | None = None) -> Callable[[Callable[P, T]], Callable[P, T]]: ...
def register_fn(name: str | Callable[..., Any] | None = None) -> Callable[..., Any]:
"""装饰器:将函数注册到全局 registry.
支持两种用法::
@register_fn # 使用函数 __name__ 作为注册名
def my_func(): ...
@register_fn("custom") # 显式指定注册名
def my_func(): ...
Parameters
----------
name : str | Callable | None
注册名或被装饰函数; 为 None 时使用函数 ``__name__``
Returns
-------
Callable
装饰器函数或被装饰函数
Raises
------
ValueError
名称已注册或无法推断函数名
"""
if callable(name):
fn = name
key = getattr(fn, "__name__", None)
if key is None:
raise ValueError("无法推断函数名, 请显式提供 name 参数")
if key in _REGISTRY:
raise ValueError(f"函数 {key!r} 已注册")
_REGISTRY[key] = fn
return fn
def decorator(fn: Callable[P, T]) -> Callable[P, T]:
key = name if name is not None else getattr(fn, "__name__", None)
if key is None:
raise ValueError("无法推断函数名, 请显式提供 name 参数")
if key in _REGISTRY:
raise ValueError(f"函数 {key!r} 已注册")
_REGISTRY[key] = fn
return fn
return decorator
def get_fn(name: str) -> Callable[..., Any]:
"""按名称获取已注册的函数.
Parameters
----------
name : str
函数名
Returns
-------
Callable
已注册的函数
Raises
------
KeyError
函数未注册
"""
if name not in _REGISTRY:
raise KeyError(f"函数 {name!r} 未注册")
return _REGISTRY[name]
def has_fn(name: str) -> bool:
"""检查函数是否已注册.
Parameters
----------
name : str
函数名
Returns
-------
bool
是否已注册
"""
return name in _REGISTRY
class FnRegistry:
"""函数注册表的面向对象访问接口."""
@staticmethod
def register(name: str | None = None) -> Callable[[Callable[..., T]], Callable[..., T]]:
"""注册装饰器, 等价于 :func:`register_fn`."""
return register_fn(name)
@staticmethod
def get(name: str) -> Callable[..., Any]:
"""获取已注册函数, 等价于 :func:`get_fn`."""
return get_fn(name)
@staticmethod
def has(name: str) -> bool:
"""检查是否已注册, 等价于 :func:`has_fn`."""
return has_fn(name)
@staticmethod
def clear() -> None:
"""清空注册表."""
_REGISTRY.clear()
@staticmethod
def names() -> list[str]:
"""返回所有已注册函数名."""
return list(_REGISTRY.keys())
+177 -42
View File
@@ -43,6 +43,7 @@ Schema 概览
字段映射
-------
- ``cmd`` / ``run``: ``cmd`` 为命令列表, ``run`` 为 shell 字符串
- ``fn``: 引用 ``register_fn`` 注册的 Python 函数名, 配合 ``args``/``kwargs`` 传参
- ``needs``: 对应 ``depends_on``
- ``if``: 条件表达式, 支持 ``success()`` / ``always()`` / ``env.VAR`` / ``env.VAR == 'x'``
- ``strategy.matrix``: 笛卡尔积展开为多个任务, 任务名追加 ``_{key}-{val}`` 后缀
@@ -51,7 +52,7 @@ Schema 概览
矩阵占位符
----------
矩阵值通过 ``${{ matrix.key }}`` 占位符注入 ``cmd`` / ``run`` / ``cwd`` / ``env`` 中::
矩阵值通过 ``${{ matrix.key }}`` 占位符注入 ``cmd`` / ``run`` / ``cwd`` / ``env`` / ``args`` / ``kwargs`` 中::
jobs:
test:
@@ -59,6 +60,36 @@ Schema 概览
strategy:
matrix:
version: ["3.8", "3.9"]
变量占位符
----------
运行时变量通过 ``${VAR}`` 占位符注入, 变量由 ``load_yaml(path, variables={...})``
或 ``parse_yaml_string(content, variables={...})`` 传入::
# Python
graph = px.load_yaml("pipeline.yaml", variables={"PROJECT_DIR": "./project"})
# YAML
jobs:
pack:
fn: pack_source
args: ["${PROJECT_DIR}", "./dist"]
函数引用
----------
``fn`` 字段引用 ``@px.register_fn`` 注册的函数, 配合 ``args``/``kwargs`` 传参::
# Python
@px.register_fn("pack_source")
def pack_source(project_dir, output_dir):
...
# YAML
jobs:
pack:
fn: pack_source
args: ["./project", "./dist"]
needs: [clean]
"""
from __future__ import annotations
@@ -72,11 +103,15 @@ from typing import Any, Mapping
import pyflowx as px
from pyflowx.conditions import BuiltinConditions
from pyflowx.registry import get_fn as _get_fn
from pyflowx.task import Condition, RetryPolicy, TaskSpec
# 矩阵占位符: ${{ matrix.key }}
_MATRIX_PLACEHOLDER = re.compile(r"\$\{\{\s*matrix\.(\w+)\s*\}\}")
# 变量占位符: ${VAR}
_VARIABLE_PLACEHOLDER = re.compile(r"\$\{(\w+)\}")
# 条件表达式
_ENV_VAR_EXISTS_RE = re.compile(r"^env\.(\w+)$")
_ENV_VAR_EQUALS_RE = re.compile(r"^env\.(\w+)\s*==\s*['\"]([^'\"]*)['\"]$")
@@ -96,13 +131,15 @@ class YamlLoadError(px.PyFlowXError):
"""YAML 加载错误."""
def load_yaml(path: str | Path) -> px.Graph:
def load_yaml(path: str | Path, variables: Mapping[str, Any] | None = None) -> px.Graph:
"""从 YAML 文件加载任务图.
Parameters
----------
path : str | Path
YAML 文件路径
variables : Mapping[str, Any] | None
运行时变量, 用于替换 ``${VAR}`` 占位符
Returns
-------
@@ -121,16 +158,18 @@ def load_yaml(path: str | Path) -> px.Graph:
content = p.read_text(encoding="utf-8")
except OSError as e:
raise YamlLoadError(f"读取文件失败: {e}") from e
return parse_yaml_string(content)
return parse_yaml_string(content, variables=variables)
def parse_yaml_string(content: str) -> px.Graph:
def parse_yaml_string(content: str, variables: Mapping[str, Any] | None = None) -> px.Graph:
"""从 YAML 字符串构建任务图.
Parameters
----------
content : str
YAML 文本
variables : Mapping[str, Any] | None
运行时变量, 用于替换 ``${VAR}`` 占位符
Returns
-------
@@ -140,7 +179,7 @@ def parse_yaml_string(content: str) -> px.Graph:
data = _parse_yaml(content)
if not isinstance(data, Mapping):
raise YamlLoadError("YAML 根节点必须是 mapping (对象)")
return _build_graph(data)
return _build_graph(data, variables=variables)
def _parse_yaml(content: str) -> Any:
@@ -156,7 +195,7 @@ def _parse_yaml(content: str) -> Any:
raise YamlLoadError(f"YAML 解析失败: {e}") from e
def _build_graph(data: Mapping[str, Any]) -> px.Graph:
def _build_graph(data: Mapping[str, Any], variables: Mapping[str, Any] | None = None) -> px.Graph:
"""从解析后的 mapping 构建 Graph."""
jobs = data.get("jobs")
if not isinstance(jobs, Mapping):
@@ -202,7 +241,7 @@ def _build_graph(data: Mapping[str, Any]) -> px.Graph:
specs: list[TaskSpec[Any]] = []
seen_names: set[str] = set()
for job_id, job_config in jobs.items():
for spec in _parse_job(job_id, job_config, job_expansions):
for spec in _parse_job(job_id, job_config, job_expansions, variables=variables):
if spec.name in seen_names:
raise YamlLoadError(f"重复任务名: {spec.name!r}")
seen_names.add(spec.name)
@@ -235,6 +274,7 @@ def _parse_job(
job_id: str,
config: Mapping[str, Any],
job_expansions: Mapping[str, list[str]],
variables: Mapping[str, Any] | None = None,
) -> list[TaskSpec[Any]]:
"""解析单个 job, 返回 spec 列表 (矩阵展开可能产生多个).
@@ -246,28 +286,16 @@ def _parse_job(
job 配置
job_expansions : Mapping[str, list[str]]
每个 job_id 到其展开名列表的映射, 用于展开矩阵依赖
variables : Mapping[str, Any] | None
运行时变量, 用于替换 ``${VAR}`` 占位符
"""
cmd_data = config.get("cmd")
run_data = config.get("run")
if cmd_data is None and run_data is None:
raise YamlLoadError(f"job {job_id!r} 必须提供 cmd 或 run 字段")
if cmd_data is not None and run_data is not None:
raise YamlLoadError(f"job {job_id!r} 不能同时提供 cmd 和 run")
cmd_data, run_data, fn_data, args_data, kwargs_data = _parse_task_type_fields(job_id, config)
matrix_combos = _expand_matrix(config.get("strategy"))
if not matrix_combos:
matrix_combos = [None]
# 展开矩阵依赖: needs 中的 job_id 替换为所有展开名
needs_raw = _as_str_tuple(config.get("needs"))
needs: list[str] = []
for dep in needs_raw:
if dep in job_expansions:
needs.extend(job_expansions[dep])
else:
needs.append(dep)
needs_tuple = tuple(needs)
needs_tuple = _expand_needs(_as_str_tuple(config.get("needs")), job_expansions)
if_data = config.get("if")
timeout = _as_float(config.get("timeout"))
@@ -285,17 +313,26 @@ def _parse_job(
if runs_on:
tags = (*tags, runs_on)
fn = _resolve_fn(job_id, fn_data)
specs: list[TaskSpec[Any]] = []
for combo in matrix_combos:
name = _matrix_name(job_id, combo)
cmd = _build_cmd(cmd_data, run_data, combo)
cwd_resolved = Path(_replace_in_str(str(cwd), combo)) if cwd is not None else None
env_resolved = _replace_in_mapping(env, combo) if env else None
cmd = _build_cmd(cmd_data, run_data, combo, variables) if fn_data is None else None
cwd_resolved = Path(_replace_in_str(str(cwd), combo, variables)) if cwd is not None else None
env_resolved = _replace_in_mapping(env, combo, variables) if env else None
conditions = _parse_condition(if_data) if if_data is not None else ()
# fn 任务的 args/kwargs 应用占位符替换
args_resolved = _replace_in_value(args_data, combo, variables) if args_data is not None else ()
kwargs_resolved = _replace_in_value(kwargs_data, combo, variables) if kwargs_data is not None else {}
spec = px.TaskSpec(
name=name,
fn=fn,
cmd=cmd,
args=tuple(args_resolved),
kwargs=dict(kwargs_resolved),
depends_on=needs_tuple,
retry=retry if retry is not None else RetryPolicy(),
timeout=timeout,
@@ -314,6 +351,67 @@ def _parse_job(
return specs
def _parse_task_type_fields(
job_id: str,
config: Mapping[str, Any],
) -> tuple[Any, Any, Any, Any, Any]:
"""解析并校验任务类型字段 (cmd/run/fn) 及参数 (args/kwargs).
Returns
-------
tuple
(cmd_data, run_data, fn_data, args_data, kwargs_data)
"""
cmd_data = config.get("cmd")
run_data = config.get("run")
fn_data = config.get("fn")
args_data = config.get("args")
kwargs_data = config.get("kwargs")
# 校验互斥: cmd/run/fn 三选一
provided = [x for x in (cmd_data, run_data, fn_data) if x is not None]
if not provided:
raise YamlLoadError(f"job {job_id!r} 必须提供 cmd/run/fn 之一")
if len(provided) > 1:
raise YamlLoadError(f"job {job_id!r} 不能同时提供 cmd/run/fn")
# fn 必须是字符串
if fn_data is not None and not isinstance(fn_data, str):
raise YamlLoadError(f"job {job_id!r} 的 fn 必须是字符串, 收到: {type(fn_data).__name__}")
# args 必须是列表, kwargs 必须是 mapping
if args_data is not None and not isinstance(args_data, list):
raise YamlLoadError(f"job {job_id!r} 的 args 必须是列表, 收到: {type(args_data).__name__}")
if kwargs_data is not None and not isinstance(kwargs_data, Mapping):
raise YamlLoadError(f"job {job_id!r} 的 kwargs 必须是 mapping, 收到: {type(kwargs_data).__name__}")
return cmd_data, run_data, fn_data, args_data, kwargs_data
def _expand_needs(
needs_raw: tuple[str, ...],
job_expansions: Mapping[str, list[str]],
) -> tuple[str, ...]:
"""展开 needs 中的矩阵依赖."""
needs: list[str] = []
for dep in needs_raw:
if dep in job_expansions:
needs.extend(job_expansions[dep])
else:
needs.append(dep)
return tuple(needs)
def _resolve_fn(job_id: str, fn_data: Any) -> Any:
"""从 registry 查找已注册的函数."""
if fn_data is None:
return None
try:
return _get_fn(fn_data)
except KeyError as e:
raise YamlLoadError(f"job {job_id!r} 引用的函数 {fn_data!r} 未注册") from e
def _expand_matrix(strategy_data: Any) -> list[Mapping[str, Any] | None]:
"""展开 strategy.matrix 为笛卡尔积列表.
@@ -373,35 +471,72 @@ def _build_cmd(
cmd_data: Any,
run_data: Any,
combo: Mapping[str, Any] | None,
variables: Mapping[str, Any] | None = None,
) -> list[str] | str:
"""构建 cmd 字段, 应用矩阵占位符替换."""
"""构建 cmd 字段, 应用占位符替换."""
if cmd_data is not None:
if isinstance(cmd_data, str):
return _replace_in_str(cmd_data, combo)
return _replace_in_str(cmd_data, combo, variables)
if isinstance(cmd_data, list):
return [_replace_in_str(str(item), combo) for item in cmd_data]
return [_replace_in_str(str(item), combo, variables) for item in cmd_data]
raise YamlLoadError(f"cmd 必须是 list 或 str, 收到: {type(cmd_data).__name__}")
# run_data 不为 None (已在外层校验)
return _replace_in_str(str(run_data), combo)
return _replace_in_str(str(run_data), combo, variables)
def _replace_in_str(text: str, combo: Mapping[str, Any] | None) -> str:
"""替换字符串中的 ${{ matrix.key }} 占位符."""
if combo is None:
return text
def _replace_in_str(
text: str,
combo: Mapping[str, Any] | None,
variables: Mapping[str, Any] | None = None,
) -> str:
"""替换字符串中的 ``${VAR}`` 变量占位符和 ``${{ matrix.key }}`` 矩阵占位符."""
# 先替换变量占位符 ${VAR}
if variables:
def _sub(match: re.Match[str]) -> str:
key = match.group(1)
if key not in combo:
raise YamlLoadError(f"矩阵占位符 matrix.{key} 未定义")
return str(combo[key])
def _sub_var(match: re.Match[str]) -> str:
key = match.group(1)
if key not in variables:
raise YamlLoadError(f"变量 {key!r} 未定义")
return str(variables[key])
return _MATRIX_PLACEHOLDER.sub(_sub, text)
text = _VARIABLE_PLACEHOLDER.sub(_sub_var, text)
# 再替换矩阵占位符 ${{ matrix.key }}
if combo is not None:
def _sub_matrix(match: re.Match[str]) -> str:
key = match.group(1)
if key not in combo:
raise YamlLoadError(f"矩阵占位符 matrix.{key} 未定义")
return str(combo[key])
text = _MATRIX_PLACEHOLDER.sub(_sub_matrix, text)
return text
def _replace_in_mapping(env: Mapping[str, str], combo: Mapping[str, Any] | None) -> dict[str, str]:
def _replace_in_mapping(
env: Mapping[str, str],
combo: Mapping[str, Any] | None,
variables: Mapping[str, Any] | None = None,
) -> dict[str, str]:
"""替换 env 映射中的占位符."""
return {k: _replace_in_str(v, combo) for k, v in env.items()}
return {k: _replace_in_str(v, combo, variables) for k, v in env.items()}
def _replace_in_value(
value: Any,
combo: Mapping[str, Any] | None,
variables: Mapping[str, Any] | None = None,
) -> Any:
"""递归替换值 (list/mapping/scalar) 中的占位符."""
if isinstance(value, str):
return _replace_in_str(value, combo, variables)
if isinstance(value, list):
return [_replace_in_value(v, combo, variables) for v in value]
if isinstance(value, Mapping):
return {k: _replace_in_value(v, combo, variables) for k, v in value.items()}
return value
def _parse_condition(expr: Any) -> tuple[Condition, ...]:
+247
View File
@@ -0,0 +1,247 @@
"""``pyflowx.registry`` 模块的单元测试."""
from __future__ import annotations
from typing import Any
import pytest
import pyflowx as px
from pyflowx.registry import _REGISTRY, FnRegistry, get_fn, has_fn, register_fn
@pytest.fixture
def clear_registry() -> Any:
"""每个测试前后清空 registry, 避免测试间相互污染."""
_REGISTRY.clear()
yield
_REGISTRY.clear()
class TestRegisterFn:
"""``register_fn`` 装饰器测试."""
@pytest.fixture(autouse=True)
def _clear(self, clear_registry: Any) -> None:
"""每个测试前后清空 registry."""
def test_register_with_explicit_name(self) -> None:
"""使用显式名称注册函数."""
@register_fn("custom_name")
def _func(x: int) -> int:
return x * 2
assert has_fn("custom_name")
assert get_fn("custom_name")(5) == 10
def test_register_without_parentheses(self) -> None:
"""无括号使用 ``@register_fn`` 时以 ``__name__`` 注册."""
@register_fn
def my_func(x: int) -> int:
return x + 1
assert has_fn("my_func")
assert get_fn("my_func")(10) == 11
def test_register_with_none_uses_func_name(self) -> None:
"""``@register_fn(None)`` 等价于使用 ``__name__``."""
@register_fn(None)
def auto_named(x: int) -> int:
return x
assert has_fn("auto_named")
def test_register_preserves_function(self) -> None:
"""装饰后函数仍可直接调用."""
@register_fn("calc")
def _calc(a: int, b: int) -> int:
return a + b
assert _calc(2, 3) == 5
def test_duplicate_registration_raises(self) -> None:
"""重复注册同名函数抛出 ValueError."""
with pytest.raises(ValueError, match="已注册"):
@register_fn("dup")
def _first() -> None:
pass
@register_fn("dup")
def _second() -> None:
pass
def test_duplicate_registration_without_parentheses(self) -> None:
"""无括号模式下重复注册也抛出 ValueError."""
with pytest.raises(ValueError, match="已注册"):
@register_fn
def shared_name() -> None:
pass
@register_fn
def shared_name() -> None: # noqa: F811
pass
class TestGetFn:
"""``get_fn`` 函数测试."""
@pytest.fixture(autouse=True)
def _clear(self, clear_registry: Any) -> None:
"""每个测试前后清空 registry."""
def test_get_registered_function(self) -> None:
"""获取已注册的函数."""
@register_fn("target")
def _target(x: int) -> int:
return x * 3
retrieved = get_fn("target")
assert retrieved is _target
def test_get_unregistered_raises_keyerror(self) -> None:
"""获取未注册的函数抛出 KeyError."""
with pytest.raises(KeyError, match="not_found"):
get_fn("not_found")
class TestHasFn:
"""``has_fn`` 函数测试."""
@pytest.fixture(autouse=True)
def _clear(self, clear_registry: Any) -> None:
"""每个测试前后清空 registry."""
def test_has_registered(self) -> None:
"""已注册返回 True."""
@register_fn("exists")
def _exists() -> None:
pass
assert has_fn("exists") is True
def test_has_not_registered(self) -> None:
"""未注册返回 False."""
assert has_fn("nope") is False
class TestFnRegistry:
"""``FnRegistry`` 类测试."""
@pytest.fixture(autouse=True)
def _clear(self, clear_registry: Any) -> None:
"""每个测试前后清空 registry."""
def test_register_method(self) -> None:
"""``FnRegistry.register`` 等价于 ``register_fn``."""
@FnRegistry.register("via_class")
def _func(x: int) -> int:
return x
assert FnRegistry.has("via_class")
assert FnRegistry.get("via_class")(7) == 7
def test_register_method_without_name(self) -> None:
"""``FnRegistry.register(None)`` 使用函数名."""
@FnRegistry.register()
def class_auto(x: int) -> int:
return x
assert FnRegistry.has("class_auto")
def test_get_method(self) -> None:
"""``FnRegistry.get`` 获取已注册函数."""
@register_fn("klass")
def _klass() -> str:
return "hello"
assert FnRegistry.get("klass")() == "hello"
def test_has_method_false(self) -> None:
"""``FnRegistry.has`` 未注册返回 False."""
assert FnRegistry.has("missing") is False
def test_clear_method(self) -> None:
"""``FnRegistry.clear`` 清空注册表."""
@register_fn("temp")
def _temp() -> None:
pass
assert FnRegistry.has("temp")
FnRegistry.clear()
assert not FnRegistry.has("temp")
def test_names_method(self) -> None:
"""``FnRegistry.names`` 返回所有注册名."""
@register_fn("alpha")
def _alpha() -> None:
pass
@register_fn("beta")
def _beta() -> None:
pass
names = FnRegistry.names()
assert "alpha" in names
assert "beta" in names
def test_names_returns_copy(self) -> None:
"""``names`` 返回列表副本, 修改不影响内部状态."""
@register_fn("orig")
def _orig() -> None:
pass
names = FnRegistry.names()
names.append("fake")
assert not FnRegistry.has("fake")
class TestOpsModules:
"""``_ops`` 子模块导入后函数自动注册的集成测试.
这些测试不使用 ``_clear_registry`` fixture, 因为 ``_ops`` 模块在首次导入时
注册函数, 清空后将无法重新注册 (模块不会再次执行).
"""
def test_all_ops_functions_registered(self) -> None:
"""导入 ``_ops`` 后所有函数应已注册."""
import inspect
from pyflowx.cli._ops import dev, files, media, system
for module in (files, dev, media, system):
for name in module.__all__:
obj = getattr(module, name)
if not inspect.isfunction(obj):
continue
assert px.has_fn(name), f"{module.__name__}.{name} 未注册"
def test_total_function_count(self) -> None:
"""注册函数总数 = 15+15+18+11 = 59."""
from pyflowx.cli._ops import dev, files, media, system # noqa: F401
all_names = px.FnRegistry.names()
assert len(all_names) == 59
def test_specific_functions_callable(self) -> None:
"""关键注册函数可调用."""
from pyflowx.cli._ops import dev, files, media, system
assert px.get_fn("get_file_timestamp") is files.get_file_timestamp
assert px.get_fn("bump_file_version") is dev.bump_file_version
assert px.get_fn("pdf_merge") is media.pdf_merge
assert px.get_fn("ssh_copy_id") is system.ssh_copy_id
+312 -2
View File
@@ -174,7 +174,7 @@ jobs:
def test_cmd_and_run_conflict(self) -> None:
"""同时提供 cmd 和 run 应报错."""
with pytest.raises(YamlLoadError, match="不能同时提供 cmd 和 run"):
with pytest.raises(YamlLoadError, match="不能同时提供 cmd/run/fn"):
parse_yaml_string(
"""
jobs:
@@ -186,7 +186,7 @@ jobs:
def test_missing_cmd_and_run(self) -> None:
"""缺少 cmd 和 run 应报错."""
with pytest.raises(YamlLoadError, match="必须提供 cmd 或 run"):
with pytest.raises(YamlLoadError, match="必须提供 cmd/run/fn 之一"):
parse_yaml_string(
"""
jobs:
@@ -1264,3 +1264,313 @@ jobs:
report = px.run(graph, strategy="sequential")
# skip_me 被跳过, downstream 因 allow-upstream-skip 仍执行
assert report.result_of("skip_me").status == TaskStatus.SKIPPED
# ---------------------------------------------------------------------- #
# fn 字段 + 变量占位符
# ---------------------------------------------------------------------- #
class TestFnField:
"""测试 fn 字段引用注册函数."""
def setup_method(self) -> None:
"""每个测试前清空 registry."""
px.FnRegistry.clear()
def teardown_method(self) -> None:
"""每个测试后清空 registry."""
px.FnRegistry.clear()
def test_fn_basic(self) -> None:
"""fn 引用注册函数."""
@px.register_fn("greet")
def _greet(name: str) -> str:
return f"hello {name}"
graph = parse_yaml_string(
"""
jobs:
greet:
fn: greet
args: ["world"]
""",
)
report = px.run(graph, strategy="sequential")
assert report.success
assert report["greet"] == "hello world"
def test_fn_with_kwargs(self) -> None:
"""fn 使用 kwargs 传参."""
@px.register_fn("add")
def _add(a: int, b: int) -> int:
return a + b
graph = parse_yaml_string(
"""
jobs:
add:
fn: add
args: [1, 2]
""",
)
report = px.run(graph, strategy="sequential")
assert report["add"] == 3
def test_fn_not_registered(self) -> None:
"""fn 引用未注册的函数应报错."""
with pytest.raises(YamlLoadError, match="未注册"):
parse_yaml_string(
"""
jobs:
t:
fn: nonexistent_fn
""",
)
def test_fn_not_string(self) -> None:
"""fn 不是字符串应报错."""
with pytest.raises(YamlLoadError, match="fn 必须是字符串"):
parse_yaml_string(
"""
jobs:
t:
fn: 123
""",
)
def test_fn_args_not_list(self) -> None:
"""args 不是列表应报错."""
with pytest.raises(YamlLoadError, match="args 必须是列表"):
parse_yaml_string(
"""
jobs:
t:
fn: some_fn
args: "not a list"
""",
)
def test_fn_kwargs_not_mapping(self) -> None:
"""kwargs 不是 mapping 应报错."""
with pytest.raises(YamlLoadError, match="kwargs 必须是 mapping"):
parse_yaml_string(
"""
jobs:
t:
fn: some_fn
kwargs: [1, 2]
""",
)
def test_fn_with_cmd_conflict(self) -> None:
"""fn 与 cmd 同时提供应报错."""
with pytest.raises(YamlLoadError, match="不能同时提供"):
parse_yaml_string(
"""
jobs:
t:
fn: some_fn
cmd: ["echo"]
""",
)
def test_fn_with_needs(self) -> None:
"""fn 任务可以有 needs 依赖."""
@px.register_fn("upper")
def _upper(s: str) -> str:
return s.upper()
graph = parse_yaml_string(
"""
jobs:
setup:
cmd: ["echo", "hello"]
process:
fn: upper
args: ["data"]
needs: [setup]
""",
)
assert graph.spec("process").depends_on == ("setup",)
report = px.run(graph, strategy="sequential")
assert report["process"] == "DATA"
def test_fn_matrix_placeholder_in_args(self) -> None:
"""fn 的 args 支持矩阵占位符."""
@px.register_fn("echo_val")
def _echo_val(val: str) -> str:
return val
graph = parse_yaml_string(
"""
jobs:
test:
fn: echo_val
args: ["${{ matrix.mode }}"]
strategy:
matrix:
mode: [release, debug]
""",
)
assert sorted(graph.names) == ["test_mode-debug", "test_mode-release"]
report = px.run(graph, strategy="sequential")
assert report["test_mode-release"] == "release"
assert report["test_mode-debug"] == "debug"
# ---------------------------------------------------------------------- #
# 变量占位符 ${VAR}
# ---------------------------------------------------------------------- #
class TestVariablePlaceholder:
"""测试 ${VAR} 变量占位符."""
def test_variable_in_cmd(self) -> None:
"""变量在 cmd 中替换."""
graph = parse_yaml_string(
"""
jobs:
echo:
cmd: ["echo", "${MSG}"]
""",
variables={"MSG": "hello"},
)
assert graph.spec("echo").cmd == ["echo", "hello"]
def test_variable_in_run(self) -> None:
"""变量在 run 中替换."""
graph = parse_yaml_string(
"""
jobs:
t:
run: "echo ${MSG} && exit 0"
""",
variables={"MSG": "hi"},
)
assert graph.spec("t").cmd == "echo hi && exit 0"
def test_variable_in_cwd(self) -> None:
"""变量在 cwd 中替换."""
graph = parse_yaml_string(
"""
jobs:
t:
cmd: ["pwd"]
cwd: "${BASE_DIR}/sub"
""",
variables={"BASE_DIR": "/tmp"},
)
spec = graph.spec("t")
assert spec.cwd is not None
assert "/tmp" in str(spec.cwd)
def test_variable_in_env(self) -> None:
"""变量在 env 中替换."""
graph = parse_yaml_string(
"""
jobs:
t:
cmd: ["echo"]
env:
PATH: "${BIN_PATH}"
""",
variables={"BIN_PATH": "/usr/local/bin"},
)
assert graph.spec("t").env == {"PATH": "/usr/local/bin"}
def test_variable_in_fn_args(self) -> None:
"""变量在 fn 的 args 中替换."""
@px.register_fn("concat")
def _concat(a: str, b: str) -> str:
return a + b
graph = parse_yaml_string(
"""
jobs:
t:
fn: concat
args: ["${A}", "${B}"]
""",
variables={"A": "foo", "B": "bar"},
)
report = px.run(graph, strategy="sequential")
assert report["t"] == "foobar"
def test_variable_undefined(self) -> None:
"""引用未定义变量应报错."""
with pytest.raises(YamlLoadError, match="变量 'UNDEFINED' 未定义"):
parse_yaml_string(
"""
jobs:
t:
cmd: ["echo", "${UNDEFINED}"]
""",
variables={"OTHER": "x"},
)
def test_multiple_variables(self) -> None:
"""多个变量同时替换."""
graph = parse_yaml_string(
"""
jobs:
t:
cmd: ["echo", "${A}", "${B}", "${A}"]
""",
variables={"A": "x", "B": "y"},
)
assert graph.spec("t").cmd == ["echo", "x", "y", "x"]
def test_variable_no_variables_provided(self) -> None:
"""未提供 variables 时, ${VAR} 原样保留."""
graph = parse_yaml_string(
"""
jobs:
t:
cmd: ["echo", "${VAR}"]
""",
)
# 没有 variables, 占位符保留
assert graph.spec("t").cmd == ["echo", "${VAR}"]
def test_variable_and_matrix_combined(self) -> None:
"""变量和矩阵占位符同时替换."""
graph = parse_yaml_string(
"""
jobs:
test:
cmd: ["echo", "${PREFIX}", "${{ matrix.mode }}"]
strategy:
matrix:
mode: [release, debug]
""",
variables={"PREFIX": "build"},
)
for name in graph.names:
spec = graph.spec(name)
cmd = spec.cmd
assert isinstance(cmd, list)
assert cmd[1] == "build"
assert cmd[2] in ("release", "debug")
def test_load_yaml_with_variables(self, tmp_path: Path) -> None:
"""load_yaml 支持 variables 参数."""
yaml_file = tmp_path / "pipeline.yaml"
yaml_file.write_text(
'jobs:\n t:\n cmd: ["echo", "${MSG}"]\n',
encoding="utf-8",
)
graph = load_yaml(yaml_file, variables={"MSG": "loaded"})
assert graph.spec("t").cmd == ["echo", "loaded"]
def test_graph_from_yaml_with_variables(self, tmp_path: Path) -> None:
"""Graph.from_yaml 支持 variables 参数."""
yaml_file = tmp_path / "pipeline.yaml"
yaml_file.write_text(
'jobs:\n t:\n cmd: ["echo", "${MSG}"]\n',
encoding="utf-8",
)
graph = px.Graph.from_yaml(yaml_file, variables={"MSG": "from_class"})
assert graph.spec("t").cmd == ["echo", "from_class"]