feat(fix-morph): optimise further to only load stuff if required

neuroml/utils.py

@@ -15,7 +15,7 @@
 import networkx
 
 import neuroml.nml.nml as schema
-from neuroml import NeuroMLDocument
+from neuroml import BiophysicalProperties, Morphology, NeuroMLDocument
 
 from . import loaders
 
@@ -348,14 +348,9 @@ def fix_external_morphs_biophys_in_cell(
     :returns: neuroml document
     :raises KeyError: if referenced morphologies/biophysics cannot be found
     """
-    if overwrite is False:
-        newdoc = copy.deepcopy(nml2_doc)
-    else:
-        newdoc = nml2_doc
-
     # get a list of morph/biophys ids being referred to by cells
     referenced_ids = []
-    for cell in newdoc.cells:
+    for cell in nml2_doc.cells:
         if load_morphology is True and cell.morphology_attr is not None:
             if cell.morphology is None:
                 referenced_ids.append(cell.morphology_attr)
@@ -378,35 +373,66 @@ def fix_external_morphs_biophys_in_cell(
 
     if len(referenced_ids) == 0:
         logger.debug("No externally referenced morphologies or biophysics")
-        return newdoc
+        return nml2_doc
+
+    if overwrite is False:
+        newdoc = copy.deepcopy(nml2_doc)
+    else:
+        newdoc = nml2_doc
 
     # load referenced ids from included files and store them in dicts
-    ext_morphs = {}
-    ext_biophys = {}
+    found = False
+    ext_morphs: Dict[str, Morphology] = {}
+    ext_biophys: Dict[str, BiophysicalProperties] = {}
+
+    # first check the same document
+    if not found and load_morphology is True and newdoc.morphology:
+        for morph in newdoc.morphology:
+            if morph.id in referenced_ids:
+                ext_morphs[morph.id] = morph
+
+            if (len(ext_morphs) + len(ext_biophys)) == len(referenced_ids):
+                logger.debug("Found all references")
+                found = True
+                break
+
+    if (
+        not found
+        and load_biophysical_properties is True
+        and newdoc.biophysical_properties
+    ):
+        for biophys in newdoc.biophysical_properties:
+            if biophys.id in referenced_ids:
+                ext_biophys[biophys.id] = biophys
+
+            if (len(ext_morphs) + len(ext_biophys)) == len(referenced_ids):
+                logger.debug("Found all references")
+                found = True
+                break
+
+    # now check included files (they need to be loaded and parsed, so this can
+    # be computationally intensive)
     # includes/morphology/biophysical_properties should generally be empty
     # lists and not None, but there may be cases where these have been removed
     # after the document was loaded
-    if newdoc.includes:
+    if not found and newdoc.includes:
         for inc in newdoc.includes:
             incdoc = loaders.read_neuroml2_file(inc.href, verbose=False, optimized=True)
+
             if load_morphology is True and incdoc.morphology:
                 for morph in incdoc.morphology:
                     if morph.id in referenced_ids:
                         ext_morphs[morph.id] = morph
+
             if load_biophysical_properties is True and incdoc.biophysical_properties:
                 for biophys in incdoc.biophysical_properties:
                     if biophys.id in referenced_ids:
                         ext_biophys[biophys.id] = biophys
 
-    if load_morphology is True and newdoc.morphology:
-        for morph in newdoc.morphology:
-            if morph.id in referenced_ids:
-                ext_morphs[morph.id] = morph
-
-    if load_biophysical_properties is True and newdoc.biophysical_properties:
-        for biophys in newdoc.biophysical_properties:
-            if biophys.id in referenced_ids:
-                ext_biophys[biophys.id] = biophys
+            if (len(ext_morphs) + len(ext_biophys)) == len(referenced_ids):
+                logger.debug("Found all references")
+                found = True
+                break
 
     # update cells by placing the morphology/biophys in them:
     # if referenced ids are not found, throw errors