path: root/lib.py

#!/usr/bin/env python
# vim: set fileencoding=UTF-8:
# Copyright 2013 Red Hat, Inc.
# Author: Jan Pokorný <jpokorny at redhat dot com>
# Distributed under GPLv2+;  generated content under CC-BY-SA 3.0
# (to view a copy, visit http://creativecommons.org/licenses/by-sa/3.0/)
"""Elements library incl. visual aspects"""

from sys import path
from os.path import expanduser, extsep
from pydot import Dot, Edge, Node, Subgraph

__all__ = []

# modify/mutate BLACKLIST in order to filter some elements out
# (not a beautiful, rather ad-hoc design, I know)
BLACKLIST = []


#
# Customization via meta level
#

def bl_test(x):
    return \
        not set((x.__class__,) + x.__class__.__bases__).intersection(BLACKLIST)


class LibMeta(type):
    def __init__(cls, name, bases, attrs):
        __all__.append(name)
        super(LibMeta, cls).__init__(name, bases, attrs)
        old_init = cls.__init__
        def new_init(self, *args, **kwargs):
            subgraphs = kwargs.pop('_subgraphs', ())
            nodes = kwargs.pop('_nodes', ())
            edges = kwargs.pop('_edges', ())
            old_init(self, *args, **kwargs)
            for s in filter(bl_test, subgraphs): s and self.add_subgraph(s)
            for n in filter(bl_test, nodes):     n and self.add_node(n)
            for e in filter(bl_test, edges):     e and self.add_edge(e)
        cls.__init__ = new_init


class LibDot(Dot):
    __metaclass__ = LibMeta
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('graph_type', 'digraph')
        super(LibDot, self).__init__(*args, **kwargs)


class LibSubgraph(Subgraph):
    __metaclass__ = LibMeta


class LibNode(Node):
    __metaclass__ = LibMeta


class LibEdge(Edge):
    __metaclass__ = LibMeta


#
# Subgraph inheritance
#


class SubgraphImportant(LibSubgraph):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', '#eaeaea')
        kwargs.setdefault('style', 'bold, filled, rounded')
        super(SubgraphImportant, self).__init__(*args, **kwargs)


class SubgraphStandard(LibSubgraph):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', '#f5f5f5')
        kwargs.setdefault('style', 'filled, rounded')
        super(SubgraphStandard, self).__init__(*args, **kwargs)


class SubgraphInvisible(LibSubgraph):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'invisible')
        super(SubgraphInvisible, self).__init__(*args, **kwargs)



#
# Node inheritance
#


class Program(LibNode):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('shape', 'box')
        kwargs.setdefault('style', 'filled')
        super(Program, self).__init__(*args, **kwargs)


class Agent(Program):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', 'lavenderblush')
        kwargs.setdefault('style', 'filled')
        super(Agent, self).__init__(*args, **kwargs)


class Library(Program):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', 'lavenderblush')
        kwargs.setdefault('style', 'filled')
        super(Library, self).__init__(*args, **kwargs)


class Executable(Program):
    pass


class Daemon(Executable):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('fillcolor', 'cornsilk')
        kwargs.setdefault('style', 'filled')
        super(Daemon, self).__init__(*args, **kwargs)


class Artefact(LibNode):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('shape', 'box3d')
        kwargs.setdefault('fillcolor', 'wheat')
        kwargs.setdefault('style', 'filled')
        super(Artefact, self).__init__(*args, **kwargs)


#
# Edge inheritance
#


class DBUS(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'dotted')
        super(DBUS, self).__init__(*args, **kwargs)


class OddjobExec(DBUS):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'oddjob exec')
        super(OddjobExec, self).__init__(*args, **kwargs)


class HTTPS(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'bold')
        super(HTTPS, self).__init__(*args, **kwargs)


class LuciHTTPS(HTTPS):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 8084')
        super(LuciHTTPS, self).__init__(*args, **kwargs)


class ApplicationSpecificProtocol(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('style', 'dashed')
        super(ApplicationSpecificProtocol, self).__init__(*args, **kwargs)


class RICCIRPC(ApplicationSpecificProtocol):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 11111')
        super(RICCIRPC, self).__init__(*args, **kwargs)


class SNMP(ApplicationSpecificProtocol):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 161')
        super(SNMP, self).__init__(*args, **kwargs)


class CIM(ApplicationSpecificProtocol):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('label', 'port 898[89]')
        super(CIM, self).__init__(*args, **kwargs)


class Consume(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'darkgreen')
        kwargs.setdefault('fontcolor', 'darkgreen')
        super(Consume, self).__init__(*args, **kwargs)


class ConsumeReversed(Consume):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('dir', 'back')
        super(ConsumeReversed, self).__init__(*reversed(args), **kwargs)


class Produce(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'orangered')
        kwargs.setdefault('fontcolor', 'orangered')
        super(Produce, self).__init__(*args, **kwargs)


class Delegate(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'navy')
        kwargs.setdefault('fontcolor', 'navy')
        kwargs.setdefault('dir', 'both')
        super(Delegate, self).__init__(*args, **kwargs)


# peer-to-peer
class Exchange(LibEdge):
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('color', 'saddlebrown')
        kwargs.setdefault('fontcolor', 'saddlebrown')
        kwargs.setdefault('dir', 'both')
        super(Exchange, self).__init__(*args, **kwargs)


# mixing

class DelegateOddjobExec(Delegate, OddjobExec):
    pass


class DelegateRICCIRPC(Delegate, RICCIRPC):
    pass


class DelegateLuciHTTPS(Delegate, LuciHTTPS):
    pass


class DelegateSNMP(Delegate, SNMP):
    pass


class DelegateCIM(Delegate, CIM):
    pass


#
# main-helpers
#


def export(fnc):
    __all__.append(fnc.__name__)
    return fnc


@export
def gen_graph(graph, blacklist=(), **kwargs):
    # can eventually restore the state
    BLACKLIST[:] = blacklist
    return graph()


@export
def xdot_graph(*args, **kwargs):
    import gtk
    import gtk.gdk
    try:
        import xdot
    except ImportError:
        #path.append(expanduser('~/wrkspc/goss-medium/jrfonseca.xdot'))
        #import xdot
        print 'missing xdot; use "pip install xdot" or equivalent'
        raise
    window = xdot.DotWindow()
    window.set_dotcode(gen_graph(*args, **kwargs).to_string())
    window.connect('destroy', gtk.main_quit)
    gtk.main()


@export
def main(graph, argv, *args, **kws):
    x, argv = (1, argv[1:]) if len(argv) > 1 and argv[1] == '-x' else (0, argv)
    output = kws.pop('output', 'sinenomine')
    if x:
        xdot_graph(graph, **kws)
    else:
        ext = argv[1] if len(argv) > 1 else 'pdf'
        fmt = {'dot': 'raw'}.get(ext, ext)
        output += extsep + ext
        gen_graph(graph, **kws).write(output, format=fmt, prog=('dot',) + args)