Source code for allensdk.internal.core.simpletree
from six import iteritems
[docs]class SimpleTree( object ):
def __init__(self, nodes,
node_id_cb,
parent_id_cb):
self.node_list = nodes
self._nodes = { node_id_cb(n):n for n in nodes }
self._parent_ids = { nid:parent_id_cb(n) for nid,n in iteritems(self._nodes) }
self._child_ids = { nid:[] for nid in self._nodes }
for nid in self._parent_ids:
pid = self._parent_ids[nid]
if pid:
self._child_ids[pid].append(nid)
[docs] def node_ids(self):
return self._nodes.keys()
[docs] def parent_id(self, nid):
try:
return self._parent_ids[nid]
except KeyError:
raise KeyError("Could not find parent for node %s" % str(nid))
[docs] def child_ids(self, nid):
try:
for cid in self._child_ids[nid]:
yield cid
except KeyError:
raise KeyError("Could not find children for node %s" % str(nid))
[docs] def ancestor_ids(self, nid):
try:
pid = nid
while pid:
yield pid
pid = self.parent_id(pid)
except:
raise KeyError("Could not find ancestors for node %s" % str(nid))
[docs] def descendant_ids(self, nid):
ids = [nid]
try:
while ids:
nid = ids.pop()
yield nid
ids += self.child_ids(nid)
except KeyError:
raise KeyError("Could not find descendants for node %s" % str(nid))
[docs] def node(self, nid):
return self._nodes[nid]
[docs] def nodes(self, nids=None):
if nids is None:
nids = self.node_ids()
for nid in nids:
yield self.node(nid)
[docs] def parent(self, nid):
return self.node(self.parent_id(nid))
[docs] def children(self, nid):
for node in self.nodes(self.child_ids(nid)):
yield node
[docs] def descendants(self, nid):
for node in self.nodes(self.descendant_ids(nid)):
yield node
[docs] def ancestors(self, nid):
for node in self.nodes(self.ancestor_ids(nid)):
yield node