deepdata/tools/test.py

"""
从 py-data-juicer 包提取算子参数，生成PostgreSQL算子初始化SQL（填充settings）
依赖：pip install data-juicer
"""
import importlib
import inspect
import json
from typing import Dict, Any, List
from data_juicer.ops.base_op import BaseOP
from data_juicer.ops import ALL_OPS

# --------------------------
# 前端表单schema映射规则（参考原sql脚本settings结构）
# 原settings样例：
# {"repeatPhraseRatio": {
#     "name": "文档词重复率",
#     "description": "xxx",
#     "type": "slider",
#     "defaultVal":0.5,
#     "min":0,"max":1,"step":0.1
# }}
# 支持类型：input / inputNumber / slider / switch / select / range
# --------------------------

def pydantic_field_to_form_schema(field) -> Dict[str, Any]:
    """
    将data‑juicer算子pydantic参数Field转为前端表单配置
    """
    finfo = field.field_info
    field_name = field.name
    desc = finfo.description or ""
    default = finfo.default if finfo.default is not None else None

    schema = {
        "name": field_name,
        "description": desc,
        "defaultVal": default,
        "required": finfo.is_required() is False
    }
    # 推断控件类型
    field_type = finfo.annotation

    if hasattr(finfo, "ge") or hasattr(finfo, "le"):
        schema["min"] = getattr(finfo, "ge", None)
        schema["max"] = getattr(finfo, "le", None)
        if hasattr(finfo, "multiple_of"):
            schema["step"] = finfo.multiple_of
        schema["type"] = "slider"
    elif field_type is bool:
        schema["type"] = "switch"
        schema["checkedLabel"] = "开启"
        schema["unCheckedLabel"] = "关闭"
    elif getattr(finfo, "enum", None):
        # 枚举 -> select下拉框
        schema["type"] = "select"
        schema["options"] = [{"label": e.value, "value": e.value} for e in finfo.enum]
    elif str(field_type).startswith("int") or str(field_type).startswith("float"):
        schema["type"] = "inputNumber"
    elif "List" in str(field_type) or "Tuple" in str(field_type):
        schema["type"] = "range"
    else:
        schema["type"] = "input"
    return schema


def extract_op_settings(op_cls: BaseOP) -> str | None:
    """
    从算子类提取所有配置参数，生成settings JSON字符串
    """
    if not hasattr(op_cls, "__pydantic_fields__"):
        return None
    pydantic_fields = op_cls.__pydantic_fields__
    settings_dict = {}
    for fname, field in pydantic_fields.items():
        # skip 内部私有参数
        if fname.startswith("_"):
            continue
        form_schema = pydantic_field_to_form_schema(field)
        settings_dict[fname] = form_schema
    if len(settings_dict) == 0:
        return None
    return json.dumps(settings_dict, ensure_ascii=False)


def gen_operator_sql(op_cls) -> str:
    """生成单条算子insert sql"""
    op_id = op_cls._name  # data‑juicer算子唯一id，例如 document_deduplicator
    op_full_name = op_cls.__name__
    # 获取注释作为描述
    doc = inspect.getdoc(op_cls) or ""
    # inputs/outputs
    inputs = ",".join(op_cls.ACCEPTED_MODALITIES)
    outputs = ",".join(op_cls.OUTPUT_MODALITIES)
    version = "1.4.4"
    settings_json = extract_op_settings(op_cls)
    file_name = ""
    file_size = None
    is_star = False
    created_by = "system"
    updated_by = "system"

    def escape_sql_str(s):
        if s is None:
            return "NULL"
        return "'" + str(s).replace("'", "''") + "'"

    cols = [
        "id", "name", "description", "version", "inputs", "outputs",
        "runtime", "settings", "file_name", "file_size", "is_star", "created_by", "updated_by"
    ]
    vals = [
        escape_sql_str(op_id),
        escape_sql_str(op_full_name),
        escape_sql_str(doc[:2000]),
        escape_sql_str(version),
        escape_sql_str(inputs),
        escape_sql_str(outputs),
        "NULL",
        escape_sql_str(settings_json),
        escape_sql_str(file_name),
        "NULL" if file_size is None else str(file_size),
        "false",
        escape_sql_str(created_by),
        escape_sql_str(updated_by),
    ]
    sql_line = f"INSERT INTO t_operator ({','.join(cols)}) VALUES ({','.join(vals)}) ON CONFLICT DO NOTHING;"
    return sql_line


def main():
    output_sql_lines: List[str] = []
    print(f"total op count: {len(ALL_OPS)}")
    for op_cls in ALL_OPS.values():
        sql = gen_operator_sql(op_cls)
        output_sql_lines.append(sql)
    # 输出完整脚本
    final_sql = "\n".join(output_sql_lines)
    with open("./data_juicer_ops_gen.sql", "w", encoding="utf‑8") as f:
        f.write("-- auto generated data‑juicer operator init sql with settings\n")
        f.write(final_sql)
    print("Generated file: data_juicer_ops_gen.sql")
    print(final_sql[:2000])


if __name__ == "__main__":
    main()