本代码用于从pubchem下载小分子(阿司匹林)，并且自动转换成pymol的pdb格式

import pubchempy as pcp
from pymol import cmd
import os
def download_sdf(compound_identifier, output_file):
    try:
        # 尝试根据标识符获取化合物
        if isinstance(compound_identifier, int):
            compound = pcp.Compound.from_cid(compound_identifier)
        else:
            results = pcp.get_compounds(compound_identifier, 'name')
            if results:
                compound = results[0]
            else:
                print(f"未找到与 {compound_identifier} 匹配的化合物。")
                return

        # 下载 SDF 文件
        sdf_content = compound.canonical_smiles
        sdf = pcp.get_sdf(compound.cid)
        with open(output_file, 'w') as file:
            file.write(sdf)
        print(f"SDF 文件已成功保存到 {output_file}")
    except Exception as e:
        print(f"下载过程中出现错误: {e}")


if __name__ == "__main__":
    # 这里可以替换为你想要查询的化合物名称或 CID
    compound_identifier = "aspirin"
    output_file = compound_identifier + ".sdf"
    download_sdf(compound_identifier, output_file)
    # 使用 PyMOL 加载 SDF 文件并进行后续处理
    cmd.load(output_file)
    # 保存为 PDB 文件（可选）
    cmd.save(compound_identifier + ".pdb", state=0)
    cmd.quit()
    # 删除临时 SDF 文件（可选）
    #os.remove(output_file)