[Python] RCSB에서 pdb 파일 다운로드하는 프로그램

# PDB ID List를 pdb.txt 파일로 입력받아 나열된 ID의 PDB파일 모두 다운로드
# 다운로드 경로 : ./PDBs 디렉토리
# 다운로드 실패 시 해당 id download_fail_list.txt에 저장


from requests import get
import os


def download_pdb(download_path = './PDBs', download_fail = './download_fail_list.txt'):
    # txt 파일에서 pdb id list 읽어오기
    with open('./pdb.txt', 'r') as pdb_file:
        pdb_id_list = [line.strip() for line in pdb_file]

    # 다운로드 경로 dir 없을 시 dir 생성
    if not os.path.exists(download_path):
        os.makedirs(download_path)

    with open(download_fail, 'w') as fail:
        # list 내 pdb id에 해당하는 파일 다운로드
        for pdb_id in pdb_id_list:
            url = f'https://files.rcsb.org/download/{pdb_id.lower()}.pdb'
            res = get(url)

            if res.status_code == 200: # res OK
                # 이진 모드(wb)로 파일에 저장
                path = os.path.join(download_path, f'{pdb_id}.pdb')
                with open(path, 'wb') as output:
                    output.write(res.content)
                print(f'Downloaded data for PDB ID {pdb_id} to {path}')

            else:
                # 다운로드 실패
                fail.write(f'{pdb_id}\n')
                print(f'Failed to download data for PDB ID: {pdb_id}')


download_pdb()