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 from xml.sax.saxutils import XMLGenerator 26 import gzip 27 import sha 28 import xml.sax 29 30 # Configurable paths and locations. 31 32 DIAGRAM_TOOLS_PATH = "/home/paulb/Software/Graphical/diagram-tools" 33 GRAPHVIZ_PATH = "/usr/bin" 34 XSLT_PROCESSOR = "/usr/bin/xsltproc" 35 36 # Graphviz "filter" programs performing layout. 37 38 FILTERS = ['circo', 'dot', 'fdp', 'neato', 'twopi'] 39 40 # Supported output formats. 41 42 IMAGE_FORMATS = ['png', 'gif'] 43 SVG_FORMATS = ['svg', 'svgz'] 44 45 OUTPUT_FORMATS = IMAGE_FORMATS + SVG_FORMATS + \ 46 ['dia', 'fig', 'hpgl', 'imap', 'mif', 'pcl', 'ps'] 47 48 # XSL transformations for SVG output. 49 50 TRANSFORMS = { 51 "notugly" : join(DIAGRAM_TOOLS_PATH, "notugly.xsl"), 52 } 53 54 55 56 # Utility functions. 57 58 def encode(s, encoding): 59 60 "Encode 's' using 'encoding' if Unicode." 61 62 if isinstance(s, unicode): 63 return s.encode(encoding) 64 else: 65 return s 66 67 class Parser(xml.sax.handler.ContentHandler): 68 69 "Common XML parsing functionality." 70 71 def parse(self, f): 72 73 "Parse content from the file object 'f' using reasonable defaults." 74 75 try: 76 parser = xml.sax.make_parser() 77 parser.setContentHandler(self) 78 parser.setErrorHandler(xml.sax.handler.ErrorHandler()) 79 parser.setFeature(xml.sax.handler.feature_external_ges, 0) 80 parser.parse(f) 81 finally: 82 f.close() 83 84 def parse_data(self, data): 85 86 "Parse the given 'data'." 87 88 f = StringIO(data) 89 try: 90 self.parse(f) 91 finally: 92 f.close() 93 94 class MetadataParser(Parser): 95 96 "Parse metadata from the svg element." 97 98 def __init__(self): 99 self.attrs = {} 100 101 def startElement(self, name, attrs): 102 if name == self.tagname: 103 self.attrs = dict(attrs) 104 105 def get_metadata(self, data, tagname): 106 107 "Process 'data', returning attributes from 'tagname'." 108 109 self.tagname = tagname 110 self.parse_data(data) 111 return self.attrs 112 113 class DocumentSelector(XMLGenerator, Parser): 114 115 "Parse a document and obtain the serialisation of the document node." 116 117 def startDocument(self): 118 pass 119 120 def get_output_identifier(text): 121 122 "Return an output identifier for the given 'text'." 123 124 return sha.new(encode(text, 'utf-8')).hexdigest() 125 126 def get_program(filter): 127 128 "Return the program for the given 'filter'." 129 130 if not filter in FILTERS: 131 return None 132 else: 133 return join(GRAPHVIZ_PATH, filter) 134 135 def transform_output(process, format, transforms): 136 137 "Transform the output from 'process' as 'format' using 'transforms'." 138 139 # No transformation can occur if the processor is missing. 140 141 if not exists(XSLT_PROCESSOR): 142 return process 143 144 # Chain transformation processors, each accepting the output of the 145 # preceding one, with the first accepting the initial Graphviz output. 146 147 for transform in transforms: 148 stylesheet = TRANSFORMS.get(transform) 149 150 # Ignore unrecognised or missing stylesheets. 151 152 if not stylesheet or not exists(stylesheet): 153 continue 154 155 # Invoke the processor, indicating standard input as the source 156 # document. 157 # Example: /usr/bin/dot /usr/local/share/diagram-tools/notugly.xsl - 158 159 process = Popen( 160 [XSLT_PROCESSOR, stylesheet, "-"], 161 shell=False, 162 stdin=process.stdout, 163 stdout=PIPE, 164 stderr=PIPE, 165 close_fds=True) 166 167 return process 168 169 def writefile(s, filename, compressed=False): 170 171 "Write 's' to the file having 'filename'." 172 173 if compressed: 174 f = gzip.open(filename, "w") 175 else: 176 f = open(filename, "w") 177 178 try: 179 f.write(s) 180 finally: 181 f.close() 182 183 184 185 # Classes for interacting with Graphviz. 186 187 class GraphvizError(Exception): 188 189 "An error produced when using Graphviz." 190 191 def __init__(self, errors): 192 self.errors = errors 193 194 class Graphviz: 195 196 "A Graphviz configuration for single or repeated invocation." 197 198 def __init__(self, filter, text): 199 200 """ 201 Employ the given 'filter' to produce a graph from the given 'text'. 202 """ 203 204 self.filter = filter 205 self.text = text 206 207 def call(self, format, transforms=None, filename=None): 208 209 """ 210 Invoke Graphviz to produce output in the given 'format'. Any 211 'transforms' are used to transform the output, if appropriate. Any 212 given 'filename' is used to write to a file. 213 """ 214 215 program = get_program(self.filter) 216 217 # Generate uncompressed SVG for later compression. 218 219 graphviz_format = format == "svgz" and "svg" or format 220 221 # Indicate a filename for direct output for non-SVG formats. 222 223 svg = format in SVG_FORMATS 224 options = filename and not svg and ["-o", filename] or [] 225 226 # Invoke the layout program, with the text to be provided on its 227 # standard input. 228 # Example: /usr/bin/dot -Tsvg -o filename 229 230 start = end = Popen( 231 [program, '-T%s' % graphviz_format] + options, 232 shell=False, 233 stdin=PIPE, 234 stdout=PIPE, 235 stderr=PIPE) 236 237 # Chain the invocation to transformations, if appropriate. 238 239 if svg and transforms: 240 end = transform_output(start, format, transforms) 241 242 # Send the graph to the filter. 243 244 start.stdin.write(encode(self.text, 'utf-8')) 245 246 if end is not start: 247 start.stdin.close() 248 249 # Obtain the eventual output. 250 251 (self.output, errors) = end.communicate() 252 253 # Obtain any metadata. 254 255 if svg: 256 parser = MetadataParser() 257 self.metadata = parser.get_metadata(self.output, "svg") 258 elif format == "cmapx": 259 parser = MetadataParser() 260 self.metadata = parser.get_metadata(self.output, "map") 261 else: 262 self.metadata = {} 263 264 # Test for errors. 265 266 if end.wait() != 0: 267 raise GraphvizError, errors 268 269 # Write the file separately, if requested. 270 271 if svg and filename: 272 writefile(self.get_output(), filename, format == "svgz") 273 274 def get_metadata(self): 275 return self.metadata 276 277 def get_output(self): 278 return self.output 279 280 def get_inline_output(self): 281 282 """ 283 Return a string containing the document element, excluding XML 284 boilerplate. 285 """ 286 287 f = StringIO() 288 parser = DocumentSelector(f, "utf-8") 289 290 try: 291 parser.parse_data(self.output) 292 return f.getvalue() 293 finally: 294 f.close() 295 296 # vim: tabstop=4 expandtab shiftwidth=4