iamt 1.1.27__py3-none-any.whl

iamt/task/task.py

@@ -0,0 +1,108 @@
+import inspect
+import importlib
+from pathlib import Path
+
+from fabric import Task as FabricTask
+
+
+class Task:
+
+    def _scan_tasks(self) -> dict[str, dict]:
+        """扫描 modules 目录下所有子目录中的 task"""
+        tasks = {}
+        modules_path = Path(__file__).parent.parent / "modules"
+
+        # region 遍历 modules 目录下的非下划线开头的子目录
+        for module_dir in modules_path.iterdir():
+            if not module_dir.is_dir() or module_dir.name.startswith("_"):
+                continue
+
+            # region 扫描子目录下不以下划线开头的 Python 文件
+            for py_file in module_dir.glob("*.py"):
+                if py_file.name.startswith("_"):
+                    continue
+
+                module_name = py_file.stem
+                import_path = f".modules.{module_dir.name}.{module_name}"
+
+                try:
+                    module = importlib.import_module(import_path, package="iamt")
+                    source_file = py_file.name
+
+                    for name, obj in vars(module).items():
+                        if isinstance(obj, FabricTask):
+                            tasks[name] = {
+                                "task": obj,
+                                "source": source_file,
+                                "module": module_dir
+                            }
+                except ImportError as e:
+                    print(f"警告: 无法导入模块 {import_path}: {e}")
+            # endregion
+        # endregion
+
+        return tasks
+
+    def ls(self):
+        """列出所有任务"""
+        print("可用任务:")
+        for name, task_info in (self._scan_tasks()).items():
+            task_obj = task_info["task"]
+            sig = inspect.signature(task_obj.body)
+            params = [p for p in sig.parameters.keys() if p != 'ctx']
+            params_str = f"({', '.join(params)})" if params else "()"
+            doc = (task_obj.body.__doc__ or "").strip().split('\n')[0]
+            location = f"@ {task_info['module'].name}/{task_info['source']}"
+            print(f"  {name}{params_str}: {doc} {location}")
+
+    def run(self):
+        """交互式选择并运行任务"""
+        from prompt_toolkit import prompt
+        from ..completers.customCompleter import CustomCompleter, CustomValidator
+        from ..client_config.client_config import ClientConfig
+
+        if not self._scan_tasks():
+            print("没有可用的任务")
+            return
+
+        # region 构建 模块名.task函数名 格式的选项列表
+        task_options = {
+            f"{task_info['module'].name}.{name}": f"{task_info['module'].name}.{name}"
+            for name, task_info in (self._scan_tasks()).items()
+        }
+        completer = CustomCompleter(task_options)
+        validator = CustomValidator(completer, error_msg="无效的任务，请从补全列表中选择")
+        # endregion
+
+        # region 模糊搜索选择
+        try:
+            selected_key = prompt(
+                "请选择要运行的任务 (Tab补全, 支持模糊搜索): ",
+                completer=completer,
+                validator=validator
+            ).strip()
+        except (KeyboardInterrupt, EOFError):
+            print("\n已取消选择")
+            return
+
+        if not selected_key or selected_key not in task_options:
+            print(f"无效的选择: {selected_key}")
+            return
+        # endregion
+
+        # region 执行选中的任务
+        task_name = selected_key.split(".")[-1]
+        task_info = (self._scan_tasks())[task_name]
+        task_obj = task_info["task"]
+        config = ClientConfig()
+        conn = config.connect()
+        if conn is None:
+            return
+        doc = (task_obj.body.__doc__ or "").strip()
+        if doc:
+            print(f"\n{doc}\n")
+        try:
+            task_obj(conn)
+        except Exception as e:
+            print(f"任务 '{task_name}' 执行失败: {e}")
+        # endregion

iamt/utils/__init__.py

@@ -0,0 +1,6 @@
+"""工具模块"""
+
+from .config_saver import save_task_config, TaskCall, get_hostname_from_iamt_config
+from .task_tracker import TaskTracker
+
+__all__ = ["save_task_config", "TaskCall", "get_hostname_from_iamt_config", "TaskTracker"]

iamt/utils/config_saver.py

@@ -0,0 +1,90 @@
+"""配置保存工具模块
+
+提供统一的 task 配置保存功能，供 DCD 和 role 执行时复用。
+"""
+
+from pathlib import Path
+from typing import TypedDict
+
+import yaml
+from rich.console import Console
+
+
+class TaskCall(TypedDict):
+    """Task 函数调用记录"""
+    task: str             # task 函数名称 (格式: module.function)
+    args: dict            # 调用参数 (不含 conn)
+
+
+class TaskConfig(TypedDict):
+    """Task 配置文件结构"""
+    hostname: str         # iamt.yaml 中的主机名
+    tasks: list[TaskCall] # 执行的 task 函数列表
+
+
+def save_task_config(
+    hostname: str,
+    tasks: list[TaskCall],
+    config_name: str,
+    suffix: str = ".dcd.yaml",
+    output_dir: Path | None = None,
+) -> Path | None:
+    """保存 task 配置到文件
+
+    Args:
+        hostname: iamt.yaml 中的主机名
+        tasks: 执行的 task 函数调用列表
+        config_name: 配置文件名 (不含后缀)
+        suffix: 配置文件后缀，默认 ".dcd.yaml"
+        output_dir: 输出目录，默认为当前工作目录
+
+    Returns:
+        保存成功返回配置文件路径，失败返回 None
+    """
+    console = Console()
+
+    if not hostname:
+        console.print("[red]hostname 不能为空[/red]")
+        return None
+
+    if not tasks:
+        console.print("[red]tasks 列表不能为空[/red]")
+        return None
+
+    config_data: TaskConfig = {
+        "hostname": hostname,
+        "tasks": tasks,
+    }
+
+    output_dir = output_dir or Path.cwd()
+    config_file = output_dir / f"{hostname}_{config_name}{suffix}"
+
+    try:
+        config_file.write_text(
+            yaml.dump(config_data, default_flow_style=False, allow_unicode=True, sort_keys=False),
+            encoding="utf-8",
+        )
+        console.print(f"[green]配置已保存到 {config_file.name}[/green]")
+        return config_file
+    except Exception as e:
+        console.print(f"[red]保存配置失败: {e}[/red]")
+        return None
+
+
+def get_hostname_from_iamt_config() -> str:
+    """从 iamt.yaml 中获取第一个 hostname
+
+    Returns:
+        hostname 字符串，未找到返回空字符串
+    """
+    from ..client_config.client_config import ClientConfig
+
+    config = ClientConfig()
+    if not config.config_file.exists():
+        return ""
+
+    try:
+        config_data = yaml.safe_load(config.config_file.read_text(encoding="utf-8"))
+        return next(iter(config_data.get("hosts", {})), "")
+    except Exception:
+        return ""

iamt/utils/task_tracker.py

@@ -0,0 +1,136 @@
+"""Task 调用追踪模块
+
+提供 Connection 包装类，用于追踪 role 执行过程中调用的 task 函数及其参数。
+"""
+
+import functools
+import importlib
+import sys
+from pathlib import Path
+from typing import Any, Callable
+
+from fabric import Connection, Task
+
+from .config_saver import TaskCall
+
+
+class TaskTracker:
+    """Task 调用追踪器
+
+    通过 monkey-patch 模块中的 task 函数，记录 role 执行过程中的所有 task 调用。
+    """
+
+    def __init__(self):
+        self._calls: list[TaskCall] = []
+        self._patched: dict[str, dict[str, Any]] = {}  # module_name -> {func_name -> original_func}
+
+    @property
+    def calls(self) -> list[TaskCall]:
+        """获取记录的 task 调用列表"""
+        return self._calls.copy()
+
+    def clear(self) -> None:
+        """清空调用记录"""
+        self._calls.clear()
+
+    def _get_task_name(self, func: Callable) -> str:
+        """从函数获取 task 名称 (格式: module.function)"""
+        module_name = func.__module__
+        func_name = func.__name__
+
+        # 从模块路径提取简短名称 (如 iamt.modules.docker.docker -> docker)
+        module_parts = module_name.split(".")
+        if "modules" in module_parts:
+            idx = module_parts.index("modules")
+            short_module = module_parts[idx + 1] if idx + 1 < len(module_parts) else module_parts[-1]
+        else:
+            short_module = module_parts[-1]
+
+        return f"{short_module}.{func_name}"
+
+    def _wrap_func(self, func: Callable, task_name: str) -> Callable:
+        """包装函数以追踪调用"""
+        tracker = self
+
+        @functools.wraps(func)
+        def wrapper(conn: Connection, *args: Any, **kwargs: Any) -> Any:
+            tracker._calls.append({
+                "task": task_name,
+                "args": kwargs.copy(),
+            })
+            return func(conn, *args, **kwargs)
+
+        return wrapper
+
+    def patch_module(self, module_name: str) -> None:
+        """Patch 指定模块中的所有 task 函数
+
+        Args:
+            module_name: 完整模块名 (如 iamt.modules.docker.docker)
+        """
+        if module_name in self._patched:
+            return
+
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError:
+            return
+
+        self._patched[module_name] = {}
+
+        for name in dir(module):
+            obj = getattr(module, name)
+            # 检查是否是在该模块定义的函数且被 @task 装饰
+            if isinstance(obj, Task) and obj.body.__module__ == module_name:
+                original_func = obj.body
+                task_name = self._get_task_name(original_func)
+                wrapped = self._wrap_func(original_func, task_name)
+                # 保存原始函数
+                self._patched[module_name][name] = original_func
+                # 替换 Task 对象的 body
+                obj.body = wrapped
+
+    def unpatch_all(self) -> None:
+        """恢复所有被 patch 的函数"""
+        for module_name, funcs in self._patched.items():
+            try:
+                module = sys.modules.get(module_name)
+                if module is None:
+                    continue
+                for name, original_func in funcs.items():
+                    obj = getattr(module, name, None)
+                    if isinstance(obj, Task):
+                        obj.body = original_func
+            except Exception:
+                pass
+        self._patched.clear()
+
+    def patch_role_module(self, role_module_name: str) -> None:
+        """根据 role 所在模块，patch 相关的 task 模块
+
+        Args:
+            role_module_name: role 函数所在的模块名 (如 docker)
+        """
+        # region patch 该模块目录下所有非下划线开头的 py 文件
+        module_dir = Path(__file__).parent.parent / "modules" / role_module_name
+        if module_dir.exists():
+            for py_file in module_dir.glob("*.py"):
+                if py_file.name.startswith("_"):
+                    continue
+                task_module = f"iamt.modules.{role_module_name}.{py_file.stem}"
+                self.patch_module(task_module)
+        # endregion
+
+    def patch_all_modules(self) -> None:
+        """Patch 所有 modules 目录下的 task 模块"""
+        modules_dir = Path(__file__).parent.parent / "modules"
+        for module_dir in modules_dir.iterdir():
+            if not module_dir.is_dir() or module_dir.name.startswith("_"):
+                continue
+            self.patch_role_module(module_dir.name)
+
+    def __enter__(self) -> "TaskTracker":
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.unpatch_all()

iamt-1.1.27.dist-info/METADATA

@@ -0,0 +1,19 @@
+Metadata-Version: 2.3
+Name: iamt
+Version: 1.1.27
+Summary: Add your description here
+Author: lipanpan
+Author-email: lipanpan <lipanpanmail@163.com>
+Requires-Dist: fabric>=3.2.2
+Requires-Dist: fire>=0.7.1
+Requires-Dist: gitpython>=3.1.46
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: prompt-toolkit>=3.0.52
+Requires-Dist: pyperclip>=1.11.0
+Requires-Dist: pyyaml>=6.0.3
+Requires-Dist: questionary>=2.1.1
+Requires-Dist: rich>=14.2.0
+Requires-Dist: toml>=0.10.2
+Requires-Python: ~=3.13.0
+Description-Content-Type: text/markdown
+