Enable compression of ModelCIF files.

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,
         )