wehub-resource-sync
259 lines
11 KiB
Python
259 lines
11 KiB
Python
import libcst as cst
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, Set, Union, List, Tuple
|
|
|
|
# ==============================================================================
|
|
# 以下所有函数和类均保持您提供的原始版本，没有任何改动
|
|
# ==============================================================================
|
|
def get_unique_module_names(imports_dict: Dict[str, str]) -> Set[str]:
|
|
"""
|
|
从字典的值中提取出所有唯一的、纯净的模块名。
|
|
它会移除前缀 '..' 和末尾的 '.'。
|
|
"""
|
|
unique_names = set()
|
|
|
|
for prefix_value in imports_dict.values():
|
|
temp_name = prefix_value
|
|
|
|
# 1. 移除开头的 '..'
|
|
if temp_name.startswith(".."):
|
|
temp_name = temp_name[2:]
|
|
|
|
# 2. 移除末尾的 '.'
|
|
final_name = temp_name.rstrip('.')
|
|
|
|
# 3. 将最终结果添加到集合中，自动保证唯一性
|
|
if final_name:
|
|
unique_names.add(final_name)
|
|
|
|
return unique_names
|
|
def get_full_name(node: Union[cst.Name, cst.Attribute, cst.ImportFrom]) -> str:
|
|
if isinstance(node, cst.Name):
|
|
return node.value
|
|
elif isinstance(node, cst.Attribute):
|
|
return get_full_name(node.value) + "." + node.attr.value
|
|
elif isinstance(node, cst.ImportFrom):
|
|
module_parts = []
|
|
if node.relative:
|
|
module_parts.append("." * len(node.relative))
|
|
if node.module:
|
|
module_parts.append(get_full_name(node.module))
|
|
return "".join(module_parts)
|
|
else:
|
|
return ""
|
|
|
|
class ModelingImportCollector(cst.CSTVisitor):
|
|
def __init__(self):
|
|
self.imports: Dict[str, str] = {} # name -> module_path
|
|
self.prefixes_before_modeling: Dict[str, str] = {}
|
|
def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
|
|
modname = get_full_name(node)
|
|
if "modeling" in modname:
|
|
modeling_index = modname.find("modeling")
|
|
prefix = modname[:modeling_index]
|
|
for alias in node.names:
|
|
name_in_scope = alias.evaluated_name
|
|
self.imports[alias.evaluated_name] = modname
|
|
self.prefixes_before_modeling[name_in_scope] = prefix
|
|
|
|
class DependencyCollector(cst.CSTVisitor):
|
|
def __init__(self):
|
|
self.names: Set[str] = set()
|
|
def visit_Name(self, node: cst.Name) -> None:
|
|
self.names.add(node.value)
|
|
|
|
class ModuleInfoCollector(cst.CSTVisitor):
|
|
def __init__(self):
|
|
self.defs: Dict[str, Union[cst.ClassDef, cst.FunctionDef, cst.Assign]] = {}
|
|
self.imports: Dict[str, Union[cst.Import, cst.ImportFrom]] = {}
|
|
self.class_stack: List[str] = []
|
|
def visit_ClassDef(self, node: cst.ClassDef) -> None:
|
|
self.defs[node.name.value] = node
|
|
self.class_stack.append(node.name.value)
|
|
def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
|
|
self.class_stack.pop()
|
|
def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
|
|
if not self.class_stack:
|
|
self.defs[node.name.value] = node
|
|
else:
|
|
fullname = ".".join(self.class_stack + [node.name.value])
|
|
self.defs[fullname] = node
|
|
def visit_Assign(self, node: cst.Assign) -> None:
|
|
if not self.class_stack:
|
|
for target_wrapper in node.targets:
|
|
if isinstance(target_wrapper.target, cst.Name):
|
|
self.defs[target_wrapper.target.value] = node
|
|
def visit_Import(self, node: cst.Import) -> None:
|
|
for alias in node.names:
|
|
name_in_scope = alias.asname.name.value if alias.asname else alias.name.value
|
|
self.imports[name_in_scope] = node
|
|
def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
|
|
for alias in node.names:
|
|
name_in_scope = alias.asname.name.value if alias.asname else alias.name.value
|
|
self.imports[name_in_scope] = node
|
|
|
|
def parse_file(file_path: str) -> Tuple[Dict, Dict, cst.Module]:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
code = f.read()
|
|
module = cst.parse_module(code)
|
|
collector = ModuleInfoCollector()
|
|
module.visit(collector)
|
|
return collector.defs, collector.imports, module
|
|
|
|
def collect_recursive(
|
|
name: str, defs: Dict[str, cst.CSTNode], imports: Dict[str, cst.CSTNode],
|
|
seen: Set[str], module: cst.Module,
|
|
) -> Tuple[Dict[str, str], Set[str], Dict[str, List[str]]]:
|
|
if name in seen or name not in defs:
|
|
return {}, set(), {}
|
|
seen.add(name)
|
|
node = defs[name]
|
|
dependencies = {name: []}
|
|
dep_collector = DependencyCollector()
|
|
node.visit(dep_collector)
|
|
results = {name: module.code_for_node(node)}
|
|
collected_imports = set()
|
|
for dep in dep_collector.names:
|
|
if dep in defs and dep not in seen:
|
|
dep_results, dep_imports , dep_deps = collect_recursive(dep, defs, imports, seen, module)
|
|
results.update(dep_results)
|
|
collected_imports.update(dep_imports)
|
|
dependencies.update(dep_deps)
|
|
dependencies[name].append(dep) # 记录依赖关系 A -> B
|
|
elif dep in imports:
|
|
import_node = imports[dep]
|
|
import_code = module.code_for_node(import_node)
|
|
collected_imports.add(import_code)
|
|
dependencies[name].append(dep)
|
|
return results, collected_imports, dependencies
|
|
|
|
def resolve_file_path(current_file: str, modpath: str) -> Path:
|
|
dots = len(modpath) - len(modpath.lstrip("."))
|
|
parts = modpath.lstrip(".").split(".")
|
|
cur_dir = Path(current_file).parent
|
|
for _ in range(dots - 1):
|
|
cur_dir = cur_dir.parent
|
|
file_path = cur_dir.joinpath(*parts).with_suffix(".py")
|
|
return file_path if file_path.exists() else None
|
|
|
|
def expand_modeling_imports(file_path: str) -> Dict[str, str]:
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
code = f.read()
|
|
module = cst.parse_module(code)
|
|
imp_collector = ModelingImportCollector()
|
|
module.visit(imp_collector)
|
|
expanded_defs = {}
|
|
all_imports = set()
|
|
seen = set()
|
|
dependencies = {}
|
|
for name, modpath in imp_collector.imports.items():
|
|
target_file = resolve_file_path(file_path, modpath)
|
|
if not target_file: continue
|
|
defs, imports, parsed_module = parse_file(str(target_file))
|
|
if name in defs:
|
|
new_defs, new_imports, new_deps = collect_recursive(name, defs, imports, seen, parsed_module)
|
|
expanded_defs.update(new_defs)
|
|
all_imports.update(new_imports)
|
|
dependencies.update(new_deps)
|
|
expanded = {}
|
|
for i, import_code in enumerate(sorted(list(all_imports))):
|
|
expanded[f"__import_{i}__"] = import_code
|
|
expanded.update(expanded_defs)
|
|
unique_modules = get_unique_module_names(imp_collector.prefixes_before_modeling)
|
|
return expanded, dependencies,unique_modules # 返回代码和依赖关系
|
|
|
|
def save_results_to_txt(result: Dict[str, str], output_file: str):
|
|
imports_to_write = []
|
|
defs_to_write = {}
|
|
for key, value in result.items():
|
|
if key.startswith("__import_"):
|
|
imports_to_write.append(value)
|
|
else:
|
|
defs_to_write[key] = value
|
|
with open(output_file, "w", encoding="utf-8") as f:
|
|
if imports_to_write:
|
|
f.write("### === Imports === ###\n")
|
|
for imp in imports_to_write:
|
|
f.write(f"{imp}\n")
|
|
f.write("\n" + "="*50 + "\n\n")
|
|
if defs_to_write:
|
|
f.write("### === Definitions === ###\n")
|
|
for k, v in sorted(defs_to_write.items()):
|
|
f.write(f"=== {k} ===\n")
|
|
f.write(f"{v}\n\n")
|
|
|
|
# ==============================================================================
|
|
# ### NEW ### 以下是为“文件重写”这一新增功能而添加的全新、独立的模块
|
|
# ==============================================================================
|
|
|
|
class ModelingImportNodeCollector(cst.CSTVisitor):
|
|
"""一个专门用于收集待删除 import 节点的新 Visitor。"""
|
|
def __init__(self):
|
|
self.nodes_to_remove: Set[cst.ImportFrom] = set()
|
|
|
|
def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
|
|
modname = get_full_name(node)
|
|
if "modeling" in modname:
|
|
self.nodes_to_remove.add(node)
|
|
|
|
class ImportRemover(cst.CSTTransformer):
|
|
"""一个独立的转换器，用于从语法树中删除指定的import节点。"""
|
|
def __init__(self, nodes_to_remove: Set[cst.ImportFrom]):
|
|
self.nodes_to_remove = nodes_to_remove
|
|
|
|
def leave_ImportFrom(
|
|
self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
|
|
) -> Union[cst.ImportFrom, cst.RemovalSentinel]:
|
|
if original_node in self.nodes_to_remove:
|
|
return cst.RemoveFromParent()
|
|
return updated_node
|
|
|
|
def remove_imports_and_rewrite(file_path: str):
|
|
"""
|
|
一个独立的函数，封装了文件读取、收集待删除节点、转换和重写的操作。
|
|
"""
|
|
# 1. 再次读取和解析文件，以启动独立的重写流程
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
code = f.read()
|
|
module = cst.parse_module(code)
|
|
|
|
# 2. 收集需要删除的节点
|
|
node_collector = ModelingImportNodeCollector()
|
|
module.visit(node_collector)
|
|
|
|
nodes_to_remove = node_collector.nodes_to_remove
|
|
if not nodes_to_remove:
|
|
print(f"No 'modeling' imports found in '{file_path}' to remove.")
|
|
return
|
|
|
|
# 3. 使用转换器生成修改后的代码
|
|
print(f"Removing {len(nodes_to_remove)} 'modeling' import(s) from '{file_path}'...")
|
|
remover = ImportRemover(nodes_to_remove)
|
|
modified_tree = module.visit(remover)
|
|
|
|
# 4. 将修改后的代码写回原文件
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
f.write(modified_tree.code)
|
|
print("File rewrite complete.")
|
|
|
|
|
|
# ==============================================================================
|
|
# ### MODIFIED ### 主程序块现在按顺序执行两个功能
|
|
# ==============================================================================
|
|
|
|
if __name__ == "__main__":
|
|
file_to_parse = "/home/hsz/PaddleFormers/PaddleFormers/paddleformers/transformers/convert/example/test_model.py"
|
|
output_filename = "modeling_imports_results.txt"
|
|
|
|
# --- 步骤 1: 执行完整的原有功能 ---
|
|
# 调用函数，其接口和返回值完全没有改变
|
|
# 同时也修正了之前版本中解包错误的bug
|
|
combined_results = expand_modeling_imports(file_to_parse)
|
|
|
|
# 保存结果，完成原有任务
|
|
save_results_to_txt(combined_results, output_filename)
|
|
print(f"Code extraction complete. Results saved to {output_filename}")
|
|
|
|
# --- 步骤 2: 在原有功能完成后，独立执行新增的功能 ---
|
|
remove_imports_and_rewrite(file_to_parse) |