diff --git a/convert_to_modelcif.py b/convert_to_modelcif.py index ffac49a34eb0ae23637ce452fbeed2897861da54..91a5176e512f6087bc734d5133be8817bd897e2a 100755 --- a/convert_to_modelcif.py +++ b/convert_to_modelcif.py @@ -4,11 +4,13 @@ file with a lot of metadata in place.""" from typing import Tuple +import gzip import hashlib import json import os import pickle import re +import shutil import sys import zipfile @@ -35,7 +37,7 @@ flags.DEFINE_integer( "Model to be converted into ModelCIF, use '--select_all' to convert all " + "models found in '--af2_output'", ) - +flags.DEFINE_bool("compress", False, "Compress the ModelCIF file using Gzip") flags.mark_flags_as_required(["ap_output"]) FLAGS = flags.FLAGS @@ -218,7 +220,8 @@ def _store_as_modelcif( structure: BioStructure, mdl_file: str, out_dir: str, - # file_prfx, compress, add_files + compress: bool = False, + # file_prfx, add_files ) -> None: """Create the actual ModelCIF file.""" system = modelcif.System( @@ -288,6 +291,8 @@ def _store_as_modelcif( encoding="ascii", ) as mmcif_fh: modelcif.dumper.write(mmcif_fh, [system]) + if compress: + _compress_cif_file(f"{mdl_file}.cif") # Create associated archive for archive in system.repositories[0].files: with zipfile.ZipFile( @@ -296,12 +301,18 @@ def _store_as_modelcif( for zfile in archive.files: cif_zip.write(zfile.path, arcname=zfile.path) os.remove(zfile.path) - # if compress: - # _compress_cif_file(mdl_fle) finally: os.chdir(oldpwd) +def _compress_cif_file(cif_file): + """Compress cif file and delete original.""" + with open(cif_file, "rb") as f_in: + with gzip.open(cif_file + ".gz", "wb") as f_out: + shutil.copyfileobj(f_in, f_out) + os.remove(cif_file) + + def _get_model_details(cmplx_name: str, data_json: dict) -> str: """Get the model description.""" ap_versions = [] @@ -422,6 +433,7 @@ def alphapulldown_model_to_modelcif( scr_file: str, out_dir: str, prj_dir: str, + compress: bool = False, ) -> None: """Convert an AlphaPulldown model into a ModelCIF formatted mmCIF file. @@ -442,7 +454,7 @@ def alphapulldown_model_to_modelcif( # read quality scores from pickle file _get_scores(modelcif_json, scr_file) - _store_as_modelcif(modelcif_json, structure, mdl_file, out_dir) + _store_as_modelcif(modelcif_json, structure, mdl_file, out_dir, compress) # ToDo: ENABLE logging.info(f"... done with '{mdl_file}'") @@ -536,7 +548,12 @@ def main(argv): ) for mdl, scrs in model_list: alphapulldown_model_to_modelcif( - complex_name, mdl, scrs, model_dir, FLAGS.ap_output + complex_name, + mdl, + scrs, + model_dir, + FLAGS.ap_output, + FLAGS.compress, )