Add option to return data when retrieving singleton/virtual scaffolds.

Author: Oliver Scott

scaffoldgraph/analysis/general.py

@@ -5,8 +5,8 @@
 """
 
 
-def get_virtual_scaffolds(scaffoldgraph):
-    """Get 'virtual' scaffolds in within a scaffold graph.
+def get_virtual_scaffolds(scaffoldgraph, data=False, default=None):
+    """Get 'virtual' scaffolds within a scaffold graph.
 
     Virtual scaffolds represent scaffolds that are not directly obtained from
     any molecule of the collection, but generated by the pruning process.
@@ -17,6 +17,13 @@ def get_virtual_scaffolds(scaffoldgraph):
     ----------
     scaffoldgraph : ScaffoldGraph
         A ScaffoldGraph object to query
+    data : str, bool, optional
+        The scaffold node attribute returned in 2-tuple (n, ddict[data]).
+        If True, return entire node attribute dict as (n, ddict).
+        If False, return just the nodes n. The default is False.
+    default : value, bool, optional
+        Value used for nodes that don't have the requested attribute.
+        Only relevant if data is not True or False.
 
     Returns
     -------
@@ -25,17 +32,22 @@ def get_virtual_scaffolds(scaffoldgraph):
 
     """
     virtual = []
-    for scaffold in scaffoldgraph.get_scaffold_nodes():
+    for scaffold, d in scaffoldgraph.get_scaffold_nodes(True):
         mol_count = 0
         for succ in scaffoldgraph.successors(scaffold):
             if scaffoldgraph.nodes[succ].get('type') == 'molecule':
                 mol_count += 1
         if mol_count == 0:
-            virtual.append(scaffold)
+            if data is False:
+                virtual.append(scaffold)
+            elif data is True:
+                virtual.append((scaffold, d))
+            else:
+                virtual.append((scaffold, d.get(data, default)))
     return virtual
 
 
-def get_singleton_scaffolds(scaffoldgraph):
+def get_singleton_scaffolds(scaffoldgraph, data=False, default=None):
     """Get singleton scaffolds within a scaffold graph.
 
     Singleton scaffolds represent scaffolds that are direct members of only
@@ -45,6 +57,13 @@ def get_singleton_scaffolds(scaffoldgraph):
     ----------
     scaffoldgraph : ScaffoldGraph
         A ScaffoldGraph object to query
+    data : str, bool, optional
+        The scaffold node attribute returned in 2-tuple (n, ddict[data]).
+        If True, return entire node attribute dict as (n, ddict).
+        If False, return just the nodes n. The default is False.
+    default : value, bool, optional
+        Value used for nodes that don't have the requested attribute.
+        Only relevant if data is not True or False.
 
     Returns
     -------
@@ -53,11 +72,16 @@ def get_singleton_scaffolds(scaffoldgraph):
 
     """
     singletons = []
-    for scaffold in scaffoldgraph.get_scaffold_nodes():
+    for scaffold, d in scaffoldgraph.get_scaffold_nodes(True):
         mol_count = 0
         for succ in scaffoldgraph.successors(scaffold):
             if scaffoldgraph.nodes[succ].get('type') == 'molecule':
                 mol_count += 1
         if mol_count == 1:
-            singletons.append(scaffold)
+            if data is False:
+                singletons.append(scaffold)
+            elif data is True:
+                singletons.append((scaffold, d))
+            else:
+                singletons.append((scaffold, d.get(data, default)))
     return singletons