1 #!/usr/bin/env python 2 3 """ 4 Graphviz utilities. 5 6 Copyright (C) 2018 Paul Boddie <paul@boddie.org.uk> 7 8 This program is free software; you can redistribute it and/or modify it under 9 the terms of the GNU General Public License as published by the Free Software 10 Foundation; either version 3 of the License, or (at your option) any later 11 version. 12 13 This program is distributed in the hope that it will be useful, but WITHOUT 14 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS 15 FOR A PARTICULAR PURPOSE. See the GNU General Public License for more 16 details. 17 18 You should have received a copy of the GNU General Public License along with 19 this program. If not, see <http://www.gnu.org/licenses/>. 20 """ 21 22 from os.path import exists, join 23 from StringIO import StringIO 24 from subprocess import Popen, PIPE 25 import gzip 26 import sha 27 import xml.sax 28 29 # Configurable paths and locations. 30 31 DIAGRAM_TOOLS_PATH = "/home/paulb/Software/Graphical/diagram-tools" 32 GRAPHVIZ_PATH = "/usr/bin" 33 XSLT_PROCESSOR = "/usr/bin/xsltproc" 34 35 # Graphviz "filter" programs performing layout. 36 37 FILTERS = ['circo', 'dot', 'fdp', 'neato', 'twopi'] 38 39 # Supported output formats. 40 41 IMAGE_FORMATS = ['png', 'gif'] 42 SVG_FORMATS = ['svg', 'svgz'] 43 44 OUTPUT_FORMATS = IMAGE_FORMATS + SVG_FORMATS + \ 45 ['dia', 'fig', 'hpgl', 'imap', 'mif', 'pcl', 'ps'] 46 47 # XSL transformations for SVG output. 48 49 TRANSFORMS = { 50 "notugly" : join(DIAGRAM_TOOLS_PATH, "notugly.xsl"), 51 } 52 53 54 55 # Utility functions. 56 57 def encode(s, encoding): 58 59 "Encode 's' using 'encoding' if Unicode." 60 61 if isinstance(s, unicode): 62 return s.encode(encoding) 63 else: 64 return s 65 66 class MetadataParser(xml.sax.handler.ContentHandler): 67 68 "Parse metadata from the svg element." 69 70 def __init__(self): 71 self.attrs = {} 72 73 def startElement(self, name, attrs): 74 if name == self.tagname: 75 self.attrs = dict(attrs) 76 77 def parse(self, f): 78 79 "Parse content from the file object 'f' using reasonable defaults." 80 81 try: 82 parser = xml.sax.make_parser() 83 parser.setContentHandler(self) 84 parser.setErrorHandler(xml.sax.handler.ErrorHandler()) 85 parser.setFeature(xml.sax.handler.feature_external_ges, 0) 86 parser.parse(f) 87 finally: 88 f.close() 89 90 def get_metadata(self, data, tagname): 91 92 "Process 'data', returning attributes from 'tagname'." 93 94 self.tagname = tagname 95 96 f = StringIO(data) 97 try: 98 self.parse(f) 99 finally: 100 f.close() 101 102 return self.attrs 103 104 def get_output_identifier(text): 105 106 "Return an output identifier for the given 'text'." 107 108 return sha.new(encode(text, 'utf-8')).hexdigest() 109 110 def get_program(filter): 111 112 "Return the program for the given 'filter'." 113 114 if not filter in FILTERS: 115 return None 116 else: 117 return join(GRAPHVIZ_PATH, filter) 118 119 def transform_output(process, format, transforms): 120 121 "Transform the output from 'process' as 'format' using 'transforms'." 122 123 # No transformation can occur if the processor is missing. 124 125 if not exists(XSLT_PROCESSOR): 126 return process 127 128 # Chain transformation processors, each accepting the output of the 129 # preceding one, with the first accepting the initial Graphviz output. 130 131 for transform in transforms: 132 stylesheet = TRANSFORMS.get(transform) 133 134 # Ignore unrecognised or missing stylesheets. 135 136 if not stylesheet or not exists(stylesheet): 137 continue 138 139 # Invoke the processor, indicating standard input as the source 140 # document. 141 # Example: /usr/bin/dot /usr/local/share/diagram-tools/notugly.xsl - 142 143 process = Popen( 144 [XSLT_PROCESSOR, stylesheet, "-"], 145 shell=False, 146 stdin=process.stdout, 147 stdout=PIPE, 148 stderr=PIPE, 149 close_fds=True) 150 151 return process 152 153 def writefile(s, filename, compressed=False): 154 155 "Write 's' to the file having 'filename'." 156 157 if compressed: 158 f = gzip.open(filename, "w") 159 else: 160 f = open(filename, "w") 161 162 try: 163 f.write(s) 164 finally: 165 f.close() 166 167 168 169 # Classes for interacting with Graphviz. 170 171 class GraphvizError(Exception): 172 173 "An error produced when using Graphviz." 174 175 def __init__(self, errors): 176 self.errors = errors 177 178 class Graphviz: 179 180 "A Graphviz configuration for single or repeated invocation." 181 182 def __init__(self, filter, text, identifier): 183 184 """ 185 Employ the given 'filter' to produce a graph from the given 'text'. The 186 output 'identifier' for the text is used to provide a filename, if 187 required. 188 """ 189 190 self.filter = filter 191 self.text = text 192 self.identifier = identifier 193 194 def call(self, format, transforms=None, filename=None): 195 196 """ 197 Invoke Graphviz to produce output in the given 'format'. Any 198 'transforms' are used to transform the output, if appropriate. Any 199 given 'filename' is used to write to a file. 200 """ 201 202 program = get_program(self.filter) 203 204 # Generate uncompressed SVG for later compression. 205 206 graphviz_format = format == "svgz" and "svg" or format 207 208 # Indicate a filename for direct output for non-SVG formats. 209 210 svg = format in SVG_FORMATS 211 options = filename and not svg and ["-o", filename] or [] 212 213 # Invoke the layout program, with the text to be provided on its 214 # standard input. 215 # Example: /usr/bin/dot -Tsvg -o filename 216 217 start = end = Popen( 218 [program, '-T%s' % graphviz_format] + options, 219 shell=False, 220 stdin=PIPE, 221 stdout=PIPE, 222 stderr=PIPE) 223 224 # Chain the invocation to transformations, if appropriate. 225 226 if svg and transforms: 227 end = transform_output(start, format, transforms) 228 229 # Send the graph to the filter. 230 231 start.stdin.write(encode(self.text, 'utf-8')) 232 233 if end is not start: 234 start.stdin.close() 235 236 # Obtain the eventual output. 237 238 (self.output, errors) = end.communicate() 239 240 # Obtain any metadata. 241 242 if svg: 243 parser = MetadataParser() 244 self.metadata = parser.get_metadata(self.output, "svg") 245 elif format == "cmapx": 246 parser = MetadataParser() 247 self.metadata = parser.get_metadata(self.output, "map") 248 else: 249 self.metadata = {} 250 251 # Test for errors. 252 253 if end.wait() != 0: 254 raise GraphvizError, errors 255 256 # Write the file separately, if requested. 257 258 if svg and filename: 259 writefile(self.get_output(), filename, format == "svgz") 260 261 def get_metadata(self): 262 return self.metadata 263 264 def get_output(self): 265 return self.output 266 267 # vim: tabstop=4 expandtab shiftwidth=4