graph.py

# -*- coding: utf-8 -*-
# First stab at implementing graphs in Python.
# Horrifically undocumented.

import operator
import textwrap


class SimpleStruct(type):
    
    @staticmethod
    def new(cls, *values):
        return tuple.__new__(cls, values)
    
    def __new__(mcls, *properties):
        attrs = {}
        for i, prop in enumerate(properties):
            attrs[prop] = property(operator.itemgetter(i))
        
        attrs['__new__'] = mcls.new
        
        return type('simple_struct', (tuple,), attrs)


class Node(SimpleStruct('name', 'value')):
    
    def __new__(cls, name='', value=None):
        return SimpleStruct.new(cls, name, value)
    
    def __repr__(self):
        if self.value is not None:
            return '<Node(%r, %r) at 0x%x>' % (self.name, self.value, id(self))
        elif not self.name:
            return '<Node at 0x%x>' % (id(self),)
        return '<Node(%r) at 0x%x>' % (self.name, id(self))


class Arc(SimpleStruct('origin', 'destination', 'weight', 'directed')):
    
    def __new__(cls, origin, destination, weight=None, directed=True):
        return SimpleStruct.new(
            cls, origin, destination, weight, directed)
    
    def __repr__(self):
        description = []
        
        if self.directed:
            description.append('<DirectedArc from: %r\n' % (self.origin,))
            indent = 17
        else:
            description.append('<Arc from: %r\n' % (self.origin,))
            indent = 9
        
        description.append('%%%ds: ' % (indent,) % ('to'))
        description.append(repr(self.destination))
        
        if self.weight is not None:
            description.append('\n%%%ds: ' % (indent,) % ('weight'))
            description.append(repr(self.weight))
        
        description.append('>')
        
        return ''.join(description)
    
    def destination_is(self, node):
        return (
            (self.destination is node) or
            ((not self.directed) and (self.origin is node)))
    
    def origin_is(self, node):
        return (
            (self.origin is node) or
            ((not self.directed) and (self.destination is node)))
    
    def node_other_than(self, node):
        if self.origin is node:
            return self.destination
        elif self.destination is node:
            return self.origin


class Graph(SimpleStruct('nodes', 'arcs')):
    
    def __new__(cls, nodes=None, arcs=None):
        if nodes is None: nodes = set()
        if arcs is None: arcs = set()
        
        return SimpleStruct.new(cls, nodes, arcs)
    
    def __repr__(self):
        return '<Graph at 0x%x>' % (id(self),)
    
    def __contains__(self, node_or_arc):
        if isinstance(node_or_arc, Node):
            return node_or_arc in self.nodes
        elif isinstance(node_or_arc, Arc):
            return node_or_arc in self.arcs
    
    def make_subset(self, with_nodes=True, with_arcs=False):
        kwargs = {}
        if with_nodes: kwargs['nodes'] = self.nodes
        if with_arcs: kwargs['arcs'] = self.arcs
        
        return type(self)(**kwargs)
    
    def is_subset_of(self, other):
        return (
            other.nodes.issuperset(self.nodes) and
            other.arcs.issuperset(self.arcs))
    
    def is_superset_of(self, other):
        return (
            other.nodes.issubset(self.nodes) and
            other.arcs.issubset(self.arcs))
    
    def incident_arcs(self, node):
        return set(arc for arc in self.arcs if arc.destination_is(node))
    
    def outcident_arcs(self, node):
        return set(arc for arc in self.arcs if arc.origin_is(node))
    
    def nodes_where(self, name=None, value=None):
        tests = [lambda node: True]
        if name is not None:
            tests.append(lambda node: operator.eq(node.name, name))
        if value is not None:
            tests.append(lambda node: operator.eq(node.value, value))
        
        return (n for n in self.nodes if all(test(n) for test in tests))
    
    def arcs_where(self, origin=None, destination=None, **kwargs):
        tests = [lambda arc: True]
        if origin is not None:
            tests.append(lambda arc: arc.origin_is(origin))
        if destination is not None:
            tests.append(lambda arc: arc.destination_is(destination))
        if 'weight' in kwargs:
            weight = kwargs['weight']
            tests.append(lambda arc: operator.eq(arc.weight, weight))
        
        return (a for a in self.arcs if all(test(a) for test in tests))
    
    def find_shortest_path(self, origin, destination):
        Label = SimpleStruct('distance', 'via', 'is_fixed')
        # Setting the default distance to infinity makes sure that if any node
        # has the default label set it will instantly be set to the first label
        # found.
        default_label = Label(float('inf'), None, False)
        labels = {origin: Label(0, None, True)}
        fixed = labels.copy() # This will eventually hold all fixed labels.
        
        while destination not in fixed:
            
            for node, label in fixed.items():
                
                for arc in self.outcident_arcs(node):
                    
                    dest = arc.node_other_than(node)
                    if dest in fixed:
                        continue
                    
                    curr_label = labels.setdefault(dest, default_label)
                    temp_distance = label.distance + (arc.weight or 1)
                    
                    if temp_distance < curr_label.distance:
                        labels[dest] = Label(temp_distance, arc, False)
            
            unfixed = [
                (node, label) for node, label in labels.items()
                    if not label.is_fixed]
            if not unfixed:
                return None
            
            next_fixed = min(unfixed, key=lambda (node, label): label.distance)
            node, temp_label = next_fixed
            labels[node] = Label(temp_label.distance, temp_label.via, True)
            fixed[node] = labels[node]
        
        path = []
        current_node = destination
        total_weight = fixed[destination].distance
        
        while current_node is not origin:
            path.insert(0, fixed[current_node].via)
            current_node = fixed[current_node].via.node_other_than(current_node)
        
        return total_weight, path