Cache PDB downloads

MolecularNodes/load.py

@@ -1,5 +1,6 @@
 import requests
 import io
+from pathlib import Path
 import bpy
 import numpy as np
 from . import coll
@@ -10,7 +11,6 @@
 from . import pkg
 from . import obj
 
-
 bpy.types.Scene.mol_pdb_code = bpy.props.StringProperty(
     name = 'pdb_code', 
     description = 'The 4-character PDB code to download', 
@@ -19,6 +19,13 @@
     subtype = 'NONE', 
     maxlen = 4
     )
+bpy.types.Scene.mol_cache_dir = bpy.props.StringProperty(
+    name = 'cache_dir',
+    description = 'Location to cache PDB files',
+    options = {'TEXTEDIT_UPDATE'},
+    default = str(Path('~', '.MolecularNodes').expanduser()),
+    subtype = 'NONE'
+)
 bpy.types.Scene.mol_import_center = bpy.props.BoolProperty(
     name = "mol_import_centre", 
     description = "Move the imported Molecule on the World Origin",
@@ -71,16 +78,18 @@
 
 
 def molecule_rcsb(
-    pdb_code,               
+    pdb_code,             
     center_molecule = False,               
     del_solvent = True,               
     include_bonds = True,   
     starting_style = 0,               
-    setup_nodes = True              
+    setup_nodes = True,
+    cache_dir = None,      
     ):
     mol, file = open_structure_rcsb(
         pdb_code = pdb_code, 
-        include_bonds=include_bonds
+        include_bonds=include_bonds,
+        cache_dir = cache_dir
         )
 
     mol_object, coll_frames = create_molecule(
@@ -167,11 +176,11 @@ def molecule_local(
     return mol_object
 
 
-def open_structure_rcsb(pdb_code, include_bonds = True):
+def open_structure_rcsb(pdb_code, cache_dir = None, include_bonds = True):
     import biotite.structure.io.mmtf as mmtf
     import biotite.database.rcsb as rcsb
 
-    file = mmtf.MMTFFile.read(rcsb.fetch(pdb_code, "mmtf"))
+    file = mmtf.MMTFFile.read(rcsb.fetch(pdb_code, "mmtf", target_path = cache_dir))
 
     # returns a numpy array stack, where each array in the stack is a model in the 
     # the file. The stack will be of length = 1 if there is only one model in the file

MolecularNodes/ui.py

@@ -29,7 +29,8 @@
             center_molecule=bpy.context.scene.mol_import_center, 
             del_solvent=bpy.context.scene.mol_import_del_solvent,
             include_bonds=bpy.context.scene.mol_import_include_bonds,
-            starting_style=bpy.context.scene.mol_import_default_style
+            starting_style=bpy.context.scene.mol_import_default_style,
+            cache_dir=bpy.context.scene.mol_cache_dir
         )
 
         bpy.context.view_layer.objects.active = mol_object
@@ -86,12 +87,15 @@
     col_main.scale_y = 1.0
     col_main.alignment = 'Expand'.upper()
     col_main.label(text = "Download from PDB")
+    col_main.prop(
+        bpy.context.scene,
+        'mol_cache_dir',
+        text = 'Cache dir')
     row_import = col_main.row()
     row_import.prop(bpy.context.scene, 'mol_pdb_code', text='PDB ID')
     row_import.operator('mol.import_protein_rcsb', text='Download', icon='IMPORT')
 
 
-
 def MOL_PT_panel_local(layout_function, ):
     col_main = layout_function.column(heading = '', align = False)
     col_main.alert = False

tests/test_load.py

@@ -175,4 +175,15 @@ def test_load_small_mol(snapshot):
 
     bond_types = mn.obj.get_attribute(obj, 'bond_type')
     edges = ''.join([str(bond_type) for bond_type in bond_types])
-    snapshot.assert_match(edges, 'asn_edges.txt')
+    snapshot.assert_match(edges, 'asn_edges.txt')
+
+def test_rcsb_cache(snapshot):
+    from pathlib import Path
+    from shutil import rmtree
+    # we want to make sure cached files are freshly downloaded, but
+    # we don't want to delete our entire real cache
+    test_cache = Path(Path.home(), '.MolecularNodesTests')
+    if test_cache.exists():
+        rmtree(test_cache)
+    _ = mn.load.molecule_rcsb('6BQN', cache_dir = test_cache)
+    assert (test_cache / '6BQN.mmtf').exists()