#!/usr/bin/env python

"""

Graphviz utilities.

Copyright (C) 2018 Paul Boddie <paul@boddie.org.uk>

This program is free software; you can redistribute it and/or modify it under

the terms of the GNU General Public License as published by the Free Software

Foundation; either version 3 of the License, or (at your option) any later

version.

This program is distributed in the hope that it will be useful, but WITHOUT

ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS

FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more

details.

You should have received a copy of the GNU General Public License along with

this program.  If not, see <http://www.gnu.org/licenses/>.

"""

from os.path import exists, join

from StringIO import StringIO

from subprocess import Popen, PIPE

from xml.sax.saxutils import XMLGenerator

import gzip

import sha

import xml.sax

# Configurable paths and locations.

DIAGRAM_TOOLS_PATH  = "/home/paulb/Software/Graphical/diagram-tools"

GRAPHVIZ_PATH       = "/usr/bin"

XSLT_PROCESSOR      = "/usr/bin/xsltproc"

# Graphviz "filter" programs performing layout.

FILTERS = ['circo', 'dot', 'fdp', 'neato', 'twopi']

# Supported output formats.

IMAGE_FORMATS = ['png', 'gif']

SVG_FORMATS = ['svg', 'svgz']

OUTPUT_FORMATS = IMAGE_FORMATS + SVG_FORMATS + \

                 ['dia', 'fig', 'hpgl', 'imap', 'mif', 'pcl', 'ps']

# XSL transformations for SVG output.

TRANSFORMS = {

    "notugly" : join(DIAGRAM_TOOLS_PATH, "notugly.xsl"),

    }

# Utility functions.

def encode(s, encoding):

    "Encode 's' using 'encoding' if Unicode."

    if isinstance(s, unicode):

        return s.encode(encoding)

    else:

        return s

class Parser(xml.sax.handler.ContentHandler):

    "Common XML parsing functionality."

    def parse(self, f):

        "Parse content from the file object 'f' using reasonable defaults."

        try:

            parser = xml.sax.make_parser()

            parser.setContentHandler(self)

            parser.setErrorHandler(xml.sax.handler.ErrorHandler())

            parser.setFeature(xml.sax.handler.feature_external_ges, 0)

            parser.parse(f)

        finally:

            f.close()

    def parse_data(self, data):

        "Parse the given 'data'."

        f = StringIO(data)

        try:

            self.parse(f)

        finally:

            f.close()

class MetadataParser(Parser):

    "Parse metadata from the svg element."

    def __init__(self):

        self.attrs = {}

    def startElement(self, name, attrs):

        if name == self.tagname:

            self.attrs = dict(attrs)

    def get_metadata(self, data, tagname):

        "Process 'data', returning attributes from 'tagname'."

        self.tagname = tagname

        self.parse_data(data)

        return self.attrs

class DocumentSelector(XMLGenerator, Parser):

    "Parse a document and obtain the serialisation of the document node."

    def startDocument(self):

        pass

def get_output_identifier(text):

    "Return an output identifier for the given 'text'."

    return sha.new(encode(text, 'utf-8')).hexdigest()

def get_program(filter):

    "Return the program for the given 'filter'."

    if not filter in FILTERS:

        return None

    else:

        return join(GRAPHVIZ_PATH, filter)

def transform_output(process, format, transforms):

    "Transform the output from 'process' as 'format' using 'transforms'."

    # No transformation can occur if the processor is missing.

    if not exists(XSLT_PROCESSOR):

        return process

    # Chain transformation processors, each accepting the output of the

    # preceding one, with the first accepting the initial Graphviz output.

    for transform in transforms:

        stylesheet = TRANSFORMS.get(transform)

        # Ignore unrecognised or missing stylesheets.

        if not stylesheet or not exists(stylesheet):

            continue

        # Invoke the processor, indicating standard input as the source

        # document.

        # Example: /usr/bin/dot /usr/local/share/diagram-tools/notugly.xsl -

        process = Popen(

            [XSLT_PROCESSOR, stylesheet, "-"],

            shell=False,

            stdin=process.stdout,

            stdout=PIPE,

            stderr=PIPE,

            close_fds=True)

    return process

def writefile(s, filename, compressed=False):

    "Write 's' to the file having 'filename'."

    if compressed:

        f = gzip.open(filename, "w")

    else:

        f = open(filename, "w")

    try:

        f.write(s)

    finally:

        f.close()

# Classes for interacting with Graphviz.

class GraphvizError(Exception):

    "An error produced when using Graphviz."

    def __init__(self, errors):

        self.errors = errors

class Graphviz:

    "A Graphviz configuration for single or repeated invocation."

    def __init__(self, filter, text):

        """

        Employ the given 'filter' to produce a graph from the given 'text'.

        """

        self.filter = filter

        self.text = text

    def call(self, format, transforms=None, filename=None):

        """

        Invoke Graphviz to produce output in the given 'format'. Any

        'transforms' are used to transform the output, if appropriate. Any

        given 'filename' is used to write to a file.

        """

        program = get_program(self.filter)

        # Generate uncompressed SVG for later compression.

        graphviz_format = format == "svgz" and "svg" or format

        # Indicate a filename for direct output for non-SVG formats.

        svg = format in SVG_FORMATS

        options = filename and not svg and ["-o", filename] or []

        # Invoke the layout program, with the text to be provided on its

        # standard input.

        # Example: /usr/bin/dot -Tsvg -o filename

        start = end = Popen(

            [program, '-T%s' % graphviz_format] + options,

            shell=False,

            stdin=PIPE,

            stdout=PIPE,

            stderr=PIPE)

        # Chain the invocation to transformations, if appropriate.

        if svg and transforms:

            end = transform_output(start, format, transforms)

        # Send the graph to the filter.

        start.stdin.write(encode(self.text, 'utf-8'))

        if end is not start:

            start.stdin.close()

        # Obtain the eventual output.

        (self.output, errors) = end.communicate()

        # Test for errors.

        if end.wait() != 0:

            raise GraphvizError, errors

        # Obtain any metadata.

        if svg:

            parser = MetadataParser()

            self.metadata = parser.get_metadata(self.output, "svg")

        elif format == "cmapx":

            parser = MetadataParser()

            self.metadata = parser.get_metadata(self.output, "map")

        else:

            self.metadata = {}

        # Write the file separately, if requested.

        if svg and filename:

            writefile(self.get_output(), filename, format == "svgz")

    def get_metadata(self):

        return self.metadata

    def get_output(self):

        return self.output

    def get_inline_output(self):

        """

        Return a string containing the document element, excluding XML

        boilerplate.

        """

        f = StringIO()

        parser = DocumentSelector(f, "utf-8")

        try:

            parser.parse_data(self.output)

            return f.getvalue()

        finally:

            f.close()

# vim: tabstop=4 expandtab shiftwidth=4