add_fk.py

import re
import sys

def add_fk_lines_advanced(input_file, output_file):
    """
    为每个表定义添加外键读取代码
    """
    with open(input_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    new_lines = []
    i = 0

    while i < len(lines):
        line = lines[i]
        new_lines.append(line)

        # 检测表定义的开始
        match = re.match(r'^(\w+)\s*=\s*create_readonly_dataframe\(', line)

        if match:
            var_name = match.group(1)

            # 跳过 tables 变量（它不是实际的表）
            if var_name == 'tables':
                i += 1
                continue

            # 找到对应的表名
            j = i
            table_name = None
            while j < len(lines):
                if 'SELECT * FROM' in lines[j]:
                    table_match = re.search(r'SELECT \* FROM ["\'](\w+)["\']', lines[j])
                    if table_match:
                        table_name = table_match.group(1)
                    break
                j += 1

            # 找到这个表定义的结束位置（找到闭合的括号）
            paren_count = 1
            j = i + 1
            while j < len(lines) and paren_count > 0:
                paren_count += lines[j].count('(') - lines[j].count(')')
                new_lines.append(lines[j])
                j = j + 1

            # 添加外键读取代码
            if table_name:
                fk_code = f'''{var_name}_fk = create_readonly_dataframe(
    pd.read_sql_query('PRAGMA foreign_key_list({table_name})', conn),
)
'''
                new_lines.append(fk_code)

            i = j
            continue

        i += 1

    with open(output_file, 'w', encoding='utf-8') as f:
        f.writelines(new_lines)

    print(f"✓ 修改完成！已保存到 {output_file}")
    print(f"✓ 共处理了 {len([l for l in new_lines if '_fk = create_readonly_dataframe' in l])} 个表的外键")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("使用方法: python add_fk.py <input_file> <output_file>")
        print("示例: python add_fk.py dataframes.py dataframes_with_fk.py")
    else:
        add_fk_lines_advanced(sys.argv[1], sys.argv[2])