taskgraph.py

#!/usr/bin/env python

from collections import deque

class TaskGraph(object):
    ''' simple acyclic, directed property graph with iteration

        use case:
            1) model nodes as dicts and dependencies as directed edges (name=occ -> name=twist)
            2) start with all nodes colored down (_st=0)
            3) iter() to return reachable, down nodes (_st=0)
            4) as nodes are started, mark them up with mark_node_up() to free child nodes for
               iteration (sets self.nodes[id]['_st']=1)
        
        >>> t = TaskGraph()
        >>> a = t.add_node({'name': 'a'})
        >>> b = t.add_node({'name': 'b'})
        >>> c = t.add_node({'name': 'c'})
        >>> t.add_edge(a, b)
        >>> t.add_edge(b, c)
        >>> 
        >>> for x in t:
        >>>     print t.get_node(x)['name']
        >>>     t.mark_node_up(x)
        c
        b
        a

        caveats:
            1) relies on monotonically increasing key ids which map to indexes
            2) does not do cycle detection
    '''
    
    def __init__(self, nodes=[], edges=[]):
        self.nodes = nodes
        self.edges = edges
        self.idx   = -1
        return

    def __iter__(self):
        while 1:
            nodes = self._get_ready_nodes()
            if nodes == None: break
            for node in nodes: yield node

    def _incr_id(self):
        self.idx = self.idx + 1
        return self.idx

    def _get_ready_nodes(self):
        ''' returns a list of nodes ready to be started '''
        red_nodes   = [ x['_id'] for x in self.nodes if x['_st'] == 0 ]
        if not red_nodes: return None
        ready_nodes = deque()

        for nid in red_nodes:
            ready, parent = True, False
            ## iterate over edges looking for ready, parents
            for edge in self.edges:
                if edge[0] == nid and self.nodes[edge[1]]['_st'] != 1: ready = False
                if edge[1] == nid: parent = True
            if ready:
                ## prioritize parents
                if parent: ready_nodes.appendleft(nid)
                else:      ready_nodes.append(nid)
        return list(ready_nodes)

    def add_node(self, node):
        ''' adds a new node to the dataset '''
        node['_id'] = self._incr_id()
        node['_st'] = 0
        self.nodes.append(node)
        return node['_id']

    def get_node(self, nid):
        ''' getter method for fetching a node by index/id '''
        return self.nodes[nid]

    def get_label(self, nid):
        ''' getter method for fetching a nodes label by id '''
        return self.get_node(nid)['_name']

    def mark_node_up(self, nid):
        ''' marks a node as up '''
        self.nodes[nid]['_st'] = 1
        return 

    def add_edge(self, a_id, b_id):
        ''' adds a new directed edge from node a to node b
            takes index ids of nodes as returned by add_node()
        '''
        if (a_id, b_id) in self.edges: return
        else:                          self.edges.append( (a_id, b_id) )
        return