code_saver.py

# fmt: off
import sys
from pathlib import Path
import shutil
import ast
import time
import logging
from typing import Set, Union, Optional
from importlib.util import find_spec
import tyro
from dataclasses import dataclass
# fmt: on

# for faster finding of local imports
IGNORE_MODULES = {"sys", "os", "copy", "typing", "tyro", "ast", "shutil", "time", "logging", "pathlib", "importlib", "dataclasses", "threading", "collections", "numpy", "torch", "gymnasium", "tyro"}

# 配置基础logger
logger = logging.getLogger("code_snapshot")
logger.setLevel(logging.WARNING)  # 默认不输出INFO日志


# ---------- 参数定义 ----------
@dataclass
class Args:
    main_script: str
    verbose: bool = False


class ImportFinder(ast.NodeVisitor):
    """AST 访问器，用于查找所有导入语句"""

    def __init__(self, current_file: Union[str, Path]):
        self.current_file: Path = Path(current_file)
        self.imports = set()

    def visit_Import(self, node):
        for alias in node.names:
            self.imports.add(alias.name)

    def visit_ImportFrom(self, node):
        if node.module is None:
            return set()
        self.imports.add(node.module)

        if node.level > 0:  # relative import
            node_parent_path = self.current_file.parents[node.level - 1]
            node_parent = node_parent_path.name
            node_module_path = node_parent_path / node.module.replace(".", "/")
            node_module = f"{node_parent}.{node.module}"
            if node_module_path.is_dir():  # from .module import *
                self.imports.add(node_module)
                for alias in node.names:
                    if alias.name == "*":
                        for sub_file in node_module_path.iterdir():
                            if sub_file.suffix == ".py":  # only import .py files
                                self.imports.add(f"{node_module}.{sub_file.stem}")
                    else:
                        self.imports.add(f"{node_module}.{alias.name}")
            elif node_module_path.with_suffix(".py").is_file():  # from .module import abc
                self.imports.add(node_module)
            else:  # maybe a single file, from file import class/function
                logger.debug(f"无法解析相对导入: {ast.dump(node)} in {self.current_file}")
                print(f"无法解析相对导入: {ast.dump(node)} in {self.current_file}")
        else:
            for alias in node.names:
                self.imports.add(f"{node.module}.{alias.name}")


def find_local_imports(file_path: Path) -> Set[str]:
    """通过AST静态分析找出所有导入的本地模块"""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            tree = ast.parse(f.read(), filename=file_path)
    except Exception as e:
        logger.error(f"读取文件失败 {file_path}: {str(e)}")
        raise

    visitor = ImportFinder(file_path)
    visitor.visit(tree)
    return visitor.imports


def collect_dependent_files(main_script: Path, project_dir_path: str, collected_paths: Set[Path] = set()) -> Set[Path]:
    """递归收集所有依赖的本地文件"""
    if main_script.suffix != ".py":
        return collected_paths
    collected_paths.add(main_script)
    imports = find_local_imports(main_script)
    logger.debug(f"在 {main_script} 中找到的导入: {imports}, 但是将要忽略模块：{IGNORE_MODULES.intersection(imports)}")
    imports -= IGNORE_MODULES  # 忽略一些常见模块
    imports.discard(None)  # 移除None值
    for module in imports:
        sys_path_back = sys.path.copy()
        modules_backup = sys.modules.copy()
        sys_path_front = []
        cur_dir = main_script.parent
        while str(cur_dir) != str(project_dir_path):
            sys_path_front.append(str(cur_dir))
            cur_dir = cur_dir.parent
        sys_path_front.append(str(project_dir_path))
        sys.path = sys_path_front + sys.path
        try:
            logger.debug(f"尝试查找模块: {module}")
            spec = find_spec(module)
        except Exception as e:  # eg: typing.Tuple
            logger.debug(f"模块未找到 {module}: {str(e)}")
            spec = None

        if spec and spec.origin and project_dir_path in spec.origin:
            origin_path = Path(spec.origin)
            if origin_path not in collected_paths:
                collected_paths.add(origin_path)

                if spec.submodule_search_locations:  # __init__.py
                    for sub_location in spec.submodule_search_locations:
                        sys.path.insert(0, sub_location)
                # 递归收集依赖
                collected_paths.update(collect_dependent_files(origin_path, project_dir_path, collected_paths=collected_paths))
            else:
                logger.debug(f"文件已收集过，跳过: {origin_path}")
        sys.path = sys_path_back
        sys.modules = modules_backup
    return collected_paths


# @threaded
def save_code_snapshot(
    main_script: Optional[Union[str, Path]] = None,
    project_dir_path: Optional[Union[str, Path]] = None,
    save_dir: Optional[Union[str, Path]] = None,
    verbose: bool = False,
) -> None:
    """
    保存代码快照，包括主脚本和所有相关依赖文件。

    :param Optional[Union[str, Path]] main_script: 主脚本路径, 默认通过 sys.argv[0] 获取当前执行的脚本路径.
    :param Optional[Union[str, Path]] project_dir_path: 项目路径，默认为当前工作目录, 只保存在该目录下的相关脚本文件.
    :param Optional[Union[str, Path]] save_dir: 保存快照的目录, 默认会自动创建一个带时间戳的目录.
    :param bool verbose: 是否输出详细日志. Defaults to False.
    :return: None
    """
    # 配置logger输出
    if verbose:
        logger.setLevel(logging.DEBUG)
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
        if not logger.handlers:
            logger.addHandler(handler)
    if not main_script:
        main_script = sys.argv[0]  # 默认使用当前执行的脚本
    main_script = Path(main_script).resolve()
    logger.debug(f"开始处理主脚本: {main_script}")

    try:
        if save_dir is None:
            timestamp = time.strftime("%Y-%m-%d_%H-%M-%S")
            save_dir = Path(f"code_snapshots/snapshot_{timestamp}")
            save_dir.mkdir(parents=True, exist_ok=True)
            logger.info(f"创建快照目录: {save_dir}")
        save_dir = Path(save_dir).resolve()
        logger.debug(f"保存目录: {save_dir}")

        if project_dir_path is None:
            project_dir_path = Path.cwd()
        project_dir_path = str(project_dir_path)

        if not Path(project_dir_path).exists():
            assert False, f"项目路径不存在: {project_dir_path}"

        if project_dir_path not in sys.path:
            sys.path.insert(0, project_dir_path)
            logger.debug(f"添加项目路径到sys.path: {project_dir_path}")
        files_to_save = collect_dependent_files(main_script, project_dir_path)
        logger.info(f"发现 {len(files_to_save)} 个相关文件")
        # 复制文件
        for filepath in files_to_save:
            try:
                # 如果file不是symlink，resolve()和absolute()结果相同
                source_path = filepath.resolve()  # 解析symbolic link后的真实路径
                abs_path = filepath.absolute()  # 文件的绝对路径
                rel_path = abs_path.relative_to(project_dir_path)
                target_path = save_dir / rel_path
                target_path.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(source_path, target_path)
                logger.debug(f"复制文件: {filepath} -> {target_path}")
            except Exception as e:
                logger.error(f"复制文件失败 {filepath}: {str(e)}")

        return

    except Exception as e:
        logger.error(f"生成快照失败: {str(e)}")
        raise


# ---------- 主程序 ----------
def main():
    """独立主函数，处理参数解析和执行流程"""
    args = tyro.cli(Args)  # 自动解析命令行参数
    snapshot_thread = save_code_snapshot(main_script=args.main_script, verbose=args.verbose)
    snapshot_thread.join()


if __name__ == "__main__":
    main()