From 47ffa8e832b1c996fac191bf95fa08a3dc0d8532 Mon Sep 17 00:00:00 2001
From: mmagnus <mag_dex@o2.pl>
Date: Thu, 10 Oct 2024 20:16:18 -0400
Subject: [PATCH] rna_tools/rna_tools_lib.py: add fetch-cif

---
 rna_tools/rna_pdb_tools.py | 11 ++++++++++-
 rna_tools/rna_tools_lib.py | 17 +++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/rna_tools/rna_pdb_tools.py b/rna_tools/rna_pdb_tools.py
index 63e7221d6..a0994a0ac 100755
--- a/rna_tools/rna_pdb_tools.py
+++ b/rna_tools/rna_pdb_tools.py
@@ -46,7 +46,7 @@
 import glob
 import os
 from rna_tools.rna_tools_lib import edit_pdb, add_header, get_version, \
-                          collapsed_view, fetch, fetch_ba, replace_chain, RNAStructure, \
+                          collapsed_view, fetch, fetch_ba, fetch_cif, replace_chain, RNAStructure, \
                           select_pdb_fragment, sort_strings, set_chain_for_struc
 from rna_tools.tools.rna_x3dna.rna_x3dna import x3DNA
 
@@ -122,6 +122,8 @@ def get_parser():
     parser.add_argument('--fetch', action='store_true', help='fetch file from the PDB db, e.g., 1xjr,\nuse \'rp\' to fetch, fetch a given join, 4w90:C or 4w90_C'  +
                         'the RNA-Puzzles standardized_dataset [around 100 MB]')
 
+    parser.add_argument('--fetch-cif',  action='store_true', help='')
+
     parser.add_argument('--fetch-ba', action='store_true',
                         help='fetch biological assembly from the PDB db')
 
@@ -894,6 +896,13 @@ def get_parser():
     if args.fetch_ba:
         fetch_ba(args.file)
 
+    if args.fetch_cif:
+        if list != type(args.file):
+            args.file = [args.file]
+        ##################################
+        for f in args.file:
+            fetch_cif(f)
+
     if args.collapsed_view or args.cv:
         collapsed_view(args)
 
diff --git a/rna_tools/rna_tools_lib.py b/rna_tools/rna_tools_lib.py
index 221c5b6be..6270679e0 100755
--- a/rna_tools/rna_tools_lib.py
+++ b/rna_tools/rna_tools_lib.py
@@ -2229,6 +2229,23 @@ def fetch_cif_ba(cif_id, path="."):
     print('ok')
     return cif_id + '_ba.cif'
 
+def fetch_cif(cif_id, path="."):
+    """fetch biological assembly cif file from RCSB.org"""
+    import urllib3
+    http = urllib3.PoolManager()
+    # try:
+    ic(cif_id)
+    response = http.request('GET', url='https://files.rcsb.org/download/' +
+                            cif_id.lower() + '.cif')
+    if not response.status == 200:
+        raise PDBFetchError()
+    txt = response.data
+
+    npath = path + os.sep + cif_id + '.cif'
+    #print('downloading... ' + npath)
+    with open(npath, 'wb') as f:
+        f.write(txt)
+    return cif_id + '.cif'
 
 def replace_chain(struc_fn, insert_fn, chain_id):
     """Replace chain of the main file (struc_fn) with some new chain (insert_fn) of given chain id.