Add metagraph.get_metanode

Author: dhimmel

hetio/hetnet.py

@@ -204,9 +204,25 @@ def __eq__(self, other):
 class MetaGraph(BaseGraph):
 
     def __init__(self):
-        """ """
         BaseGraph.__init__(self)
 
+    def get_metanode(self, metanode):
+        """
+        Return the metanode specified by the input, which can be either a:
+         - MetaNode (passthrough)
+         - metanode kind (str)
+         - metanode abbreviation (str)
+        """
+        if isinstance(metanode, MetaNode):
+            return metanode
+        if metanode in self.node_dict:
+            return self.get_node(metanode)
+        # Assume metanode must be an abbreviation
+        return self.get_node(self.abbrev_to_kind[metanode])
+
+    def get_metaedge(self, metanode):
+        pass
+
     @staticmethod
     def from_edge_tuples(metaedge_tuples, kind_to_abbrev=None):
         """Create a new metagraph defined by its edges."""
@@ -231,6 +247,7 @@ def from_edge_tuples(metaedge_tuples, kind_to_abbrev=None):
     def set_abbreviations(self, kind_to_abbrev):
         """Add abbreviations as an attribute for metanodes and metaedges"""
         self.kind_to_abbrev = kind_to_abbrev
+        self.abbrev_to_kind = {v: k for k, v in kind_to_abbrev.items()}
         for kind, metanode in self.node_dict.items():
             metanode.abbrev = kind_to_abbrev[kind]
         for metaedge in self.edge_dict.values():
@@ -281,10 +298,9 @@ def extract_metapaths(self, source, target=None, max_length=4):
         the metagraph from source to target. If target is None (default), then
         metapaths to any target node are returned.
         """
-        if not isinstance(source, MetaNode):
-            source = self.node_dict[source]
-        if target and not isinstance(target, MetaNode):
-            target = self.node_dict[target]
+        source = self.get_metanode(source)
+        if target:
+            target = self.get_metanode(target)
 
         assert max_length >= 0
         if max_length == 0:

test/graph_test.py

@@ -111,6 +111,14 @@ def test_disase_gene_example():
         ('Gene', 'Gene', 'interaction', 'both'),
     ]
     metagraph = hetio.hetnet.MetaGraph.from_edge_tuples(metaedge_tuples)
+
+    # Test metagraph getter methods
+    gene_metanode = metagraph.node_dict['Gene']
+    assert metagraph.get_metanode(gene_metanode) == gene_metanode
+    assert metagraph.get_metanode('Gene') == gene_metanode
+    assert metagraph.get_metanode('G') == gene_metanode
+
+    # Create graph
     graph = hetio.hetnet.Graph(metagraph)
     nodes = dict()