如何高效从4000个Python文件的ScriptInfo装饰器提取标题值？

Great question—handling 4000 Python files with inconsistent decorator syntax and variable naming is no small task, and cutting down from two passes to one will save you a ton of time. Let’s dive into a more efficient approach using Python’s Abstract Syntax Tree (AST) that lets you extract the title in a single pass per file.

Why AST Instead of String Matching/Two Passes? #

String regex or manual parsing is fragile here (variable names can be anything, decorator parameters are unpredictable). The AST parses your code into a structured tree, so you can reliably:

1. Track variable-to-value mappings as you traverse the file
2. Identify the @ScriptInfo decorator and extract the variable it uses for the title
3. Tie those two pieces together in one linear pass, no need to re-read the file

Single-Pass Implementation #

Here’s a concrete script that does this:

import ast
import os
from multiprocessing import Pool  # Optional: For parallel processing of files

def extract_script_title(file_path):
    target_var = None
    title_value = None
    var_value_map = {}

    with open(file_path, 'r', encoding='utf-8') as f:
        tree = ast.parse(f.read())

    for node in ast.walk(tree):
        # First, track all string constant variable assignments
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name):
                    var_name = target.id
                    # Only capture direct string assignments (adjust if titles use expressions)
                    if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
                        var_value_map[var_name] = node.value.value
                        # If this is the variable we're looking for from the decorator, exit early
                        if var_name == target_var:
                            title_value = node.value.value
                            break
            if title_value:
                break

        # Second, check for the ScriptInfo decorator and get its target variable
        elif isinstance(node, ast.FunctionDef):
            for decorator in node.decorator_list:
                if (isinstance(decorator, ast.Call) and 
                    isinstance(decorator.func, ast.Name) and 
                    decorator.func.id == 'ScriptInfo'):
                    # Check positional arguments for a variable reference
                    for arg in decorator.args:
                        if isinstance(arg, ast.Name):
                            target_var = arg.id
                            # If we already saw this variable's value, grab it immediately
                            if target_var in var_value_map:
                                title_value = var_value_map[target_var]
                                break
                    # Check keyword arguments for a variable reference (e.g., title=my_title)
                    for kw in decorator.keywords:
                        if isinstance(kw.value, ast.Name):
                            target_var = kw.value.id
                            if target_var in var_value_map:
                                title_value = var_value_map[target_var]
                                break
                    if title_value:
                        break
            if title_value:
                break

    return title_value

# Example usage (parallel processing for 4000 files)
if __name__ == "__main__":
    root_dir = "/path/to/your/python/files"
    py_files = [os.path.join(root, f) for root, _, files in os.walk(root_dir) for f in files if f.endswith('.py')]

    # Use parallel processing to speed up handling 4000 files
    with Pool() as pool:
        results = pool.map(extract_script_title, py_files)
    
    # Print or save results
    for file_path, title in zip(py_files, results):
        if title:
            print(f"File: {os.path.basename(file_path)} | Title: {title}")

Key Advantages #

1. Single pass per file: We traverse the AST once, tracking variables and checking decorators in sequence—no need to read files twice.
2. Robustness: Handles any variable name for the title, and works whether the decorator uses positional or keyword arguments.
3. Speed: Early exits as soon as we find the title value, avoiding unnecessary node processing. For 4000 files, adding parallel processing (via multiprocessing.Pool) will cut down total runtime significantly.

Edge Cases to Consider #

• If your title values come from expressions (e.g., my_title = get_title_from_db() instead of a direct string), you’ll need to extend the code to evaluate those expressions (use ast.literal_eval for safe cases, or a sandboxed evaluator if needed).
• If a file has multiple @ScriptInfo decorators, this code grabs the first one—adjust the logic if you need to extract all titles.

内容的提问来源于stack exchange，提问作者Erick Tompson