Visualising Project Haystack graphs using pyhaystack and graphviz

haystack2dot.py

#!/usr/bin/env python

# Dump a Project Haystack tree to a graphviz `dot` file.
# Example usage:
# python haystack2dot.py --client-type widesky --client-uri http://localhost:8084 \
#        --client-username user@example.com --client-password password \
#        --client-id aaaaaaaa --client-secret bbbbbbbb \
#        --filter 'equipRef==@vrtsofteng.db1_ac1 and (point or modbusBlock or modbusReserved)' \
#        --filter 'device2Ref==@vrtsofteng.db1_ac1' \
#        --filter 'device2Ref==@vrtsofteng.vrtems_gw' \
#        --not-val wsgServiceConfig \
#        --not-ref siteRef > graph.dot \
#    && dot -Tpng -ograph.png -Grankdir=RL graph.dot
# (C) 2017 VRT Systems

# Rev 2: add --node-style and --ref-style

import pyhaystack.client
import hszinc
import argparse
import logging
import textwrap

import re

TAG_RE = re.compile(r'^([a-z][a-zA-Z0-9_]*)(.*)$')
OP_RE=re.compile(r'^ *(<|<=|!=|==|=|>=|>) *("([^\\"]|\\[\\"])*"|\'([^\\\']|\\[\\\'])*\'|[^;]*)(;?.*|)$')

class StyleFilter(object):
    def __init__(self, filter_expr, attributes_str):
        self.filter_match = eval('lambda t : %s' % filter_expr, {}, {})
        self.attributes = dict([
            attrval.split('=',1) for attrval in attributes_str.split(';')
        ])

    def match(self, entity):
        return self.filter_match(entity.tags)


def main(*args, **kwargs):
    parser = argparse.ArgumentParser()
    parser.add_argument('--client-type', dest='client_type', type=str,
            help='Project Haystack Server Type')
    parser.add_argument('--client-uri', dest='client_uri', type=str,
            help='Project Haystack URI')
    parser.add_argument('--client-username', dest='client_username', type=str,
            help='Project Haystack username')
    parser.add_argument('--client-password', dest='client_password', type=str,
            help='Project Haystack password')
    parser.add_argument('--client-id', dest='client_id', type=str,
            help='Project Haystack Client ID')
    parser.add_argument('--client-secret', dest='client_secret', type=str,
            help='Project Haystack secret')
    parser.add_argument('--client-arg', dest='client_args', action='append', type=str, nargs=2,
            help='Set arbitrary Haystack client options')

    parser.add_argument('--id', dest='ids', action='append', type=str, default=[],
            help='Include this entity ID in the graph')
    parser.add_argument('--filter', dest='filters', action='append', type=str, default=[],
            help='Include entities matching this filter in the graph')

    parser.add_argument('--not-id', dest='not_ids', action='append', type=str, default=[],
            help='Exclude this entity ID in the graph')
    parser.add_argument('--not-filter', dest='not_filters', action='append', type=str, default=[],
            help='Exclude entities matching this filter in the graph')

    parser.add_argument('--tag', dest='tags', action='append', type=str, default=[],
            help='Include this tag in the graph')
    parser.add_argument('--not-tag', dest='not_tags', action='append', type=str, default=[],
            help='Exclude this tag in the graph')
    parser.add_argument('--not-val', dest='not_vals', action='append', type=str, default=[],
            help='Suppress the value of this tag')

    parser.add_argument('--ref', dest='refs', action='append', type=str, default=[],
            help='Follow this ref when encountered in the graph')
    parser.add_argument('--not-ref', dest='not_refs', action='append', type=str, default=[],
            help='Do not follow this ref when encountered in the graph')

    parser.add_argument('--node-style', dest='node_style', action='append', type=str, default=[],
            nargs=2, metavar=('PY_EXPR','STYLE'),
            help='Apply styling to nodes if the tags match the given expression')
    parser.add_argument('--ref-style', dest='ref_style', action='append', type=str, default=[],
            nargs=2, metavar=('TAG_NAME','STYLE'),
            help='Apply styling to refs matching the given name')

    args = parser.parse_args(*args, **kwargs)

    logging.basicConfig(level=logging.DEBUG)

    node_styles = [
            StyleFilter(filter_str, attr_str)
            for filter_str, attr_str
            in args.node_style
    ]

    ref_style = {}
    for (ref, r_style) in args.ref_style:
        ref_style[ref] = dict([attrval.split('=',1) for attrval in r_style.split(';')])

    client_args = {}
    if args.client_args:
        client_args.update(dict(args.client_args))

    for key, arg in (   ('implementation', 'client_type'),
                        ('uri', 'client_uri'),
                        ('username', 'client_username'),
                        ('password', 'client_password'),
                        ('client_id', 'client_id'),
                        ('client_secret', 'client_secret')  ):
        val = getattr(args, arg)
        if val:
            client_args[key] = val

    client = pyhaystack.client.get_instance(**client_args)

    # Entity cache
    entity = {}

    # Build up the include/exclude list
    not_ids = set(args.not_ids)
    ids = set(args.ids)

    for filter_expr in args.not_filters:
        logging.info('Searching filter: %s', filter_expr)
        find_op = client.find_entity(filter_expr)
        find_op.wait()

        for en in find_op.result.keys():
            logging.debug('Adding to exclude list: %s', en)
            not_ids.add(en)

    # Process excludes for explicit ID list
    ids -= not_ids

    # Fetch the IDs in the explicit ID list
    logging.info('Fetching explicitly listed IDs')
    fetch_op = client.get_entity(list(ids))
    fetch_op.wait()
    entity.update(fetch_op.result)

    # Fetch the entities that match the given filters.
    for filter_expr in args.filters:
        logging.info('Searching filter: %s', filter_expr)
        find_op = client.find_entity(filter_expr)
        find_op.wait()

        for en, e in find_op.result.items():
            if en in not_ids:
                logging.debug('Ignoring excluded entity %s', e.id)
            else:
                logging.debug('Adding to include list: %s', e.id)
                entity[en] = e
                ids.add(en)

    # Work through the list of IDs found
    tags = set(args.tags) - set(args.not_tags)
    not_refs = set(args.not_refs)
    not_vals = set(args.not_vals)
    refs = set(args.refs) - not_refs
    visited = set()
    node = {}
    edge = {}
    todo = list(ids)

    while todo:
        new_ids = []
        for eid in todo:
            if eid in visited:
                continue

            e = entity[eid]

            logging.debug('Inspecting entity: %s', e)
            visited.add(eid)

            # First, process the non-ref tags (and not-followed refs)
            node_data = []
            e_refs = set()
            for tag in (tags or list(e.tags.keys())):
                if tag not in e.tags:
                    # Not present
                    continue

                # dis is handled
                if tag == 'dis':
                    continue

                val = e.tags[tag]
                if val is hszinc.MARKER:
                    # Just list the tag name
                    node_data.append(tag)
                    continue

                if ((tag in refs) or (not refs)) \
                        and isinstance(val, hszinc.Ref) \
                        and (val.name not in not_ids):
                    # We handle refs specially
                    if tag not in not_refs:
                        e_refs.add(tag)
                    continue

                if tag in not_vals:
                    # Suppress the value for this tag
                    node_data.append('%s=...' % tag)
                    continue

                node_data.append('%s=%s' % (tag,
                    hszinc.dump_scalar(val, mode=hszinc.MODE_ZINC)))

            node_data.sort()
            node_def = dict(
                    shape='oval',
                    label='%s\nid: %s\n\n%s' % (
                        e.tags['dis'], e.id,
                        '\n'.join(textwrap.TextWrapper().wrap(', '.join(node_data)))
                    )
            )
            node[eid] = node_def

            for node_style in node_styles:
                if node_style.match(e):
                    node_def.update(node_style.attributes)

            for ref in (refs or e_refs):
                if ref not in e.tags:
                    # Not present
                    continue

                val = e.tags[ref]
                if val.name in not_ids:
                    continue

                edge_def = dict(label=ref)
                edge_def.update(ref_style.get(ref, {}))
                edge['"%s" -> "%s"' % (eid, val.name)] = edge_def

                if val.name not in visited:
                    logging.debug('Will follow %s[%s] -> %s', e.id, ref, val)
                    new_ids.append(val.name)

        # Visit the new IDs and add them to the todo list.
        todo = new_ids
        fetch_op = client.get_entity(todo)
        fetch_op.wait()
        entity.update(fetch_op.result)

    # Generate the Graphviz file
    print ('digraph model {')
    for nid, a in node.items():
        print ('  "%s" [%s];' % (
            nid, ', '.join(
                ['%s="%s"' % (an, ar.replace('\\','\\\\').replace('"','\\"'))
                    for an, ar in a.items()]
            )))
    for eid, a in edge.items():
        if a is not None:
            print ('  %s [%s];' % (
                eid, ', '.join(
                ['%s="%s"' % (an, ar.replace('\\','\\\\').replace('"','\\"'))
                    for an, ar in a.items()]
            )))
        else:
            print ('  %s' % eid)
    print ('}')

if __name__ == '__main__':
    main()

http://static.vk4msl.id.au/pyhaystack/2017/12/07-haystack2dot/graph.png is example output from this script.